"""Message handling for Claude Agents SDK.

This module provides a simplified message handler that forwards Claude SDK
messages directly to the backend via adk.messages.send().

The recommended approach is:

    from terminaluse.lib import adk

    async for message in query(prompt="...", options=options):
        await adk.messages.send(task_id=task_id, content=message)

adk.messages.send() accepts Claude SDK message types directly and handles:
- Automatic forwarding to the backend via raw_events API
- Backend transformation via ClaudeAdapter
- UI streaming to connected clients
"""

from __future__ import annotations

from typing import Any

from claude_agent_sdk import (
    AssistantMessage,
    ResultMessage,
    SystemMessage,
    TextBlock,
)

from terminaluse.lib import adk
from terminaluse.lib.utils.logging import make_logger

logger = make_logger(__name__)


class ClaudeMessageHandler:
    """Handles Claude SDK messages and forwards them to TerminalUse.

    Simplified handler focused on:
    - Forwarding messages to backend via adk.messages.send()
    - Extracting session_id from SystemMessage/ResultMessage
    - Extracting usage and cost from ResultMessage
    - Serializing responses for Temporal

    Note: Tool lifecycle events (requests/responses) are handled by
    TemporalStreamingHooks, not this class.
    """

    def __init__(
        self,
        task_id: str | None,
        trace_id: str | None,
        parent_span_id: str | None,
    ):
        self.task_id = task_id
        self.trace_id = trace_id
        self.parent_span_id = parent_span_id

        # Message tracking
        self.messages: list[Any] = []
        self.serialized_messages: list[dict] = []

        # Result data
        self.session_id: str | None = None
        self.usage_info: dict | None = None
        self.cost_info: float | None = None

    async def initialize(self):
        """Initialize the handler. No-op in the simplified version."""
        pass

    async def handle_message(self, message: Any):
        """Process a single message from Claude SDK."""
        self.messages.append(message)
        msg_num = len(self.messages)

        # Debug logging (verbose - only for troubleshooting)
        logger.debug(f"[{msg_num}] Message type: {type(message).__name__}")
        if isinstance(message, AssistantMessage):
            block_types = [type(b).__name__ for b in message.content]
            logger.debug(f"   [{msg_num}] Content blocks: {block_types}")

        # Forward message to backend via adk.messages.send()
        if self.task_id:
            await adk.messages.send(task_id=self.task_id, content=message)

        # Route to specific handlers for metadata extraction
        if isinstance(message, AssistantMessage):
            await self._handle_assistant_message(message, msg_num)
        elif isinstance(message, SystemMessage):
            await self._handle_system_message(message)
        elif isinstance(message, ResultMessage):
            await self._handle_result_message(message)

    async def _handle_assistant_message(self, message: AssistantMessage, _msg_num: int):
        """Handle AssistantMessage - extract text content for serialization.

        Note: The actual message streaming is handled by adk.messages.send().
        This method only extracts text for the serialized response.
        """
        # Collect text for final response
        text_content = []
        for block in message.content:
            if isinstance(block, TextBlock):
                text_content.append(block.text)

        if text_content:
            self.serialized_messages.append({"role": "assistant", "content": "\n".join(text_content)})

    async def _handle_system_message(self, message: SystemMessage):
        """Handle system message - extract session_id."""
        if message.subtype == "init":
            self.session_id = message.data.get("session_id")
            logger.debug(f"Session initialized: {self.session_id[:16] if self.session_id else 'unknown'}...")
        else:
            logger.debug(f"SystemMessage: {message.subtype}")

    async def _handle_result_message(self, message: ResultMessage):
        """Handle result message - extract usage and cost."""
        self.usage_info = message.usage
        self.cost_info = message.total_cost_usd

        # Update session_id if available
        if message.session_id:
            self.session_id = message.session_id

        logger.info(f"Cost: ${self.cost_info:.4f}, Duration: {message.duration_ms}ms, Turns: {message.num_turns}")

    async def cleanup(self):
        """Clean up. No-op in the simplified version."""
        pass

    def get_results(self) -> dict[str, Any]:
        """Get final results for Temporal."""
        return {
            "messages": self.serialized_messages,
            "task_id": self.task_id,
            "session_id": self.session_id,
            "usage": self.usage_info,
            "cost_usd": self.cost_info,
        }
