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

from __future__ import annotations

import asyncio
import json
import time
import uuid
from pathlib import Path
from typing import Any

from loguru import logger

from flowly.bus.events import InboundMessage
from flowly.bus.queue import MessageBus
from flowly.providers.base import LLMProvider
from flowly.agent.tools.registry import ToolRegistry
from flowly.agent.tools.filesystem import ReadFileTool, WriteFileTool, EditFileTool, ListDirTool
from flowly.agent.tools.shell import ExecTool
from flowly.exec.types import ExecConfig
from flowly.agent.tools.web import WebSearchTool, WebFetchTool
from flowly.agent.subagent_registry import SubagentRegistry, SubagentRunRecord
from flowly.agent.subagent_announce_queue import AnnounceItem, get_or_create_queue


class SubagentManager:
    """
    Manages background subagent execution.

    Subagents are lightweight agent instances that run in the background
    to handle specific tasks. They share the same LLM provider but have
    isolated context and a focused system prompt.
    """

    def __init__(
        self,
        provider: LLMProvider,
        workspace: Path,
        bus: MessageBus,
        model: str | None = None,
        brave_api_key: str | None = None,
        exec_config: ExecConfig | None = None,
        registry: SubagentRegistry | None = None,
    ):
        self.provider = provider
        self.workspace = workspace
        self.bus = bus
        self.model = model or provider.get_default_model()
        self.brave_api_key = brave_api_key
        self.exec_config = exec_config or ExecConfig()
        self._registry = registry or SubagentRegistry()
        self._running_tasks: dict[str, asyncio.Task[None]] = {}
        self._busy_sessions: set[str] = set()

    # ------------------------------------------------------------------
    # Busy tracking (called by AgentLoop around message processing)
    # ------------------------------------------------------------------

    def mark_busy(self, session_key: str) -> None:
        self._busy_sessions.add(session_key)

    def mark_idle(self, session_key: str) -> None:
        self._busy_sessions.discard(session_key)

    def is_busy(self, session_key: str) -> bool:
        return session_key in self._busy_sessions

    # ------------------------------------------------------------------
    # Crash recovery
    # ------------------------------------------------------------------

    def resume_pending(self) -> None:
        """Called on startup: announce any tasks that were running when process crashed."""
        pending = self._registry.pending()
        if not pending:
            return
        logger.warning(f"[SubagentManager] Found {len(pending)} pending task(s) from previous run — marking as failed")
        for record in pending:
            self._registry.update(
                record.run_id,
                ended_at=time.time(),
                outcome="error",
                error="Process restarted before task completed",
            )
            asyncio.create_task(
                self._announce(record.run_id, "The process restarted before this task completed.", "error")
            )

    # ------------------------------------------------------------------
    # Public spawn
    # ------------------------------------------------------------------

    async def spawn(
        self,
        task: str,
        label: str | None = None,
        origin_channel: str = "cli",
        origin_chat_id: str = "direct",
        model: str | None = None,
        cleanup: str = "keep",
        timeout_seconds: int | None = None,
        is_subagent_caller: bool = False,
    ) -> str:
        """
        Spawn a subagent to execute a task in the background.

        Args:
            task: The task description for the subagent.
            label: Optional human-readable label for the task.
            origin_channel: The channel to announce results to.
            origin_chat_id: The chat ID to announce results to.
            model: Optional model override for this subagent.
            cleanup: "keep" or "delete" (registry cleanup policy).
            timeout_seconds: Optional execution timeout.
            is_subagent_caller: True if spawner is itself a subagent (loop prevention).

        Returns:
            JSON status string.
        """
        # Loop prevention: subagents cannot spawn subagents
        if is_subagent_caller:
            logger.warning("[SubagentManager] Blocked recursive spawn attempt")
            return json.dumps({
                "status": "forbidden",
                "error": "Subagents cannot spawn other subagents",
            })

        run_id = str(uuid.uuid4())
        child_session_key = f"subagent:{run_id}"
        parent_session_key = f"{origin_channel}:{origin_chat_id}"
        display_label = label or (task[:40] + "..." if len(task) > 40 else task)
        resolved_model = model or self.model

        record = SubagentRunRecord(
            run_id=run_id,
            child_session_key=child_session_key,
            parent_session_key=parent_session_key,
            parent_channel=origin_channel,
            parent_chat_id=origin_chat_id,
            task=task,
            label=display_label,
            model=resolved_model,
            cleanup=cleanup,
            created_at=time.time(),
        )
        self._registry.register(record)

        bg_task = asyncio.create_task(
            self._run_subagent(run_id, task, display_label, origin_channel, origin_chat_id,
                               resolved_model, timeout_seconds)
        )
        self._running_tasks[run_id] = bg_task
        bg_task.add_done_callback(lambda _: self._running_tasks.pop(run_id, None))

        logger.info(f"[SubagentManager] Spawned [{run_id[:8]}]: {display_label}")
        return json.dumps({
            "status": "accepted",
            "run_id": run_id[:8],
            "label": display_label,
            "message": f"Background task '{display_label}' started. I'll notify you when it completes.",
        })

    # ------------------------------------------------------------------
    # Internal execution (core loop preserved from original)
    # ------------------------------------------------------------------

    async def _run_subagent(
        self,
        run_id: str,
        task: str,
        label: str,
        origin_channel: str,
        origin_chat_id: str,
        model: str,
        timeout_seconds: int | None,
    ) -> None:
        """Execute the subagent task and announce the result."""
        logger.info(f"[SubagentManager] [{run_id[:8]}] starting: {label}")
        self._registry.update(run_id, started_at=time.time())

        async def _do_run() -> str:
            # Build subagent tools (no message tool, no spawn tool)
            tools = ToolRegistry()
            tools.register(ReadFileTool())
            tools.register(WriteFileTool())
            tools.register(EditFileTool())
            tools.register(ListDirTool())
            tools.register(ExecTool(config=self.exec_config, working_dir=str(self.workspace)))
            tools.register(WebSearchTool(api_key=self.brave_api_key))
            tools.register(WebFetchTool())

            # Build messages with subagent-specific prompt
            system_prompt = self._build_subagent_prompt(task)
            messages: list[dict[str, Any]] = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": task},
            ]

            # Run agent loop (limited iterations) — identical to original
            max_iterations = 15
            iteration = 0
            final_result: str | None = None

            while iteration < max_iterations:
                iteration += 1

                response = await self.provider.chat(
                    messages=messages,
                    tools=tools.get_definitions(),
                    model=model,
                )

                if response.has_tool_calls:
                    # Add assistant message with tool calls
                    tool_call_dicts = [
                        {
                            "id": tc.id,
                            "type": "function",
                            "function": {
                                "name": tc.name,
                                "arguments": json.dumps(tc.arguments),
                            },
                        }
                        for tc in response.tool_calls
                    ]
                    messages.append({
                        "role": "assistant",
                        "content": response.content or "",
                        "tool_calls": tool_call_dicts,
                    })

                    # Execute tools
                    for tool_call in response.tool_calls:
                        logger.debug(f"[SubagentManager] [{run_id[:8]}] tool: {tool_call.name}")
                        result = await tools.execute(tool_call.name, tool_call.arguments)
                        messages.append({
                            "role": "tool",
                            "tool_call_id": tool_call.id,
                            "name": tool_call.name,
                            "content": result,
                        })
                else:
                    final_result = response.content
                    break

            return final_result or "Task completed but no final response was generated."

        try:
            if timeout_seconds:
                final_result = await asyncio.wait_for(_do_run(), timeout=timeout_seconds)
            else:
                final_result = await _do_run()

            logger.info(f"[SubagentManager] [{run_id[:8]}] completed successfully")
            self._registry.update(run_id, ended_at=time.time(), outcome="ok")
            await self._announce(run_id, final_result, "ok")

        except asyncio.TimeoutError:
            logger.warning(f"[SubagentManager] [{run_id[:8]}] timed out")
            self._registry.update(run_id, ended_at=time.time(), outcome="timeout")
            await self._announce(run_id, "(task timed out)", "timeout")

        except Exception as e:
            logger.error(f"[SubagentManager] [{run_id[:8]}] failed: {e}")
            self._registry.update(run_id, ended_at=time.time(), outcome="error", error=str(e))
            await self._announce(run_id, f"Error: {e}", "error")

    async def _announce(self, run_id: str, result: str, status: str) -> None:
        """Deliver result to parent agent — via queue if parent is busy."""
        record = self._registry.get(run_id)
        if not record:
            return

        status_text = {
            "ok": "completed successfully",
            "error": "failed",
            "timeout": "timed out",
        }.get(status, status)

        duration = ""
        if record.started_at and record.ended_at:
            secs = int(record.ended_at - record.started_at)
            duration = f" (runtime: {secs}s)"

        announce_content = (
            f"[Background task '{record.label}' {status_text}{duration}]\n\n"
            f"Task: {record.task}\n\n"
            f"Result:\n{result}\n\n"
            "Summarize this naturally for the user. Keep it brief (1-2 sentences). "
            "Do not mention technical details like 'subagent' or task IDs."
        )

        parent_session_key = record.parent_session_key

        async def _send(prompt: str) -> None:
            msg = InboundMessage(
                channel="system",
                sender_id="subagent",
                chat_id=parent_session_key,
                content=prompt,
            )
            await self.bus.publish_inbound(msg)
            logger.debug(f"[SubagentManager] [{run_id[:8]}] announced to {parent_session_key}")

        if self.is_busy(parent_session_key):
            # Parent is currently processing — queue the announce
            q = get_or_create_queue(parent_session_key, _send)
            await q.enqueue(AnnounceItem(
                prompt=announce_content,
                summary=f"{record.label} — {status_text}",
            ))
            logger.debug(f"[SubagentManager] [{run_id[:8]}] queued announce (parent busy)")
        else:
            await _send(announce_content)

        self._registry.update(run_id, announced=True)

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _build_subagent_prompt(self, task: str) -> str:
        """Build a focused system prompt for the subagent."""
        return f"""# Subagent

You are a subagent spawned by the main agent to complete a specific task.

## Your Task
{task}

## Rules
1. Stay focused - complete only the assigned task, nothing else
2. Your final response will be reported back to the main agent
3. Do not initiate conversations or take on side tasks
4. Be concise but informative in your findings

## What You Can Do
- Read and write files in the workspace
- Execute shell commands
- Search the web and fetch web pages
- Complete the task thoroughly

## What You Cannot Do
- Send messages directly to users (no message tool available)
- Spawn other subagents
- Access the main agent's conversation history

## Workspace
Your workspace is at: {self.workspace}

When you have completed the task, provide a clear summary of your findings or actions."""

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