"""Reactive state store for the chat module.

Provides centralized state management with subscriber notifications.
Widgets subscribe to state changes and react accordingly.
"""

from dataclasses import replace
from typing import Callable
from uuid import UUID

from alloy_runtime_types.dtos.sessions import MessageResponse, SessionSummary
from alloy_runtime_sdk.logging.config import get_logger

from cli.tui.chat.types import ChatPhase, ChatState, RenderMode, SessionContext

logger = get_logger(__name__)


class ChatStore:
    """Reactive state store for chat screen.

    Follows the observable store pattern:
    - State is accessed via .state property (immutable snapshot)
    - Mutations happen through methods that create new state
    - Subscribers are notified on every state change

    Usage:
        store = ChatStore()

        # Subscribe to changes
        def on_change(state: ChatState):
            print(f"Phase: {state.phase}")

        unsubscribe = store.subscribe(on_change)

        # Mutate state
        store.set_phase(ChatPhase.SENDING)  # triggers on_change

        # Unsubscribe when done
        unsubscribe()
    """

    def __init__(self) -> None:
        """Initialize with default state."""
        self._state = ChatState()
        self._subscribers: list[Callable[[ChatState], None]] = []

    @property
    def state(self) -> ChatState:
        """Get current immutable state snapshot."""
        return self._state

    def subscribe(self, callback: Callable[[ChatState], None]) -> Callable[[], None]:
        """Subscribe to state changes.

        Args:
            callback: Function called with new state on every change.

        Returns:
            Unsubscribe function - call to stop receiving updates.
        """
        self._subscribers.append(callback)
        return lambda: self._subscribers.remove(callback)

    def _notify(self) -> None:
        """Notify all subscribers of state change."""
        for subscriber in self._subscribers:
            subscriber(self._state)

    def _update(self, **changes: object) -> None:
        """Update state with changes and notify subscribers.

        Creates a new ChatState with the given field changes.
        """
        self._state = replace(self._state, **changes)
        self._notify()

    # === Phase Management ===

    def set_phase(self, phase: ChatPhase) -> None:
        """Set the current chat phase."""
        old_phase = self._state.phase
        if old_phase != phase:
            logger.debug(
                "store_phase_changed", old_phase=old_phase.name, new_phase=phase.name
            )
        self._update(phase=phase)

    def is_busy(self) -> bool:
        """Check if the chat is in a busy state (can't accept new actions)."""
        return self._state.phase != ChatPhase.IDLE

    # === Session Management ===

    def set_session(
        self,
        session_id: UUID,
        session_name: str | None,
        agent_id: UUID | None,
        agent_name: str | None,
        message_count: int,
        provider_key: str | None = None,
        model_name: str | None = None,
        system_instruction: str | None = None,
    ) -> None:
        """Set the active session context.

        Args:
            session_id: The session ID.
            session_name: Optional session name.
            agent_id: Agent ID (for agent mode).
            agent_name: Agent name (for agent mode).
            message_count: Number of messages in session.
            provider_key: Provider key (for direct model mode).
            model_name: Model name (for direct model mode).
            system_instruction: System instruction (for direct model mode).
        """
        session = SessionContext(
            session_id=session_id,
            session_name=session_name,
            agent_id=agent_id,
            agent_name=agent_name,
            message_count=message_count,
            provider_key=provider_key,
            model_name=model_name,
            system_instruction=system_instruction,
        )
        logger.info(
            "store_session_set",
            session_id=str(session_id),
            session_name=session_name,
            agent_name=agent_name,
            model_name=model_name,
            message_count=message_count,
        )
        self._update(session=session)

    def clear_session(self) -> None:
        """Clear the active session."""
        logger.debug("store_session_cleared")
        self._update(session=None, messages=())

    def set_sessions(self, sessions: list[SessionSummary]) -> None:
        """Set the session list."""
        self._update(sessions=tuple(sessions))

    def set_session_search_query(self, query: str) -> None:
        """Set the session search query."""
        self._update(session_search_query=query)

    # === Message Management ===

    def set_messages(self, messages: list[MessageResponse]) -> None:
        """Set the message list for current session."""
        logger.debug("store_messages_set", message_count=len(messages))
        self._update(messages=tuple(messages))

    # === Error Management ===

    def set_error(self, error: str | None) -> None:
        """Set or clear error state."""
        if error:
            logger.warning("store_error_set", error=error)
        self._update(error=error)

    def clear_error(self) -> None:
        """Clear error state."""
        self._update(error=None)

    # === Convenience Methods ===

    def start_generation(self, user_message: str | None = None) -> None:
        """Transition to sending state with optional pending user message."""
        logger.info(
            "store_generation_started",
            has_user_message=user_message is not None,
            message_preview=user_message,
        )
        self._update(
            phase=ChatPhase.SENDING,
            error=None,
            pending_user_message=user_message,
        )

    def finish_generation(self, *, clear_pending_user_message: bool = True) -> None:
        """Transition from sending back to idle.

        Args:
            clear_pending_user_message: Whether to clear the optimistic pending user
                message immediately. Callers that reload persisted messages after a
                successful generation can leave this False until reload completes.
        """
        logger.info(
            "store_generation_finished",
            pending_user_message=self._state.pending_user_message,
            clear_pending_user_message=clear_pending_user_message,
        )
        changes: dict[str, object] = {"phase": ChatPhase.IDLE}
        if clear_pending_user_message:
            changes["pending_user_message"] = None
        self._update(**changes)

    def toggle_thinking_collapsed(self) -> None:
        """Toggle whether thinking blocks are collapsed."""
        self._update(thinking_collapsed=not self._state.thinking_collapsed)

    def set_thinking_collapsed(self, collapsed: bool) -> None:
        """Set whether thinking blocks are collapsed."""
        self._update(thinking_collapsed=collapsed)

    def clear_pending_user_message(self) -> None:
        """Clear the pending user message."""
        self._update(pending_user_message=None)

    def reset(self) -> None:
        """Reset to initial state."""
        self._state = ChatState()
        self._notify()

    # === UI State ===

    def toggle_sidebar(self) -> None:
        """Toggle sidebar visibility."""
        self._update(sidebar_visible=not self._state.sidebar_visible)

    def set_sidebar_visible(self, visible: bool) -> None:
        """Set sidebar visibility."""
        self._update(sidebar_visible=visible)

    def toggle_render_mode(self) -> None:
        """Toggle between markdown and plain text rendering."""
        current = self._state.render_mode
        new_mode = (
            RenderMode.PLAIN if current == RenderMode.MARKDOWN else RenderMode.MARKDOWN
        )
        logger.debug(
            "store_render_mode_toggled", old_mode=current.name, new_mode=new_mode.name
        )
        self._update(render_mode=new_mode)

    def set_render_mode(self, mode: RenderMode) -> None:
        """Set the render mode."""
        self._update(render_mode=mode)
