"""Message handling for chat sessions."""

import asyncio
from pathlib import Path
from typing import Any

from langchain_core.messages import AnyMessage, HumanMessage
from langchain_core.runnables import RunnableConfig
from langgraph.types import Command, Interrupt

from src.agents.context import AgentContext
from src.cli.bootstrap.initializer import initializer
from src.cli.builders import MessageContentBuilder
from src.cli.handlers import CompressionHandler, InterruptHandler
from src.cli.theme import console, theme
from src.core.logging import get_logger
from src.utils.compression import should_auto_compress

logger = get_logger(__name__)


class MessageDispatcher:
    """Dispatch user message processing and AI response streaming."""

    def __init__(self, session):
        """Initialize with reference to CLI session."""
        self.session = session
        self.interrupt_handler = InterruptHandler()
        self.message_builder = MessageContentBuilder(Path(session.context.working_dir))

    async def dispatch(self, content: str) -> None:
        """Dispatch user message and get AI response."""
        try:
            reference_mapping = self.session.prefilled_reference_mapping.copy()
            self.session.prefilled_reference_mapping.clear()

            message_content, image_refs = self.message_builder.build(content)

            reference_mapping.update(image_refs)

            human_message = HumanMessage(
                content=message_content,
                short_content=content,
                additional_kwargs={"reference_mapping": reference_mapping},
            )
            ctx = self.session.context
            agent_context = AgentContext(
                approval_mode=ctx.approval_mode,
                working_dir=ctx.working_dir,
                input_cost_per_mtok=ctx.input_cost_per_mtok,
                output_cost_per_mtok=ctx.output_cost_per_mtok,
                tool_output_max_tokens=ctx.tool_output_max_tokens,
            )

            graph_config = RunnableConfig(
                configurable={"thread_id": ctx.thread_id},
                recursion_limit=ctx.recursion_limit,
            )

            await self._stream_response(
                {"messages": [human_message]},
                graph_config,
                agent_context,
            )

        except Exception as e:
            console.print_error(f"Error processing message: {e}")
            logger.debug("Message processing error", exc_info=True)

    async def _stream_response(
        self,
        input_data: dict[str, Any],
        config: RunnableConfig,
        context: AgentContext,
    ) -> None:
        """Stream with automatic interrupt handling loop."""
        current_input: dict[str, Any] | Command = input_data
        rendered_messages: set[str] = set()

        while True:
            interrupted = False
            cancelled = False

            try:
                with console.console.status(
                    f"[{theme.spinner_color}]Randomizing...[/{theme.spinner_color}]"
                ) as status:
                    async for chunk in self.session.graph.astream(
                        current_input,
                        config,
                        context=context,
                        stream_mode="updates",
                        subgraphs=True,
                    ):
                        interrupts = self._extract_interrupts(chunk)
                        if interrupts:
                            status.stop()
                            resume_value = await self.interrupt_handler.handle(
                                interrupts
                            )
                            # Map resume value to interrupt ID format
                            if isinstance(resume_value, dict):
                                current_input = Command(resume=resume_value)
                            else:
                                current_input = Command(
                                    resume={interrupts[0].id: resume_value}
                                )
                            interrupted = True
                            break
                        else:
                            await self._process_chunk(chunk, rendered_messages)

            except (asyncio.CancelledError, KeyboardInterrupt):
                # User cancelled the generation
                cancelled = True

            if cancelled or not interrupted:
                # Cancelled or no interrupts encountered, streaming completed
                break

    @staticmethod
    def _extract_interrupts(chunk) -> list[Interrupt] | None:
        """Extract interrupt data from chunk."""
        if isinstance(chunk, tuple) and len(chunk) == 2:
            _, data = chunk
            return data["__interrupt__"] if "__interrupt__" in data else None
        elif isinstance(chunk, dict):
            return chunk["__interrupt__"] if "__interrupt__" in chunk else None
        return None

    async def _process_chunk(self, chunk, rendered_messages) -> None:
        """Process a streaming chunk and display separate blocks."""
        if isinstance(chunk, tuple) and len(chunk) == 2:
            # Handle subgraphs=True format: (namespace, data)
            namespace, data = chunk

            # Process different types of updates
            for node_name, node_data in data.items():
                if not isinstance(node_data, dict):
                    continue

                # Update token/cost tracking from any node that has this data
                # (middleware nodes emit these fields first, then get merged into agent state)
                await self._update_token_tracking(node_data)

                # Render messages if present
                if node_data and "messages" in node_data and node_data["messages"]:
                    messages = node_data["messages"]
                    last_message: AnyMessage = messages[-1]

                    message_id = (
                        f"{last_message.id or id(last_message)}_{last_message.type}"
                    )

                    # Skip if we've already rendered this message
                    if message_id in rendered_messages:
                        continue

                    # Mark this message as rendered
                    rendered_messages.add(message_id)

                    # Render the message
                    self.session.renderer.render_message(last_message)

    async def _update_token_tracking(self, node_data: dict[str, Any]) -> None:
        """Update session context with token tracking data if present in node."""
        token_fields = {
            "current_input_tokens",
            "current_output_tokens",
            "total_cost",
        }

        # Check if any token tracking fields are present
        if not any(field in node_data for field in token_fields):
            return

        # Extract and update context
        updates = {
            field: node_data.get(field) for field in token_fields if field in node_data
        }

        if updates:
            self.session.update_context(**updates)
            # Check if auto-compression should be triggered after token update
            await self._check_auto_compression()

    async def _check_auto_compression(self) -> None:
        """Check if auto-compression should be triggered and execute if needed."""
        try:
            ctx = self.session.context
            config_data = await initializer.load_agents_config(ctx.working_dir)
            agent_config = config_data.get_agent_config(ctx.agent)

            if not agent_config or not agent_config.compression:
                return

            compression_config = agent_config.compression

            if not compression_config.auto_compress_enabled:
                return

            context_window = agent_config.llm.context_window
            current_tokens = ctx.current_input_tokens or 0

            if should_auto_compress(
                current_tokens,
                context_window,
                compression_config.auto_compress_threshold,
            ):
                usage_pct = int(
                    (current_tokens / context_window * 100) if context_window else 0
                )

                with console.console.status(
                    f"[{theme.spinner_color}]Context at {usage_pct}%, auto-compressing to new thread...[/{theme.spinner_color}]"
                ):
                    compression_handler = CompressionHandler(self.session)
                    await compression_handler.handle()

        except Exception as e:
            logger.debug(f"Auto-compression check failed: {e}", exc_info=True)
