"""Textual UI adapter for agent execution."""
# ruff: noqa: PLR0912, PLR0915, ANN401, PLR2004, BLE001, TRY203
# This module has complex streaming logic ported from execution.py

from __future__ import annotations

import asyncio
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any

from langchain_core.messages import AIMessage, HumanMessage
from langgraph.types import Command

from ieum.file_ops import FileOpTracker
from ieum.image_utils import create_multimodal_content
from ieum.input import ImageTracker, parse_file_mentions
from ieum.observability import (
    create_trace_metadata,
    flush_langfuse_handler,
    get_langfuse_handler,
)
from ieum.stream_handler import BaseStreamHandler
from ieum.widgets.messages import (
    AssistantMessage,
    DiffMessage,
    SystemMessage,
    ToolCallMessage,
)

if TYPE_CHECKING:
    from collections.abc import Callable


class TokenExtractor:
    """Optimized token extraction with provider detection caching.

    Detects the provider format on first successful extraction and
    uses the fast path for subsequent extractions.
    """

    # Provider format types
    UNKNOWN = 0
    OPENAI = 1  # usage_metadata.total_tokens
    ANTHROPIC = 2  # usage_metadata.{input_tokens, output_tokens}
    OLLAMA = 3  # response_metadata.{prompt_eval_count, eval_count}
    GOOGLE = 4  # usage_metadata.{prompt_token_count, candidates_token_count}
    GROQ = 5  # response_metadata.usage.{prompt_tokens, completion_tokens}
    OPENROUTER = 6  # response_metadata.{prompt_tokens, completion_tokens}

    def __init__(self) -> None:
        self._detected_format = self.UNKNOWN

    def extract(self, message: Any) -> int:
        """Extract token usage from a message.

        Uses cached provider format for fast path if already detected.
        """
        # Fast path: use detected format
        if self._detected_format != self.UNKNOWN:
            return self._extract_fast(message)

        # Slow path: try all formats and cache the one that works
        return self._extract_and_detect(message)

    def _extract_fast(self, message: Any) -> int:
        """Fast extraction using cached provider format."""
        fmt = self._detected_format

        if fmt == self.OPENAI:
            if hasattr(message, "usage_metadata") and message.usage_metadata:
                return message.usage_metadata.get("total_tokens", 0)

        elif fmt == self.ANTHROPIC:
            if hasattr(message, "usage_metadata") and message.usage_metadata:
                usage = message.usage_metadata
                return usage.get("input_tokens", 0) + usage.get("output_tokens", 0)

        elif fmt == self.OLLAMA:
            if hasattr(message, "response_metadata") and message.response_metadata:
                resp = message.response_metadata
                return resp.get("prompt_eval_count", 0) + resp.get("eval_count", 0)

        elif fmt == self.GOOGLE:
            if hasattr(message, "usage_metadata") and message.usage_metadata:
                usage = message.usage_metadata
                return usage.get("prompt_token_count", 0) + usage.get("candidates_token_count", 0)

        elif fmt == self.GROQ:
            if hasattr(message, "response_metadata") and message.response_metadata:
                nested = message.response_metadata.get("usage", {})
                if isinstance(nested, dict):
                    return nested.get("prompt_tokens", 0) + nested.get("completion_tokens", 0)

        elif fmt == self.OPENROUTER:
            if hasattr(message, "response_metadata") and message.response_metadata:
                resp = message.response_metadata
                return resp.get("prompt_tokens", 0) + resp.get("completion_tokens", 0)

        return 0

    def _extract_and_detect(self, message: Any) -> int:
        """Try all formats and cache the one that works."""
        # Try OpenAI format (most common)
        if hasattr(message, "usage_metadata") and message.usage_metadata:
            usage = message.usage_metadata
            total = usage.get("total_tokens", 0)
            if total:
                self._detected_format = self.OPENAI
                return total

            # Anthropic style
            input_toks = usage.get("input_tokens", 0)
            output_toks = usage.get("output_tokens", 0)
            if input_toks or output_toks:
                self._detected_format = self.ANTHROPIC
                return input_toks + output_toks

            # Google style
            prompt_toks = usage.get("prompt_token_count", 0)
            candidates_toks = usage.get("candidates_token_count", 0)
            if prompt_toks or candidates_toks:
                self._detected_format = self.GOOGLE
                return prompt_toks + candidates_toks

        # Try response_metadata formats
        if hasattr(message, "response_metadata") and message.response_metadata:
            resp_meta = message.response_metadata

            # Ollama style
            prompt_toks = resp_meta.get("prompt_eval_count", 0)
            eval_toks = resp_meta.get("eval_count", 0)
            if prompt_toks or eval_toks:
                self._detected_format = self.OLLAMA
                return prompt_toks + eval_toks

            # Groq nested usage
            if "usage" in resp_meta:
                nested_usage = resp_meta["usage"]
                if isinstance(nested_usage, dict):
                    prompt_toks = nested_usage.get("prompt_tokens", 0)
                    completion_toks = nested_usage.get("completion_tokens", 0)
                    if prompt_toks or completion_toks:
                        self._detected_format = self.GROQ
                        return prompt_toks + completion_toks

            # OpenRouter style
            prompt_toks = resp_meta.get("prompt_tokens", 0)
            completion_toks = resp_meta.get("completion_tokens", 0)
            if prompt_toks or completion_toks:
                self._detected_format = self.OPENROUTER
                return prompt_toks + completion_toks

        # Direct token_usage attribute (rare)
        if hasattr(message, "token_usage") and message.token_usage:
            token_usage = message.token_usage
            if isinstance(token_usage, dict):
                total = token_usage.get("total_tokens", 0)
                if total:
                    return total
                return token_usage.get("prompt_tokens", 0) + token_usage.get("completion_tokens", 0)

        return 0

    def reset(self) -> None:
        """Reset provider detection (e.g., when switching models)."""
        self._detected_format = self.UNKNOWN


