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

import asyncio
import uuid
from pathlib import Path

from loguru import logger

from asedclaw.agent.tools.filesystem import FilesystemToolkit
from asedclaw.agent.tools.memory import MemoryToolkit
from asedclaw.agent.tools.registry import ToolRegistry
from asedclaw.agent.tools.shell import ExecTool
from asedclaw.agent.tools.web import WebToolkit
from asedclaw.bus.events import InboundMessage
from asedclaw.bus.queue import MessageBus
from asedclaw.config.schema import ExecToolConfig
from asedclaw.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,
        exec_config: "ExecToolConfig | None" = None,
        restrict_to_workspace: bool = False,
    ):
        from asedclaw.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.exec_config = exec_config or ExecToolConfig()
        self.restrict_to_workspace = restrict_to_workspace
        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,
    ) -> 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}

        bg_task = asyncio.create_task(
            self._run_subagent(task_id, task, display_label, origin)
        )
        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],
    ) -> None:
        """Execute the subagent task and announce the result."""
        logger.info("Subagent [{}] starting task: {}", task_id, label)

        try:
            # Build subagent tools (no message tool, no spawn tool)
            tools = ToolRegistry()
            allowed_dir = self.workspace if self.restrict_to_workspace else None
            tools.register_many(
                FilesystemToolkit(workspace=self.workspace, allowed_dir=allowed_dir),
                MemoryToolkit(workspace=self.workspace),
                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,
                ),
                WebToolkit(api_key=self.brave_api_key, proxy=self.web_proxy),
            )
            system_prompt = self._build_subagent_prompt()
            final_result = await self._run_with_agno(task, system_prompt, tools)

            if final_result is None:
                final_result = "Task completed but no final response was generated."

            logger.info("Subagent [{}] completed successfully", task_id)
            await self._announce_result(task_id, label, task, final_result, origin, "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, "error")

    async def _run_with_agno(
        self,
        task: str,
        system_prompt: str,
        registry: ToolRegistry,
    ) -> str | None:
        """Run a subagent task via agno Agent."""
        try:
            from agno.agent import RunEvent
            from agno.agent.agent import Agent
            from agno.models.openai import OpenAILike
        except Exception:
            logger.exception("Subagent agno runtime unavailable")
            return None

        model_kwargs = {"id": self.model}
        if getattr(self.provider, "api_key", None):
            model_kwargs["api_key"] = self.provider.api_key
        if getattr(self.provider, "api_base", None):
            model_kwargs["base_url"] = self.provider.api_base
        if self.temperature is not None:
            model_kwargs["temperature"] = self.temperature

        agent = Agent(
            model=OpenAILike(**model_kwargs),
            tools=registry.get_agno_tool_instances(),
            instructions=system_prompt,
            markdown=True,
            stream=True,
        )

        final_content = ""
        async for event in agent.arun(task, stream=True, stream_events=True):
            if getattr(event, "event", None) == RunEvent.run_content:
                delta = getattr(event, "content", None)
                if isinstance(delta, str) and delta:
                    final_content += delta
            elif getattr(event, "event", None) == RunEvent.run_completed:
                content = getattr(event, "content", None)
                if isinstance(content, str) and content.strip():
                    final_content = content

        return final_content or None

    async def _announce_result(
        self,
        task_id: str,
        label: str,
        task: str,
        result: str,
        origin: dict[str, str],
        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."""

        # 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,
        )

        await self.bus.publish_inbound(msg)
        logger.debug("Subagent [{}] announced result to {}:{}", task_id, origin['channel'], origin['chat_id'])
    
    def _build_subagent_prompt(self) -> str:
        """Build a focused system prompt for the subagent."""
        from asedclaw.agent.context import ContextBuilder
        from asedclaw.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}"""]

        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)
