"""Subagent manager for background task execution."""

import asyncio
import json
import uuid
from contextlib import AsyncExitStack
from pathlib import Path
from typing import Any

from loguru import logger

from madpod_agent.agent.investigation import (
    EvidenceLog,
    InvestigationFindings,
    InvestigationProfile,
    build_provisional_response,
    build_retry_instruction,
    build_turn_policy,
    classify_request,
    record_tool_event,
    verify_investigation,
)
from madpod_agent.agent.tools.browser import register_browser_tools
from madpod_agent.agent.tools.filesystem import EditFileTool, ListDirTool, ReadFileTool, WriteFileTool
from madpod_agent.agent.tools.registry import ToolRegistry
from madpod_agent.agent.tools.shell import ExecTool
from madpod_agent.agent.tools.web import WebFetchTool, WebSearchTool
from madpod_agent.bus.events import InboundMessage
from madpod_agent.bus.queue import MessageBus
from madpod_agent.config.schema import BrowserToolsConfig, ExecToolConfig, WebSearchConfig
from madpod_agent.providers.base import LLMProvider


class SubagentManager:
    """Manages background subagent execution."""

    def __init__(
        self,
        provider: LLMProvider,
        workspace: Path,
        bus: MessageBus,
        model: str | None = None,
        temperature: float = 0.7,
        max_tokens: int = 4096,
        reasoning_effort: str | None = None,
        brave_api_key: str | None = None,
        web_proxy: str | None = None,
        web_search_config: WebSearchConfig | None = None,
        summary_provider: LLMProvider | None = None,
        browser_config: BrowserToolsConfig | None = None,
        exec_config: "ExecToolConfig | None" = None,
        restrict_to_workspace: bool = False,
        mcp_servers: dict | None = None,
        summary_model: str | None = None,
    ):
        from madpod_agent.config.schema import ExecToolConfig
        self.provider = provider
        self.workspace = workspace
        self.bus = bus
        self.model = model or provider.get_default_model()
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.reasoning_effort = reasoning_effort
        self.brave_api_key = brave_api_key
        self.web_proxy = web_proxy
        self.web_search_config = web_search_config
        self.summary_provider = summary_provider or provider
        self.browser_config = browser_config or BrowserToolsConfig()
        self.exec_config = exec_config or ExecToolConfig()
        self.restrict_to_workspace = restrict_to_workspace
        self.mcp_servers = mcp_servers or {}
        self.summary_model = summary_model
        self._running_tasks: dict[str, asyncio.Task[None]] = {}
        self._session_tasks: dict[str, set[str]] = {}  # session_key -> {task_id, ...}

    async def spawn(
        self,
        task: str,
        label: str | None = None,
        origin_channel: str = "cli",
        origin_chat_id: str = "direct",
        session_key: str | None = None,
        origin_message_id: str | None = None,
        origin_metadata: dict[str, Any] | None = None,
    ) -> str:
        """Spawn a subagent to execute a task in the background."""
        task_id = str(uuid.uuid4())[:8]
        display_label = label or task[:30] + ("..." if len(task) > 30 else "")
        origin = {
            "channel": origin_channel,
            "chat_id": origin_chat_id,
            "session_key": session_key or f"{origin_channel}:{origin_chat_id}",
        }
        origin_metadata = dict(origin_metadata or {})
        if origin_message_id and not str(origin_metadata.get("assignment_message_id") or "").strip():
            origin_metadata["assignment_message_id"] = origin_message_id
        if origin_message_id and not str(origin_metadata.get("message_id") or "").strip():
            origin_metadata["message_id"] = origin_message_id

        bg_task = asyncio.create_task(
            self._run_subagent(task_id, task, display_label, origin, origin_metadata)
        )
        self._running_tasks[task_id] = bg_task
        if session_key:
            self._session_tasks.setdefault(session_key, set()).add(task_id)

        def _cleanup(_: asyncio.Task) -> None:
            self._running_tasks.pop(task_id, None)
            if session_key and (ids := self._session_tasks.get(session_key)):
                ids.discard(task_id)
                if not ids:
                    del self._session_tasks[session_key]

        bg_task.add_done_callback(_cleanup)

        logger.info("Spawned subagent [{}]: {}", task_id, display_label)
        return f"Subagent [{display_label}] started (id: {task_id}). I'll notify you when it completes."

    async def _run_subagent(
        self,
        task_id: str,
        task: str,
        label: str,
        origin: dict[str, str],
        origin_metadata: dict[str, Any],
    ) -> None:
        """Execute the subagent task and announce the result."""
        logger.info("Subagent [{}] starting task: {}", task_id, label)

        try:
            findings = await self._execute_task_loop(
                task,
                profile=classify_request(task),
                read_only=False,
            )

            logger.info("Subagent [{}] completed successfully", task_id)
            await self._announce_result(
                task_id,
                label,
                task,
                findings.summary,
                origin,
                origin_metadata,
                "ok",
            )

        except Exception as e:
            error_msg = f"Error: {str(e)}"
            logger.error("Subagent [{}] failed: {}", task_id, e)
            await self._announce_result(
                task_id,
                label,
                task,
                error_msg,
                origin,
                origin_metadata,
                "error",
            )

    async def investigate(self, task: str) -> InvestigationFindings:
        """Run a bounded read-only investigation pass and return the findings."""
        profile = classify_request(task)
        return await self._execute_task_loop(
            self._build_investigation_task(task),
            profile=profile,
            read_only=True,
        )

    async def _build_tools(self, *, read_only: bool) -> tuple[ToolRegistry, AsyncExitStack | None]:
        """Build the tool registry for normal or investigation-only subagents."""
        tools = ToolRegistry()
        allowed_dir = self.workspace if self.restrict_to_workspace else None
        tools.register(ReadFileTool(workspace=self.workspace, allowed_dir=allowed_dir))
        tools.register(ListDirTool(workspace=self.workspace, allowed_dir=allowed_dir))
        if not read_only:
            tools.register(WriteFileTool(workspace=self.workspace, allowed_dir=allowed_dir))
            tools.register(EditFileTool(workspace=self.workspace, allowed_dir=allowed_dir))
        tools.register(ExecTool(
            working_dir=str(self.workspace),
            timeout=self.exec_config.timeout,
            restrict_to_workspace=self.restrict_to_workspace,
            path_append=self.exec_config.path_append,
        ))

        mcp_stack: AsyncExitStack | None = None
        if self.mcp_servers:
            from madpod_agent.agent.tools.mcp import connect_mcp_servers

            mcp_stack = AsyncExitStack()
            await mcp_stack.__aenter__()
            await connect_mcp_servers(self.mcp_servers, tools, mcp_stack)

        register_browser_tools(
            tools,
            workspace=self.workspace,
            config=self.browser_config,
        )
        tools.register(WebSearchTool(
            api_key=self.brave_api_key,
            proxy=self.web_proxy,
            provider=self.summary_provider,
            workspace=self.workspace,
            search_config=self.web_search_config,
            summary_model=self.summary_model,
        ))
        tools.register(WebFetchTool(
            proxy=self.web_proxy,
            workspace=self.workspace,
            search_config=self.web_search_config,
        ))
        return tools, mcp_stack

    async def _execute_task_loop(
        self,
        task: str,
        *,
        profile: InvestigationProfile,
        read_only: bool,
    ) -> InvestigationFindings:
        """Run the isolated subagent loop and return evidence-backed findings."""
        tools, mcp_stack = await self._build_tools(read_only=read_only)
        messages: list[dict[str, Any]] = [
            {
                "role": "system",
                "content": self._build_subagent_prompt(
                    profile=profile,
                    read_only=read_only,
                ),
            },
            {"role": "user", "content": task},
        ]
        evidence = EvidenceLog()
        forced_tool_retry = False
        verification_retry = False
        final_result: str | None = None

        try:
            for iteration in range(1, 16):
                response = await self.provider.chat(
                    messages=messages,
                    tools=tools.get_definitions(),
                    model=self.model,
                    temperature=self.temperature,
                    max_tokens=self.max_tokens,
                    reasoning_effort=self.reasoning_effort,
                )

                if response.has_tool_calls:
                    tool_call_dicts = [
                        {
                            "id": tc.id,
                            "type": "function",
                            "function": {
                                "name": tc.name,
                                "arguments": json.dumps(tc.arguments, ensure_ascii=False),
                            },
                        }
                        for tc in response.tool_calls
                    ]
                    messages.append({
                        "role": "assistant",
                        "content": response.content or "",
                        "tool_calls": tool_call_dicts,
                    })
                    for tool_call in response.tool_calls:
                        args_str = json.dumps(tool_call.arguments, ensure_ascii=False)
                        logger.debug("Subagent executing: {} with arguments: {}", tool_call.name, args_str)
                        result = await tools.execute(tool_call.name, tool_call.arguments)
                        record_tool_event(
                            evidence,
                            tool_name=tool_call.name,
                            arguments=tool_call.arguments,
                            result=result,
                        )
                        messages.append({
                            "role": "tool",
                            "tool_call_id": tool_call.id,
                            "name": tool_call.name,
                            "content": result,
                        })
                    continue

                clean = (response.content or "").strip()
                if response.finish_reason == "error":
                    final_result = clean or "The investigation failed because the model call returned an error."
                    break

                if profile.is_investigative and not evidence.tool_summaries and not forced_tool_retry:
                    forced_tool_retry = True
                    messages.append({
                        "role": "user",
                        "content": build_retry_instruction(
                            profile,
                            missing=["gather concrete evidence before answering"],
                            evidence=evidence,
                        ),
                    })
                    continue

                verified, missing = verify_investigation(profile, evidence)
                if profile.is_investigative and not verified and not verification_retry:
                    verification_retry = True
                    messages.append({
                        "role": "user",
                        "content": build_retry_instruction(
                            profile,
                            missing=missing,
                            evidence=evidence,
                        ),
                    })
                    continue

                final_result = clean
                if profile.is_investigative and not verified:
                    final_result = build_provisional_response(
                        clean,
                        missing=missing,
                        evidence=evidence,
                    )
                break
            else:
                final_result = "The investigation hit its iteration limit before it could finish."
        finally:
            if mcp_stack is not None:
                await mcp_stack.aclose()

        return InvestigationFindings(
            summary=final_result or "Task completed but no final response was generated.",
            evidence=evidence,
        )

    async def _announce_result(
        self,
        task_id: str,
        label: str,
        task: str,
        result: str,
        origin: dict[str, str],
        origin_metadata: dict[str, Any],
        status: str,
    ) -> None:
        """Announce the subagent result to the main agent via the message bus."""
        status_text = "completed successfully" if status == "ok" else "failed"

        announce_content = f"""[Subagent '{label}' {status_text}]

Task: {task}

Result:
{result}

Summarize this naturally for the user. Keep it brief (1-2 sentences). Do not mention technical details like "subagent" or task IDs."""

        source = dict(origin_metadata or {})
        assignment_message_id = str(
            source.get("assignment_message_id")
            or source.get("directive_id")
            or source.get("message_id")
            or ""
        ).strip() or None
        upstream_in_reply_to = str(source.get("in_reply_to") or "").strip() or None
        metadata: dict[str, Any] = {
            "message_id": uuid.uuid4().hex,
        }
        for key in ("round_id", "subtask_id"):
            value = source.get(key)
            if value:
                metadata[key] = value
        if assignment_message_id:
            metadata["assignment_message_id"] = assignment_message_id
            metadata["directive_id"] = str(source.get("directive_id") or assignment_message_id)
            metadata["in_reply_to"] = assignment_message_id
        elif upstream_in_reply_to:
            metadata["in_reply_to"] = upstream_in_reply_to
        if upstream_in_reply_to and upstream_in_reply_to != metadata.get("in_reply_to"):
            metadata["origin_in_reply_to"] = upstream_in_reply_to

        # Inject as system message to trigger main agent
        msg = InboundMessage(
            channel="system",
            sender_id="subagent",
            chat_id=f"{origin['channel']}:{origin['chat_id']}",
            content=announce_content,
            metadata=metadata,
            session_key_override=origin.get("session_key"),
        )

        await self.bus.publish_inbound(msg)
        logger.debug("Subagent [{}] announced result to {}:{}", task_id, origin['channel'], origin['chat_id'])

    @staticmethod
    def _build_investigation_task(task: str) -> str:
        """Turn a broad task into a read-only investigation brief."""
        return (
            "Gather findings for the main agent. Do not edit files or claim the task is complete.\n\n"
            "Return a concise findings report with concrete evidence such as file paths, commands, and source URLs when available.\n\n"
            f"Task:\n{task}"
        )

    def _build_subagent_prompt(
        self,
        *,
        profile: InvestigationProfile,
        read_only: bool,
    ) -> str:
        """Build a focused system prompt for the subagent."""
        from madpod_agent.agent.context import ContextBuilder
        from madpod_agent.agent.skills import SkillsLoader

        time_ctx = ContextBuilder._build_runtime_context(None, None)
        parts = [f"""# Subagent

{time_ctx}

You are a subagent spawned by the main agent to complete a specific task.
Stay focused on the assigned task. Your final response will be reported back to the main agent.

## Workspace
{self.workspace}"""]
        policy = build_turn_policy(profile)
        if policy:
            parts.append(policy)
        if read_only:
            parts.append(
                "## Investigation Mode\n\n"
                "This is a bounded investigation pass. Stay read-only, gather evidence, and return findings for the main agent instead of claiming the whole task is finished."
            )
        parts.append(
            """## Browser Guidance

If `browser_*` tools are available, you can open Chrome/Chromium and interact with pages directly.
Prefer `browser_*` for interactive or dynamic pages and explicit browser-opening requests.
Prefer `web_search` / `web_fetch` for cheap simple lookup or citation gathering.
Re-run `browser_snapshot` after navigation or page changes because refs can change."""
        )

        skills_summary = SkillsLoader(self.workspace).build_skills_summary()
        if skills_summary:
            parts.append(f"## Skills\n\nRead SKILL.md with read_file to use a skill.\n\n{skills_summary}")

        return "\n\n".join(parts)

    async def cancel_by_session(self, session_key: str) -> int:
        """Cancel all subagents for the given session. Returns count cancelled."""
        tasks = [self._running_tasks[tid] for tid in self._session_tasks.get(session_key, [])
                 if tid in self._running_tasks and not self._running_tasks[tid].done()]
        for t in tasks:
            t.cancel()
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)
        return len(tasks)

    def get_running_count(self) -> int:
        """Return the number of currently running subagents."""
        return len(self._running_tasks)