# Global token extractor instance (cached per session)
_token_extractor = TokenExtractor()


def _extract_token_usage(message: Any) -> int:
    """Extract token usage from a message in a provider-agnostic way.

    Uses cached provider format detection for optimized extraction.

    Supports multiple LLM providers with different response formats:
    - OpenAI: usage_metadata.total_tokens
    - Anthropic: usage_metadata.{input_tokens, output_tokens}
    - Ollama: response_metadata.{prompt_eval_count, eval_count}
    - Google: usage_metadata.{prompt_token_count, candidates_token_count}
    - Groq: response_metadata.usage.{prompt_tokens, completion_tokens}
    - OpenRouter: response_metadata.{prompt_tokens, completion_tokens}

    Args:
        message: The LLM message chunk to extract tokens from

    Returns:
        Total tokens (input + output) or 0 if not available
    """
    return _token_extractor.extract(message)


class TextualUIAdapter(BaseStreamHandler):
    """Adapter for rendering agent output to Textual widgets.

    Inherits from BaseStreamHandler to reuse shared event processing logic.
    Implements the on_* callbacks to render Textual UI widgets.
    """

    def __init__(
        self,
        mount_message: Callable,
        update_status: Callable[[str], None],
        request_approval_fn: Callable,  # async callable returning Future
        on_auto_approve_enabled: Callable[[], None] | None = None,
        scroll_to_bottom: Callable[[], None] | None = None,
        show_thinking: Callable[[], None] | None = None,
        hide_thinking: Callable[[], None] | None = None,
        request_interview_fn: Callable | None = None,  # async callable returning Future
    ) -> None:
        """Initialize the adapter.

        Args:
            mount_message: Async callable to mount a message widget
            update_status: Callable to update the status bar message
            request_approval_fn: Callable that returns a Future for HITL approval
            on_auto_approve_enabled: Callback when auto-approve is enabled
            scroll_to_bottom: Callback to scroll chat to bottom
            show_thinking: Callback to show/reposition thinking spinner
            hide_thinking: Callback to hide thinking spinner
            request_interview_fn: Callable that returns a Future for interview answers
        """
        self._mount_message = mount_message
        self._update_status = update_status
        self._request_approval_fn = request_approval_fn
        self._on_auto_approve_enabled = on_auto_approve_enabled
        self._scroll_to_bottom = scroll_to_bottom
        self._show_thinking = show_thinking
        self._hide_thinking = hide_thinking
        self._request_interview_fn = request_interview_fn

        # State tracking
        self._current_assistant_message: AssistantMessage | None = None
        self._current_tool_messages: dict[str, ToolCallMessage] = {}
        self._pending_text = ""
        self._token_tracker: Any = None
        self._thread_id: str | None = None

        # Track assistant messages per namespace for multi-subagent streaming
        self._assistant_message_by_namespace: dict[tuple, Any] = {}
        # Last completed assistant content (for plan capture)
        self._last_completed_content: str = ""

    # --- Legacy compatibility aliases ---
    @property
    def _request_approval(self):
        return self._request_approval_fn

    @property
    def _request_interview(self):
        return self._request_interview_fn

    def set_token_tracker(self, tracker: Any) -> None:
        """Set the token tracker for usage tracking."""
        self._token_tracker = tracker

    # --- BaseStreamHandler callback implementations ---

    async def on_chat_token(self, token: str, ns_key: tuple) -> None:
        # Get or create assistant message for this namespace
        current_msg = self._assistant_message_by_namespace.get(ns_key)
        if current_msg is None:
            # Hide thinking spinner when assistant starts responding
            if self._hide_thinking:
                await self._hide_thinking()
            current_msg = AssistantMessage()
            await self._mount_message(current_msg)
            self._assistant_message_by_namespace[ns_key] = current_msg

        # Append just the new text chunk for smoother streaming
        await current_msg.append_content(token)

        # Sticky scroll
        if self._scroll_to_bottom:
            self._scroll_to_bottom()

    async def on_tool_start(
        self, tool_id: str, tool_name: str, tool_input: dict, ns_key: tuple,
    ) -> None:
        # Hide thinking spinner before showing tool call
        if self._hide_thinking:
            await self._hide_thinking()

        # Mount tool call message
        tool_msg = ToolCallMessage(tool_name, tool_input)
        await self._mount_message(tool_msg)
        self._current_tool_messages[tool_id] = tool_msg

        # Sticky scroll
        if self._scroll_to_bottom:
            self._scroll_to_bottom()

    async def on_tool_end(
        self, tool_id: str, tool_name: str, output: str, status: str,
    ) -> None:
        # Reshow thinking spinner after tool result
        if self._show_thinking:
            await self._show_thinking()

        if tool_id in self._current_tool_messages:
            tool_msg = self._current_tool_messages[tool_id]
            if status == "success":
                tool_msg.set_success(output)
            else:
                tool_msg.set_error(output or "Error")
            del self._current_tool_messages[tool_id]

    async def on_file_changed(self, record: Any) -> None:
        if record.diff:
            await self._mount_message(
                DiffMessage(record.diff, record.display_path)
            )

    async def on_thinking_token(self, token: str) -> None:
        pass  # TUI does not render thinking tokens (uses spinner instead)

    async def on_thinking_end(self) -> None:
        pass  # TUI uses spinner, no explicit thinking end needed

    async def on_message_complete(self, ns_key: tuple, tokens: int) -> None:
        # Capture assistant content before flushing/popping
        msg_widget = self._assistant_message_by_namespace.get(ns_key)
        if msg_widget:
            content = getattr(msg_widget, "_content", "")
            if content:
                self._last_completed_content = content
        if msg_widget and self._thread_id:
            content = getattr(msg_widget, "_content", "")
            if content:
                from ieum.sessions import save_conversation_message
                asyncio.create_task(
                    save_conversation_message(self._thread_id, "assistant", content, tokens)
                )

        await _flush_assistant_text_ns(
            self, "", ns_key, self._assistant_message_by_namespace,
            tokens=tokens, force_stop=True,
        )
        self._assistant_message_by_namespace.pop(ns_key, None)

    async def on_human_message(self, ns_key: tuple) -> None:
        current_msg = self._assistant_message_by_namespace.get(ns_key)
        if current_msg is not None:
            await current_msg.stop_stream()
            self._assistant_message_by_namespace.pop(ns_key, None)

    async def on_stream_error(self, error: str) -> None:
        await self._mount_message(SystemMessage(f"Error: {error}"))

    async def on_interrupted(self) -> None:
        await self._mount_message(SystemMessage("Interrupted by user"))

    async def on_tool_rejected_cleanup(self) -> None:
        for tool_msg in self._current_tool_messages.values():
            tool_msg.set_rejected()
        self._current_tool_messages.clear()

    async def on_rejection_message(self) -> None:
        await self._mount_message(
            SystemMessage("Command rejected. Tell the agent what you'd like instead.")
        )

    async def request_approval(
        self, action_requests: list[dict], assistant_id: str | None,
    ) -> asyncio.Future:
        return await self._request_approval_fn(action_requests, assistant_id)

    async def request_interview(
        self, interview_id: str, questions: list[dict],
    ) -> asyncio.Future:
        if self._request_interview_fn:
            return await self._request_interview_fn(questions)
        # No interview handler - return resolved future with empty answers
        loop = asyncio.get_running_loop()
        fut: asyncio.Future = loop.create_future()
        fut.set_result({})
        return fut

    async def on_approval_auto_approved(self) -> None:
        if self._on_auto_approve_enabled:
            self._on_auto_approve_enabled()

    async def on_pre_stream(self) -> None:
        if self._show_thinking:
            await self._show_thinking()
        if self._token_tracker:
            self._token_tracker.hide()

    async def on_post_stream(self, tokens: int) -> None:
        if self._token_tracker:
            if tokens:
                breakdown = getattr(self, "_token_breakdown", {})
                self._token_tracker.add(
                    tokens,
                    0,
                    input_tokens=breakdown.get("input_tokens", 0),
                    output_tokens=breakdown.get("output_tokens", 0),
                )
            else:
                self._token_tracker.show()

    async def on_interview_complete(self, interview_id: str) -> None:
        pass  # TUI does not need explicit interview completion notification

    async def on_tool_call_flush_text(self, ns_key: tuple) -> None:
        current_msg = self._assistant_message_by_namespace.get(ns_key)
        if current_msg is not None:
            await current_msg.stop_stream()
            self._assistant_message_by_namespace.pop(ns_key, None)

    async def on_tool_running(self) -> None:
        for tool_msg in self._current_tool_messages.values():
            tool_msg.set_running()


