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

import asyncio
import json
import os
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.archive import ArchiveTool
from madpod_agent.agent.tools.browser import register_browser_tools
from madpod_agent.agent.tools.chart import ChartTool
from madpod_agent.agent.tools.diff import DiffTool
from madpod_agent.agent.tools.download import DownloadTool
from madpod_agent.agent.tools.extract import register_extract_tools
from madpod_agent.agent.tools.filesystem import EditFileTool, ListDirTool, ReadFileTool, WriteFileTool
from madpod_agent.agent.tools.http_request import HTTPRequestTool
from madpod_agent.agent.tools.pdf import PDFReadTool
from madpod_agent.agent.tools.python_exec import PythonExecTool
from madpod_agent.agent.tools.registry import ToolRegistry
from madpod_agent.agent.tools.shell import ExecTool
from madpod_agent.agent.tools.spreadsheet import SpreadsheetTool
from madpod_agent.agent.tools.transcribe import TranscribeTool
from madpod_agent.agent.tools.web import WebFetchTool, WebSearchTool
from madpod_agent.agent.tools.web_crawl import WebCrawlTool
from madpod_agent.agent.tools.file_convert import FileConvertTool
from madpod_agent.agent.tools.json_query import JsonQueryTool
from madpod_agent.agent.tools.regex_tool import RegexTool
from madpod_agent.agent.tools.sqlite_tool import SqliteTool
from madpod_agent.agent.tools.template_render import TemplateRenderTool
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."""

    @staticmethod
    def _env_flag(name: str) -> bool:
        return str(os.environ.get(name, "")).strip().lower() in {"1", "true", "yes", "on"}

    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
        disable_filesystem_tools = self._env_flag("MADPOD_AGENT_DISABLE_FILESYSTEM_TOOLS")
        disable_exec_tool = self._env_flag("MADPOD_AGENT_DISABLE_EXEC_TOOL")
        disable_browser_tools = self._env_flag("MADPOD_AGENT_DISABLE_BROWSER_TOOLS")
        if not disable_filesystem_tools:
            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))
        if not disable_exec_tool:
            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)

        if not disable_browser_tools:
            register_browser_tools(
                tools,
                workspace=self.workspace,
                config=self.browser_config,
            )
            browser_nav = tools.get("browser_navigate")
            if browser_nav is not None:
                register_extract_tools(tools, browser_nav.backend)
        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,
        ))
        tools.register(DownloadTool(workspace=self.workspace, proxy=self.web_proxy))
        tools.register(PDFReadTool(workspace=self.workspace, allowed_dir=allowed_dir))
        tools.register(TranscribeTool(workspace=self.workspace, allowed_dir=allowed_dir))
        tools.register(PythonExecTool(
            workspace=self.workspace,
            timeout=self.exec_config.timeout,
            allowed_dir=allowed_dir,
        ))
        # Utility tools
        tools.register(HTTPRequestTool(proxy=self.web_proxy))
        tools.register(WebCrawlTool(proxy=self.web_proxy))
        tools.register(SpreadsheetTool(workspace=self.workspace, allowed_dir=allowed_dir))
        tools.register(ChartTool(workspace=self.workspace))
        tools.register(DiffTool(workspace=self.workspace, allowed_dir=allowed_dir))
        tools.register(ArchiveTool(workspace=self.workspace, allowed_dir=allowed_dir))
        tools.register(SqliteTool(workspace=self.workspace, allowed_dir=allowed_dir))
        tools.register(TemplateRenderTool(workspace=self.workspace, allowed_dir=allowed_dir))
        tools.register(FileConvertTool(workspace=self.workspace, allowed_dir=allowed_dir))
        tools.register(RegexTool())
        tools.register(JsonQueryTool(workspace=self.workspace))
        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)
