"""Textual UI adapter — consumes session events and renders widgets."""

from __future__ import annotations

import asyncio
import logging
from collections.abc import AsyncGenerator, Awaitable, Callable
from typing import Any

from icarus.session.events import (
    AutoApproveEnabled,
    SessionEvent,
    StatusUpdate,
    TextComplete,
    TextDelta,
    TokenUsage,
    ToolCallStart,
    ToolResult,
    ToolsApproved,
    ToolsRejected,
)
from icarus_tui.widgets.messages import (
    AppMessage,
    AssistantMessage,
    DiffMessage,
    ToolCallMessage,
)

logger = logging.getLogger(__name__)


class TextualUIAdapter:
    """Maps :class:`SessionEvent` objects to Textual widgets.

    Constructed with the same callbacks as before; the only change is that
    the 700-line ``execute_task_textual`` function is gone — replaced by
    :meth:`run`, which iterates over a ``stream_turn()`` async generator.
    """

    def __init__(
        self,
        mount_message: Callable[..., Awaitable[None]],
        on_auto_approve_enabled: Callable[[], None] | None = None,
        scroll_to_bottom: Callable[[], None] | None = None,
        set_spinner: Callable[[str | None], Awaitable[None]] | None = None,
        set_active_message: Callable[[str | None], None] | None = None,
        sync_message_content: Callable[[str, str], None] | None = None,
    ) -> None:
        self._mount_message = mount_message
        self._on_auto_approve_enabled = on_auto_approve_enabled
        self._scroll_to_bottom = scroll_to_bottom
        self._set_spinner = set_spinner
        self._set_active_message = set_active_message
        self._sync_message_content = sync_message_content

        self._current_tool_messages: dict[str, ToolCallMessage] = {}
        self._current_assistant_msg: AssistantMessage | None = None
        self._token_tracker: Any = None
        self._tokens_updated: bool = False

    # -- public helpers (unchanged API) --

    def set_token_tracker(self, tracker: Any) -> None:  # noqa: ANN401
        self._token_tracker = tracker

    def finalize_pending_tools_with_error(self, error: str) -> None:
        """Mark all pending tool widgets as error and clear tracking."""
        for tool_msg in list(self._current_tool_messages.values()):
            tool_msg.set_error(error)
        self._current_tool_messages.clear()
        if self._set_active_message:
            self._set_active_message(None)

    # -- main entry point --

    async def run(
        self,
        event_stream: AsyncGenerator[SessionEvent, None],
    ) -> None:
        """Consume *event_stream* and render each event as a Textual widget."""
        self._tokens_updated = False
        if self._token_tracker:
            self._token_tracker.hide()

        try:
            async for event in event_stream:
                await self._handle_event(event)
        except (asyncio.CancelledError, KeyboardInterrupt):
            if self._set_active_message:
                self._set_active_message(None)
            await self._mount_message(AppMessage("Interrupted by user"))
            for tool_msg in list(self._current_tool_messages.values()):
                tool_msg.set_rejected()
            self._current_tool_messages.clear()
        finally:
            # Restore token display — either updated via TokenUsage event
            # during _handle_event, or restored to previous value if no
            # tokens were captured (e.g. cancelled before LLM responded).
            if self._token_tracker and not self._tokens_updated:
                self._token_tracker.show()

    # -- event dispatch --

    async def _handle_event(self, event: SessionEvent) -> None:  # noqa: C901, PLR0912
        match event:
            case StatusUpdate(status=status):
                if self._set_spinner:
                    await self._set_spinner(status)

            case TextDelta(text=text, message_id=msg_id):
                if self._current_assistant_msg is None:
                    if self._set_active_message:
                        self._set_active_message(msg_id)
                    self._current_assistant_msg = AssistantMessage(id=msg_id)
                    await self._mount_message(self._current_assistant_msg)
                await self._current_assistant_msg.append_content(text)
                if self._scroll_to_bottom:
                    self._scroll_to_bottom()

            case TextComplete(full_text=full_text, message_id=msg_id):
                if self._current_assistant_msg is None:
                    msg = AssistantMessage(full_text, id=msg_id)
                    await self._mount_message(msg)
                    await msg.write_initial_content()
                    self._current_assistant_msg = msg
                else:
                    await self._current_assistant_msg.stop_stream()
                if (
                    self._sync_message_content
                    and self._current_assistant_msg
                    and self._current_assistant_msg.id
                ):
                    self._sync_message_content(
                        self._current_assistant_msg.id,
                        self._current_assistant_msg._content,
                    )
                if self._set_active_message:
                    self._set_active_message(None)
                self._current_assistant_msg = None

            case ToolCallStart(tool_id=tid, tool_name=name, args=args):
                tool_msg = ToolCallMessage(name, args)
                await self._mount_message(tool_msg)
                self._current_tool_messages[tid] = tool_msg
                if self._scroll_to_bottom:
                    self._scroll_to_bottom()

            case ToolResult(
                tool_id=tid, status=status, output=output,
                diff=diff, diff_path=diff_path,
            ):
                if tid in self._current_tool_messages:
                    tool_msg = self._current_tool_messages[tid]
                    if status == "success":
                        tool_msg.set_success(output)
                    else:
                        tool_msg.set_error(output or "Error")
                    self._current_tool_messages.pop(tid, None)
                if diff:
                    await self._mount_message(DiffMessage(diff, diff_path))

            case ToolsApproved():
                for tool_msg in self._current_tool_messages.values():
                    tool_msg.set_running()

            case ToolsRejected(message=msg):
                for tool_msg in list(self._current_tool_messages.values()):
                    tool_msg.set_rejected()
                self._current_tool_messages.clear()
                await self._mount_message(AppMessage(msg))

            case AutoApproveEnabled():
                if self._on_auto_approve_enabled:
                    self._on_auto_approve_enabled()

            case TokenUsage(total_tokens=total):
                if self._token_tracker:
                    self._token_tracker.add(total, 0)
                    self._tokens_updated = True
