# 中文注释：
# 文件：echobot/agent/subagent.py
# 说明：智能体核心逻辑，包括主循环、上下文、记忆与任务队列。

"""Subagent Manager - 后台子任务管理
=================================================

此模块管理后台运行的子任务（Subagent）。

子任务是什么：
- 子任务是轻量级的 Agent 实例，在后台运行
- 用于处理耗时任务，不阻塞主 Agent
- 完成后通过消息总线通知主 Agent

使用场景：
- 复杂的研究任务
- 大规模代码重构
- 长时间运行的分析任务

主要类：
- SubagentManager: 子任务管理器

使用示例：
    manager = SubagentManager(provider, workspace, bus)
    task_id = await manager.spawn("Research X topic", label="research")
"""

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

from loguru import logger

from echobot.bus.events import InboundMessage
from echobot.bus.queue import MessageBus
from echobot.providers.base import LLMProvider
from echobot.agent.tools.registry import ToolRegistry
from echobot.agent.tools.filesystem import ReadFileTool, WriteFileTool, ListDirTool
from echobot.agent.tools.shell import ExecTool
from echobot.agent.tools.web import WebSearchTool, WebFetchTool


class SubagentManager:
    """
    后台子任务管理器
    
    子任务是轻量级的 Agent 实例，在后台运行以处理耗时任务。
    完成后通过消息总线将结果通知主 Agent。
    
    设计特点：
    1. 独立上下文：子任务有自己的系统提示词
    2. 受限工具集：没有消息和 spawn 工具
    3. 结果路由：结果会自动发送到指定渠道
    
    主要功能：
    - spawn(): 创建后台子任务
    - get_running_count(): 获取运行中的子任务数量
    
    使用示例：
        manager = SubagentManager(provider, workspace, bus)
        task_id = await manager.spawn(
            task="Research AI news",
            label="research",
            origin_channel="telegram",
            origin_chat_id="12345"
        )
    """
    
    def __init__(
        self,
        provider: LLMProvider,
        workspace: Path,
        bus: MessageBus,
        model: str | None = None,
        brave_api_key: str | None = None,
    ):
        """
        初始化子任务管理器
        
        Args:
            provider: LLM 提供商
            workspace: 工作空间路径
            bus: 消息总线
            model: 使用的模型
            brave_api_key: Brave Search API 密钥
        """
        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._running_tasks: dict[str, asyncio.Task[None]] = {}
    
    async def spawn(
        self,
        task: str,
        label: str | None = None,
        origin_channel: str = "cli",
        origin_chat_id: str = "direct",
    ) -> str:
        """
        创建后台子任务
        
        Args:
            task: 子任务描述
            label: 可读标签（用于显示）
            origin_channel: 结果发送的目标渠道
            origin_chat_id: 结果发送的目标聊天 ID
        
        Returns:
            str: 状态消息
        """
        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
        
        # 完成后清理
        bg_task.add_done_callback(lambda _: self._running_tasks.pop(task_id, None))
        
        logger.info(f"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:
        """
        执行子任务
        
        流程：
        1. 构建工具集（无消息工具）
        2. 构建子任务专用的系统提示词
        3. 运行 Agent 循环
        4. 完成后通知主 Agent
        
        Args:
            task_id: 子任务 ID
            task: 任务描述
            label: 显示标签
            origin: 原始目标信息
        """
        logger.info(f"Subagent [{task_id}] starting task: {label}")
        
        try:
            # 构建工具集（无消息工具，无 spawn 工具）
            tools = ToolRegistry()
            # 子任务工具同样绑定工作区边界，避免后台任务越界访问。
            tools.register(ReadFileTool(root_dir=self.workspace))
            tools.register(WriteFileTool(root_dir=self.workspace))
            tools.register(ListDirTool(root_dir=self.workspace))
            tools.register(ExecTool(working_dir=str(self.workspace), root_dir=self.workspace))
            tools.register(WebSearchTool(api_key=self.brave_api_key))
            tools.register(WebFetchTool())
            
            # 构建消息（使用子任务专用提示词）
            system_prompt = self._build_subagent_prompt(task)
            messages: list[dict[str, Any]] = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": task},
            ]
            
            # 运行 Agent 循环（限制迭代次数）
            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=self.model,
                )
                
                if response.has_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,
                    })
                    
                    # 执行工具
                    for tool_call in response.tool_calls:
                        logger.debug(f"Subagent [{task_id}] executing: {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
            
            if final_result is None:
                final_result = "Task completed but no final response was generated."
            
            logger.info(f"Subagent [{task_id}] completed successfully")
            await self._announce_result(task_id, label, task, final_result, origin, "ok")
            
        except Exception as e:
            error_msg = f"Error: {str(e)}"
            logger.error(f"Subagent [{task_id}] failed: {e}")
            await self._announce_result(task_id, label, task, error_msg, origin, "error")
    
    async def _announce_result(
        self,
        task_id: str,
        label: str,
        task: str,
        result: str,
        origin: dict[str, str],
        status: str,
    ) -> None:
        """
        通知子任务结果
        
        通过消息总线将结果发送给主 Agent，主 Agent 会将结果转换为自然语言通知用户。
        
        Args:
            task_id: 子任务 ID
            label: 显示标签
            task: 原始任务描述
            result: 执行结果
            origin: 原始目标信息
            status: 状态（"ok" 或 "error"）
        """
        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."""
        
        # 注入系统消息以触发主 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(f"Subagent [{task_id}] announced result to {origin['channel']}:{origin['chat_id']}")
    
    def _build_subagent_prompt(self, task: str) -> str:
        """
        构建子任务专用的系统提示词
        
        与主 Agent 不同，子任务的提示词更加专注：
        - 明确的任务目标
        - 禁止发起对话或子任务
        - 简洁但信息丰富
        
        Args:
            task: 任务描述
        
        Returns:
            str: 格式化的系统提示词
        """
        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:
        """
        获取当前运行中的子任务数量
        
        Returns:
            int: 运行中的子任务数量
        """
        return len(self._running_tasks)
