"""In-flight chat-turn state — the foundation for the "is the agent still
working?" UX that every mainstream chat app has and AgntSpace historically
lacked.

Before this module existed, a user who closed the chat tab mid-stream came
back to a conversation that looked dead: the WebSocket had dropped, the
streamed reply was gone (no incremental persist), and the UI had no way to
ask "hey, is the agent still doing something for this conversation?".

The TurnRegistry is a process-local dict keyed by conversation id. chat_stream
opens a ``TurnState`` entry at turn start, mutates it through the phases
(``thinking`` → ``tool`` → ``streaming`` → ``done``), and removes it on
completion. A new ``GET /api/conversations/{id}/status`` route reads the
entry and returns it to the UI so the chat page can render a "still working"
indicator + current tool name.

State is deliberately NOT persisted to disk:
  - Turn state is ephemeral — if the server restarts, in-flight turns are
    gone and the client should see a "did not complete" sentinel.
  - No cross-restart guarantee is needed; the value is giving the UI a
    live signal while the turn is running.

Thread-safety: all chat turns run in the asyncio event loop so there is a
single implicit lock. No explicit asyncio.Lock needed — dict mutations
within a single task are atomic.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from datetime import UTC, datetime

log = logging.getLogger(__name__)


def _now_iso() -> str:
    return datetime.now(UTC).isoformat()


@dataclass
class TurnState:
    """Live state of an in-flight chat turn.

    Mutated by chat_stream as the turn progresses; read by the status
    endpoint. ``message_id`` is set once the draft assistant row is
    created (during phase=streaming) so the UI can re-hydrate partial
    content. ``tools_attempted`` and ``tools_succeeded`` feed the
    progress indicator.
    """

    conversation_id: str
    started_at: str = field(default_factory=_now_iso)
    phase: str = "thinking"          # thinking | tool | streaming | done | error | cancelled
    message_id: str | None = None    # draft assistant message id once streaming begins
    last_tool_name: str | None = None
    tools_attempted: list[str] = field(default_factory=list)
    tools_succeeded: int = 0
    tools_failed: int = 0
    last_chunk_at: str | None = None
    round_index: int = 0              # which tool-loop round we're in
    error: str | None = None

    def to_dict(self) -> dict:
        return {
            "conversation_id": self.conversation_id,
            "started_at": self.started_at,
            "phase": self.phase,
            "message_id": self.message_id,
            "last_tool_name": self.last_tool_name,
            "tools_attempted": list(self.tools_attempted),
            "tools_succeeded": self.tools_succeeded,
            "tools_failed": self.tools_failed,
            "last_chunk_at": self.last_chunk_at,
            "round_index": self.round_index,
            "error": self.error,
        }


class TurnRegistry:
    """Per-conversation in-flight turn tracker.

    Usage pattern::

        turn = registry.start_turn(conv_id)
        try:
            turn.phase = "thinking"
            ... complete_with_tools ...
            turn.phase = "tool"
            turn.note_tool("create_kb", succeeded=True)
            turn.phase = "streaming"
            turn.message_id = draft_mid
            ... stream chunks ...
            turn.phase = "done"
        finally:
            registry.end_turn(conv_id)

    Safe to call ``end_turn`` even if the conversation was never
    registered (no-op). Safe to call ``get_status`` on a conversation
    that never had a turn (returns None).
    """

    def __init__(self) -> None:
        self._turns: dict[str, TurnState] = {}

    def start_turn(self, conversation_id: str) -> TurnState:
        """Begin a new turn. Replaces any existing turn for this conv.

        Typical usage is one turn per conversation at a time — chat_stream
        is single-threaded per WS connection. If an old turn is still
        registered (previous request that errored out without cleaning
        up), we overwrite it. Log at debug so unexpected collisions are
        visible without spamming the log.
        """
        if conversation_id in self._turns:
            log.debug(
                "TurnRegistry: replacing stale turn for %s (phase=%s)",
                conversation_id, self._turns[conversation_id].phase,
            )
        state = TurnState(conversation_id=conversation_id)
        self._turns[conversation_id] = state
        return state

    def get_status(self, conversation_id: str) -> TurnState | None:
        return self._turns.get(conversation_id)

    def active_conversations(self) -> list[str]:
        """Return the conversation ids with an in-flight turn. Used by
        the /api/conversations list endpoint to tag rows that are
        currently being processed."""
        return list(self._turns.keys())

    def note_tool(
        self, conversation_id: str, tool_name: str, *, succeeded: bool,
    ) -> None:
        """Record a tool call on the turn state. No-op if the turn is
        gone (e.g. end_turn already called because of a race).
        """
        state = self._turns.get(conversation_id)
        if state is None:
            return
        state.tools_attempted.append(tool_name)
        state.last_tool_name = tool_name
        if succeeded:
            state.tools_succeeded += 1
        else:
            state.tools_failed += 1

    def note_chunk(self, conversation_id: str) -> None:
        """Bump ``last_chunk_at`` so the UI can detect a stuck stream
        (no chunk in the last N seconds → show a "slow" hint).
        """
        state = self._turns.get(conversation_id)
        if state is not None:
            state.last_chunk_at = _now_iso()

    def set_phase(
        self,
        conversation_id: str,
        phase: str,
        *,
        message_id: str | None = None,
    ) -> None:
        """Advance the phase. Optionally attach a draft message id
        (streaming phase always sets this so the UI can rehydrate).
        """
        state = self._turns.get(conversation_id)
        if state is None:
            return
        state.phase = phase
        if message_id is not None:
            state.message_id = message_id

    def set_error(self, conversation_id: str, error: str) -> None:
        state = self._turns.get(conversation_id)
        if state is None:
            return
        state.phase = "error"
        state.error = error[:500]

    def end_turn(self, conversation_id: str) -> None:
        """Remove the turn from the registry. Idempotent — safe to call
        from a finally block whether or not start_turn fired.
        """
        self._turns.pop(conversation_id, None)
