"""In-memory model for the chat panel.

Converts raw WebSocket events into logical chat messages that renderers understand.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Literal


@dataclass
class ChatItem:
    """Logical item shown in the chat history."""

    kind: Literal[
        "user_message",
        "assistant_message",
        "thinking_block",
        "tool_started",
        "tool_call",
        "error",
        "status",
        "planning_status",
        "plan_ready",
        "turn_meta",
        "context_summary",
        "recall",
    ]
    content: str = ""
    meta: dict = field(default_factory=dict)
    id: str = ""


class ChatModel:
    """Accumulate stream deltas and tool events into ChatItems."""

    def __init__(self) -> None:
        self.items: list[ChatItem] = []
        self._current_assistant: ChatItem | None = None
        self._current_thinking: ChatItem | None = None

    def add_user_message(self, text: str) -> None:
        self.items.append(ChatItem(kind="user_message", content=text))
        self._current_assistant = None

    def load_history(self, messages: list[dict]) -> None:
        """Replace the item list with the session's stored messages.

        Maps persisted Message dicts (as returned by GET /sessions/{id}) back
        into the same ChatItems the live stream produces, so a resumed session
        shows its past conversation. Display-only context markers
        (is_display=False: the context-only user message, compressor summaries,
        compression events) are skipped — only what the user would have seen
        live is rebuilt.
        """
        self.items.clear()
        self._current_assistant = None
        self._current_thinking = None

        # tool_call_id -> tool arguments, so a later ``role=="tool"`` message can
        # recover the args the assistant sent (the persisted tool message only
        # carries the result, not the call args). Mirrors the live path where the
        # tool_call WS event's meta already contains ``args``.
        args_by_id: dict[str, dict] = {}

        for m in messages:
            if not m.get("is_display", True):
                continue
            role = m.get("role")
            if role == "user":
                self.items.append(ChatItem(kind="user_message", content=m.get("content") or ""))
            elif role == "assistant":
                thinking = m.get("thinking")
                if thinking:
                    self.items.append(ChatItem(kind="thinking_block", content=thinking))
                content = m.get("content") or ""
                if m.get("is_plan"):
                    self.items.append(ChatItem(kind="plan_ready", content=content, meta={"is_subagent": False}))
                elif content:
                    # Match the live stream: stream_delta (assistant text) arrives
                    # BEFORE tool_started, so the text bubble sits above the tool
                    # cards, not below them. Emitting tool_started first would strand
                    # the answer under its own tool calls on resume.
                    self.items.append(ChatItem(kind="assistant_message", content=content))
                for tc in m.get("tool_calls") or []:
                    tcid = tc.get("id")
                    if tcid:
                        args_by_id[tcid] = tc.get("arguments") or {}
                    self.items.append(
                        ChatItem(
                            kind="tool_started",
                            meta={"tool": tc.get("name", ""), "args": tc.get("arguments") or {}},
                        )
                    )
                # The whole-turn duration is stamped only on the assistant message
                # that closed the run (intermediate tool-tour assistant messages get
                # no elapsed_seconds), so this reproduces the single turn_meta line
                # the live stream_end appends at the end of a turn.
                elapsed = m.get("elapsed_seconds")
                if elapsed is not None:
                    self.items.append(ChatItem(kind="turn_meta", meta={"elapsed_seconds": elapsed}))
            elif role == "tool":
                tcid = m.get("tool_call_id")
                args = args_by_id.pop(tcid, {}) if tcid else {}
                self.items.append(
                    ChatItem(
                        kind="tool_call",
                        meta={
                            "tool": m.get("name") or "",
                            "args": args,
                            "result": m.get("content") or "",
                            "success": True,
                            "metadata": m.get("metadata") or {},
                        },
                    )
                )

    def handle_ws_event(self, msg: dict) -> ChatItem | None:
        msg_type = msg.get("type")

        if msg_type == "stream_start":
            # A new turn starts. Do NOT eagerly create an assistant bubble here —
            # that would place the (possibly empty) answer at the TOP, above the
            # tools/thinking that follow, where scroll_end hides it. Instead reset
            # the current-assistant pointer and let the first stream_delta create
            # the bubble lazily, at the position where text actually arrives.
            # Reset the thinking pointer too: if the previous turn broke without a
            # ``thinking_end`` (dropped socket, stop without the event), a dangling
            # ``_current_thinking`` would make the new turn's first ``thinking_delta``
            # append to the previous turn's reasoning block, gluing two turns'
            # reasoning together.
            self._current_assistant = None
            self._current_thinking = None
            return None

        if msg_type == "thinking_delta":
            if self._current_thinking is None:
                self._current_thinking = ChatItem(kind="thinking_block", content="")
                self.items.append(self._current_thinking)
            self._current_thinking.content += msg.get("delta", "")
            return None

        if msg_type == "thinking_end":
            self._current_thinking = None
            return None

        if msg_type == "turn_thinking":
            # A complete thinking block from a (sub-agent) non-streaming turn.
            # Stored as a thinking_block with an is_subagent flag so the renderer
            # can style/indent it as nested sub-agent reasoning.
            item = ChatItem(
                kind="thinking_block",
                content=msg.get("thinking", "") or "",
                meta={"is_subagent": bool(msg.get("is_subagent", False))},
            )
            self.items.append(item)
            return item

        if msg_type == "stream_delta":
            if self._current_assistant is None:
                self._current_assistant = ChatItem(kind="assistant_message", content="")
                self.items.append(self._current_assistant)
            self._current_assistant.content += msg.get("delta", "")
            return None

        if msg_type == "tool_started":
            # Drop the current assistant pointer so the next stream_delta opens a
            # fresh bubble BELOW the tool cards (the final answer ends up at the
            # bottom, visible after scroll_end — not stranded at the top).
            self._current_assistant = None
            item = ChatItem(kind="tool_started", meta=msg)
            self.items.append(item)
            return item

        if msg_type == "tool_call":
            self._current_assistant = None
            item = ChatItem(kind="tool_call", meta=msg)
            self.items.append(item)
            return item

        if msg_type == "error":
            item = ChatItem(kind="error", content=msg.get("message", ""))
            self.items.append(item)
            return item

        if msg_type == "status":
            item = ChatItem(kind="status", content=msg.get("content", ""))
            self.items.append(item)
            return item

        if msg_type == "planning_status":
            label = msg.get("label", "") or ""
            is_subagent = bool(msg.get("is_subagent", False))
            phase = msg.get("phase")
            # Non-subagent planning is a transient rolling indicator: update the
            # last pending planning line in place so the phases (Analysis →
            # Execution plan → Plan review) share one line instead of stacking.
            # Subagent planning is appended as its own dim line.
            if not is_subagent and self.items and self.items[-1].kind == "planning_status" and not bool(
                self.items[-1].meta.get("is_subagent", False)
            ):
                self.items[-1].content = label
                self.items[-1].meta = {"phase": phase, "is_subagent": False}
                return self.items[-1]
            item = ChatItem(
                kind="planning_status",
                content=label,
                meta={"phase": phase, "is_subagent": is_subagent},
            )
            self.items.append(item)
            return item

        if msg_type == "plan_ready":
            plan = msg.get("plan", "") or ""
            is_subagent = bool(msg.get("is_subagent", False))
            # The transient non-subagent planning indicator is consumed by the
            # finalized plan card.
            if not is_subagent and self.items and self.items[-1].kind == "planning_status" and not bool(
                self.items[-1].meta.get("is_subagent", False)
            ):
                self.items.pop()
            item = ChatItem(
                kind="plan_ready",
                content=plan,
                meta={"is_subagent": is_subagent},
            )
            self.items.append(item)
            return item

        if msg_type == "stream_stopped":
            self._current_assistant = None
            self._current_thinking = None
            item = ChatItem(kind="status", content="Generation stopped by user")
            self.items.append(item)
            return item

        if msg_type == "stream_end":
            # Purge assistant/thinking bubbles that never received any text so the
            # conversation does not show empty "Navi" panels.
            self.items = [
                it
                for it in self.items
                if not (
                    it.kind in ("assistant_message", "thinking_block") and not it.content
                )
            ]
            # Record the total time the agent spent on this request (all steps)
            # as a dim metadata line below the answer. elapsed_seconds is the
            # authoritative value measured by the backend over the whole turn.
            self.items.append(
                ChatItem(kind="turn_meta", meta={"elapsed_seconds": msg.get("elapsed_seconds")})
            )
            self._current_assistant = None
            self._current_thinking = None
            return None

        if msg_type == "context_compressed":
            # The compressor replaced the discarded turns with a summary. Surface
            # that text in the chat as a distinct block so the user can see what
            # was kept — for both forced /compact and mid-turn auto-compress.
            item = ChatItem(
                kind="context_summary",
                content=msg.get("summary", "") or "",
                meta={
                    "messages_before": msg.get("messages_before"),
                    "messages_after": msg.get("messages_after"),
                },
            )
            self.items.append(item)
            return item

        if msg_type in ("heartbeat", "session_sync"):
            return None

        if msg_type == "recall_update":
            # schedule_recall lifecycle (scheduled/fired/cancelled/skipped/
            # rescheduled), pushed out-of-band by the scheduler/orchestrator.
            item = ChatItem(kind="recall", meta=msg)
            self.items.append(item)
            return item

        # Unknown event — store as status for debugging.
        item = ChatItem(kind="status", content=f"{msg_type}: {msg}")
        self.items.append(item)
        return item