def _build_interrupted_ai_message(
    pending_text_by_namespace: dict[tuple, str],
    current_tool_messages: dict[str, Any],
) -> AIMessage | None:
    """Build an AIMessage capturing interrupted state (text + tool calls).

    Args:
        pending_text_by_namespace: Dict of accumulated text by namespace
        current_tool_messages: Dict of tool_id -> ToolCallMessage widget

    Returns:
        AIMessage with accumulated content and tool calls, or None if empty
    """
    main_ns_key = ()
    accumulated_text = pending_text_by_namespace.get(main_ns_key, "").strip()

    # Reconstruct tool_calls from displayed tool messages
    tool_calls = []
    for tool_id, tool_widget in current_tool_messages.items():
        tool_calls.append(
            {
                "id": tool_id,
                "name": tool_widget._tool_name,
                "args": tool_widget._args,
            }
        )

    if not accumulated_text and not tool_calls:
        return None

    return AIMessage(
        content=accumulated_text,
        tool_calls=tool_calls if tool_calls else [],
    )


async def execute_task_textual(
    user_input: str,
    agent: Any,
    assistant_id: str | None,
    session_state: Any,
    adapter: TextualUIAdapter,
    backend: Any = None,
    image_tracker: ImageTracker | None = None,
) -> None:
    """Execute a task with output directed to Textual UI.

    Uses BaseStreamHandler.process_stream for the core event loop,
    delegating to the TextualUIAdapter callbacks for TUI-specific rendering.

    Args:
        user_input: The user's input message
        agent: The LangGraph agent to execute
        assistant_id: The agent identifier
        session_state: Session state with auto_approve flag
        adapter: The TextualUIAdapter for UI operations
        backend: Optional backend for file operations
        image_tracker: Optional tracker for images
    """
    # Parse file mentions and inject content if any
    prompt_text, mentioned_files = parse_file_mentions(user_input)

    # Max file size to embed inline (256KB, matching mistral-vibe)
    # Larger files get a reference instead - use read_file tool to view them
    max_embed_bytes = 256 * 1024

    if mentioned_files:
        context_parts = [prompt_text, "\n\n## Referenced Files\n"]
        for file_path in mentioned_files:
            try:
                file_size = file_path.stat().st_size
                if file_size > max_embed_bytes:
                    # File too large - include reference instead of content
                    size_kb = file_size // 1024
                    context_parts.append(
                        f"\n### {file_path.name}\n"
                        f"Path: `{file_path}`\n"
                        f"Size: {size_kb}KB (too large to embed, use read_file tool to view)"
                    )
                else:
                    content = file_path.read_text()
                    context_parts.append(
                        f"\n### {file_path.name}\nPath: `{file_path}`\n```\n{content}\n```"
                    )
            except Exception as e:
                context_parts.append(f"\n### {file_path.name}\n[Error reading file: {e}]")
        final_input = "\n".join(context_parts)
    else:
        final_input = prompt_text

    # Include images in the message content
    images_to_send = []
    if image_tracker:
        images_to_send = image_tracker.get_images()
    if images_to_send:
        message_content = create_multimodal_content(final_input, images_to_send)
    else:
        message_content = final_input

    thread_id = session_state.thread_id
    adapter._thread_id = thread_id
    configurable: dict[str, str] = {"thread_id": thread_id}

    # Support checkpoint restore: if a restore checkpoint_id is set, include it
    # so LangGraph resumes from that specific checkpoint
    restore_cp = getattr(session_state, "_restore_checkpoint_id", None)
    if restore_cp:
        configurable["checkpoint_id"] = restore_cp
        session_state._restore_checkpoint_id = None  # consume once

    config = {
        "configurable": configurable,
        "metadata": {
            "assistant_id": assistant_id,
            "agent_name": assistant_id,
            "updated_at": datetime.now(UTC).isoformat(),
        }
        if assistant_id
        else {},
    }

    # Add Langfuse callback for observability if configured
    langfuse_handler = get_langfuse_handler(
        session_id=thread_id,
        metadata=create_trace_metadata(
            assistant_id=assistant_id or "unknown",
            model_name=None,  # Model name not available here
        ),
    )

    if langfuse_handler:
        config["callbacks"] = [langfuse_handler]

    captured_input_tokens = 0

    # Show thinking spinner and hide token display
    await adapter.on_pre_stream()

    # Auto-commit callback
    def _auto_commit_hook(record: Any) -> None:
        if not getattr(session_state, "auto_commit", False):
            return
        try:
            from ieum.git.auto_commit import AutoCommitManager
            path_str = record.args.get("file_path") or record.args.get("path") or ""
            if not path_str:
                return
            mgr = AutoCommitManager(repo_path=session_state._cwd if hasattr(session_state, "_cwd") else None)
            mgr.stage_files([path_str])
            mgr.commit()
        except Exception:
            pass  # Auto-commit is best-effort

    file_op_tracker = FileOpTracker(
        assistant_id=assistant_id,
        backend=backend,
        on_write_complete=_auto_commit_hook,
    )

    # Clear images from tracker after creating the message
    if image_tracker:
        image_tracker.clear()

    stream_input: dict | Command = {"messages": [{"role": "user", "content": message_content}]}

    def on_auto_approve():
        session_state.auto_approve = True

    def on_mark_hitl_approved(action_request: dict) -> None:
        tool_name = action_request.get("name")
        if tool_name in {"write_file", "edit_file"}:
            args = action_request.get("args", {})
            if isinstance(args, dict):
                file_op_tracker.mark_hitl_approved(tool_name, args)

    try:
        should_continue = True
        current_input: Any = stream_input
        while should_continue:
            should_continue, next_input, tokens = await adapter.process_stream(
                agent=agent,
                stream_input=current_input,
                config=config,
                file_op_tracker=file_op_tracker,
                session_auto_approve=session_state.auto_approve,
                assistant_id=assistant_id,
                approval_policy=getattr(session_state, "approval_policy", "manual"),
                on_auto_approve_enabled=on_auto_approve,
                on_mark_hitl_approved=on_mark_hitl_approved,
            )
            captured_input_tokens = max(captured_input_tokens, tokens)
            if next_input is not None:
                current_input = next_input

    except asyncio.CancelledError:
        await adapter._mount_message(SystemMessage("Interrupted by user"))

        # Save accumulated state before marking tools as rejected
        try:
            # Build pending text dict from accumulated assistant message content
            pending_text: dict[tuple, str] = {}
            for ns_key, msg_widget in adapter._assistant_message_by_namespace.items():
                content = getattr(msg_widget, "_content", "")
                if content:
                    pending_text[ns_key] = content
            interrupted_msg = _build_interrupted_ai_message(
                pending_text,
                adapter._current_tool_messages,
            )
            if interrupted_msg:
                await agent.aupdate_state(config, {"messages": [interrupted_msg]})

            cancellation_msg = HumanMessage(
                content="[SYSTEM] Task interrupted by user. Previous operation was cancelled."
            )
            await agent.aupdate_state(config, {"messages": [cancellation_msg]})
        except Exception:  # noqa: S110
            pass  # State update is best-effort

        # Mark tools as rejected AFTER saving state
        for tool_msg in list(adapter._current_tool_messages.values()):
            tool_msg.set_rejected()
        adapter._current_tool_messages.clear()

        # Report tokens even on interrupt
        await adapter.on_post_stream(captured_input_tokens)

        # Flush Langfuse handler on interrupt
        flush_langfuse_handler(langfuse_handler)
        return

    except KeyboardInterrupt:
        await adapter._mount_message(SystemMessage("Interrupted by user"))

        # Save accumulated state before marking tools as rejected
        try:
            pending_text_kb: dict[tuple, str] = {}
            for ns_key, msg_widget in adapter._assistant_message_by_namespace.items():
                content = getattr(msg_widget, "_content", "")
                if content:
                    pending_text_kb[ns_key] = content
            interrupted_msg = _build_interrupted_ai_message(
                pending_text_kb,
                adapter._current_tool_messages,
            )
            if interrupted_msg:
                await agent.aupdate_state(config, {"messages": [interrupted_msg]})

            cancellation_msg = HumanMessage(
                content="[SYSTEM] Task interrupted by user. Previous operation was cancelled."
            )
            await agent.aupdate_state(config, {"messages": [cancellation_msg]})
        except Exception:  # noqa: S110
            pass  # State update is best-effort

        # Mark tools as rejected AFTER saving state
        for tool_msg in list(adapter._current_tool_messages.values()):
            tool_msg.set_rejected()
        adapter._current_tool_messages.clear()

        # Report tokens even on interrupt
        await adapter.on_post_stream(captured_input_tokens)

        # Flush Langfuse handler on interrupt
        flush_langfuse_handler(langfuse_handler)
        return

    # Update token tracker
    await adapter.on_post_stream(captured_input_tokens)

    # Flush Langfuse handler to ensure all traces are sent
    flush_langfuse_handler(langfuse_handler)


async def _flush_assistant_text_ns(
    adapter: TextualUIAdapter,
    text: str,
    ns_key: tuple,
    assistant_message_by_namespace: dict[tuple, Any],
    tokens: int = 0,
    force_stop: bool = False,
) -> None:
    """Flush accumulated assistant text for a specific namespace.

    Finalizes the streaming by stopping the MarkdownStream.
    If no message exists yet, creates one with the full content.

    Args:
        adapter: The UI adapter
        text: The accumulated text
        ns_key: The namespace key
        assistant_message_by_namespace: Dict tracking messages by namespace
        tokens: Token count to display on the message
        force_stop: If True, stop stream even if text is empty
    """
    current_msg = assistant_message_by_namespace.get(ns_key)

    if current_msg is None and not text.strip():
        return

    if current_msg is None:
        # No message was created during streaming - create one with full content
        current_msg = AssistantMessage(text)
        await adapter._mount_message(current_msg)
        await current_msg.write_initial_content()
        assistant_message_by_namespace[ns_key] = current_msg
    elif force_stop or text.strip():
        # Stop the stream to finalize the content
        await current_msg.stop_stream()

    # Set token count on the message
    if tokens > 0:
        current_msg.set_tokens(tokens)
