"""Spawns Claude Code sessions via claude-agent-sdk."""

import asyncio
from pathlib import Path
from typing import Any, Awaitable, Callable, Dict, List, Optional

from mercury_launcher.log import get_logger

logger = get_logger("agent_runner")

# Fallback labels when no input details are available
_TOOL_LABELS: Dict[str, str] = {
    "Read": "Reading file...",
    "Edit": "Editing file...",
    "Write": "Writing file...",
    "Bash": "Running command...",
    "Glob": "Searching files...",
    "Grep": "Searching code...",
    "WebFetch": "Fetching web page...",
    "WebSearch": "Searching the web...",
    "Task": "Running subagent...",
    "mcp__mercury__mercury_send_message": "Sending message...",
    "mcp__mercury__mercury_wait_for_messages": "Waiting for messages...",
    "mcp__mercury__mercury_list_agents": "Listing agents...",
    "mcp__mercury__mercury_read_thread": "Reading thread...",
}


def _tool_display_name(tool_name: str, tool_input: Optional[Dict[str, Any]] = None) -> str:
    """Convert an SDK tool name + input to a descriptive status label."""
    inp = tool_input or {}

    # Extract details from input where useful
    if tool_name in ("Read", "Edit", "Write"):
        fp = inp.get("file_path", "")
        if fp:
            short = fp.rsplit("/", 1)[-1] if "/" in fp else fp
            verb = {"Read": "Reading", "Edit": "Editing", "Write": "Writing"}[tool_name]
            return f"{verb} {short}"
    elif tool_name == "Bash":
        cmd = inp.get("command", "")
        if cmd:
            short = cmd[:60] + ("..." if len(cmd) > 60 else "")
            return f"Running {short}"
    elif tool_name in ("Grep", "Glob"):
        pattern = inp.get("pattern", "")
        if pattern:
            return f"Searching {pattern[:40]}"
    elif tool_name == "mcp__mercury__mercury_send_message":
        to = inp.get("to", "")
        if to:
            return f"Messaging {to}"
    elif tool_name == "Task":
        desc = inp.get("description", "")
        if desc:
            return f"Subagent: {desc[:40]}"

    # Fallback to static labels
    if tool_name in _TOOL_LABELS:
        return _TOOL_LABELS[tool_name]
    if tool_name.startswith("mcp__mercury__"):
        short = tool_name.removeprefix("mcp__mercury__mercury_").replace("_", " ")
        return f"{short.capitalize()}..."
    return f"Using {tool_name}..."


class AgentRunner:
    """Manages Claude Code session lifecycle for a single Mercury agent."""

    def __init__(
        self,
        agent_name: str,
        mercury_url: str,
        mcp_api_key: str,
        working_directory: Path,
        model: Optional[str] = None,
        max_budget_usd: float = 2.0,
        on_status: Optional[Callable[[str], Awaitable[None]]] = None,
        on_activity: Optional[Callable[[str, Optional[str]], Awaitable[None]]] = None,
    ):
        self.agent_name = agent_name
        self.mercury_url = mercury_url
        self.mcp_api_key = mcp_api_key
        self.working_directory = working_directory
        self.model = model
        self.max_budget_usd = max_budget_usd
        self._on_status = on_status
        self._on_activity = on_activity  # (label, text_content) for rich terminal + chat activity
        self._active = False

    @property
    def is_active(self) -> bool:
        return self._active

    async def run_session(
        self,
        system_prompt: str,
        session_prompt: str,
    ) -> Dict[str, Any]:
        """Spawn a Claude Code session and wait for it to complete.

        Returns a dict with session results (cost, turns, duration, etc).
        """
        if self._active:
            logger.warning("[%s] Session already active, skipping", self.agent_name)
            return {"error": "session_already_active"}

        self._active = True
        try:
            return await self._spawn_session(system_prompt, session_prompt)
        finally:
            self._active = False

    async def _spawn_session(
        self,
        system_prompt: str,
        session_prompt: str,
    ) -> Dict[str, Any]:
        """Internal: spawn and run the Claude Code session."""
        from claude_agent_sdk import query, ClaudeAgentOptions, ResultMessage
        from claude_agent_sdk.types import AssistantMessage, TextBlock, ToolUseBlock

        # Configure MCP server for Mercury
        mcp_servers = {
            "mercury": {
                "type": "http",
                "url": self.mercury_url,
                "headers": {"x-api-key": self.mcp_api_key},
            }
        }

        options = ClaudeAgentOptions(
            system_prompt={
                "type": "preset",
                "preset": "claude_code",
                "append": system_prompt,
            },
            permission_mode="acceptEdits",
            cwd=str(self.working_directory),
            mcp_servers=mcp_servers,
            allowed_tools=[
                "mcp__mercury__*",  # All Mercury MCP tools
                "Read", "Write", "Edit", "Glob", "Grep", "Bash",
            ],
        )

        if self.model:
            options.model = self.model

        logger.info("[%s] Spawning Claude Code session", self.agent_name)

        result_info: Dict[str, Any] = {}

        try:
            async for message in query(prompt=session_prompt, options=options):
                # Surface tool use and text as status/activity updates
                if isinstance(message, AssistantMessage):
                    for block in message.content:
                        if isinstance(block, ToolUseBlock):
                            tool_input = block.input if hasattr(block, "input") and isinstance(block.input, dict) else {}
                            label = _tool_display_name(block.name, tool_input)
                            if self._on_status:
                                try:
                                    await self._on_status(label)
                                except Exception:
                                    pass
                            if self._on_activity:
                                try:
                                    await self._on_activity(label, None)
                                except Exception:
                                    pass
                        elif isinstance(block, TextBlock) and block.text:
                            text = block.text[:120] + ("..." if len(block.text) > 120 else "")
                            if self._on_activity:
                                try:
                                    await self._on_activity("Thinking", text)
                                except Exception:
                                    pass

                if isinstance(message, ResultMessage):
                    result_info = {
                        "status": getattr(message, "subtype", "unknown"),
                        "cost_usd": getattr(message, "total_cost_usd", 0),
                        "num_turns": getattr(message, "num_turns", 0),
                        "duration_ms": getattr(message, "duration_ms", 0),
                    }
                    logger.info(
                        "[%s] Session complete: status=%s, cost=$%.4f, turns=%d, duration=%dms",
                        self.agent_name,
                        result_info.get("status"),
                        result_info.get("cost_usd", 0),
                        result_info.get("num_turns", 0),
                        result_info.get("duration_ms", 0),
                    )

        except Exception as e:
            logger.error("[%s] Session error: %s", self.agent_name, e, exc_info=True)
            result_info = {"status": "error", "error": str(e)}

        return result_info
