"""子Agent委派引擎 — 替代 DeerFlow SubagentExecutor + task_tool"""

import json
import logging
import threading
import uuid
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any, Optional

from openai import OpenAI
from app.agent.messages import AIMessage, HumanMessage, SystemMessage
from app.agent.tools.decorator import tool, get_tool_schema

logger = logging.getLogger(__name__)


# ==================== 数据结构 ====================

class SubagentStatus(str, Enum):
    """子Agent任务状态"""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    TIMED_OUT = "timed_out"


@dataclass
class SubagentResult:
    """子Agent任务执行结果"""
    task_id: str
    status: SubagentStatus = SubagentStatus.PENDING
    result: str = ""
    error: str = ""
    started_at: Optional[datetime] = None
    completed_at: Optional[datetime] = None
    ai_messages: list[dict] = field(default_factory=list)


@dataclass
class SubagentConfig:
    """子Agent配置"""
    name: str
    description: str
    system_prompt: str
    tools: list[str] = field(default_factory=list)          # whitelist (empty = all allowed)
    disallowed_tools: list[str] = field(default_factory=lambda: ["task"])
    model: str = "inherit"                                   # "inherit" = use parent model
    max_turns: int = 50
    timeout_seconds: int = 300


# ==================== 子Agent注册表 ====================

SUBAGENT_REGISTRY: dict[str, SubagentConfig] = {
    "general-purpose": SubagentConfig(
        name="general-purpose",
        description="通用子Agent，用于法律分析、法规检索、案例研究等复杂任务",
        system_prompt="""你是综合执法智能助手的子Agent，专门负责完成父Agent委派的具体任务。

你的职责：
1. 严格按照任务描述完成指定的分析、检索或研究工作
2. 使用可用的工具（如法规搜索、案例搜索等）获取必要信息
3. 提供结构化、准确、有法律依据的分析结果
4. 如果信息不足，明确说明缺失的部分

工作原则：
- 专注完成委派任务，不扩展任务范围
- 引用法律条文时必须准确，标注出处
- 分析结果必须客观、有据可依
- 使用中文回复""",
        tools=[],               # all tools allowed
        disallowed_tools=["task"],  # prevent recursive subagent spawning
        model="inherit",
        max_turns=50,
        timeout_seconds=300,
    ),
}


# ==================== 全局状态 ====================

_background_tasks: dict[str, SubagentResult] = {}
_bg_lock = threading.Lock()
_executor_pool = ThreadPoolExecutor(max_workers=3, thread_name_prefix="subagent")


# ==================== SubagentEngine ====================

class SubagentEngine:
    """子Agent委派引擎 — 管理子Agent的创建、执行和结果收集"""

    def __init__(self, parent_llm: OpenAI, all_tools: list, parent_model_name: str = ""):
        """
        Args:
            parent_llm: 父Agent使用的 OpenAI 客户端实例
            all_tools: 所有可用工具列表
            parent_model_name: 父Agent使用的模型名称（用于inherit）
        """
        self.parent_llm = parent_llm
        self.all_tools = all_tools
        self.parent_model_name = parent_model_name
        # Extract base_url and api_key from the OpenAI client for creating new instances
        self.parent_base_url = str(parent_llm.base_url) if parent_llm.base_url else ""
        self.parent_api_key = parent_llm.api_key or ""

    def _filter_tools(self, config: SubagentConfig) -> list:
        """根据配置过滤工具白名单/黑名单"""
        filtered = []
        for t in self.all_tools:
            # 黑名单优先
            if t.name in config.disallowed_tools:
                continue
            # 白名单为空表示全部允许
            if not config.tools or t.name in config.tools:
                filtered.append(t)
        return filtered

    def _resolve_model(self, config: SubagentConfig) -> OpenAI:
        """解析子Agent使用的模型"""
        if config.model == "inherit":
            return self.parent_llm
        # 创建新模型实例
        return OpenAI(
            base_url=self.parent_base_url,
            api_key=self.parent_api_key,
        )

    def execute(
        self,
        task_description: str,
        task_prompt: str,
        subagent_type: str = "general-purpose",
        task_id: str = "",
    ) -> SubagentResult:
        """同步执行子Agent任务

        Args:
            task_description: 任务简短描述
            task_prompt: 完整的任务提示词
            subagent_type: 子Agent类型（对应 SUBAGENT_REGISTRY 中的 key）
            task_id: 任务ID（为空则自动生成）

        Returns:
            SubagentResult 执行结果
        """
        if not task_id:
            task_id = str(uuid.uuid4())

        config = SUBAGENT_REGISTRY.get(subagent_type, SUBAGENT_REGISTRY["general-purpose"])
        result = SubagentResult(task_id=task_id, status=SubagentStatus.RUNNING,
                                started_at=datetime.now())

        with _bg_lock:
            _background_tasks[task_id] = result

        try:
            llm = self._resolve_model(config)
            tools = self._filter_tools(config)

            # Build OpenAI tool schemas
            tool_schemas = [get_tool_schema(t) for t in tools] if tools else None

            # Build messages
            system_prompt = config.system_prompt
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": task_prompt},
            ]

            # ReAct loop
            model_name = self.parent_model_name if config.model == "inherit" else config.model
            for _ in range(config.max_turns):
                response = llm.chat.completions.create(
                    model=model_name,
                    messages=messages,
                    tools=tool_schemas,
                    temperature=0.5,
                )
                choice = response.choices[0]
                msg = choice.message

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

                    # Execute tools
                    tool_map = {t.name: t for t in tools}
                    for tc in msg.tool_calls:
                        tool_name = tc.function.name
                        try:
                            tool_args = json.loads(tc.function.arguments)
                        except json.JSONDecodeError:
                            tool_args = {}
                        tool_func = tool_map.get(tool_name)
                        if tool_func:
                            try:
                                tool_result = tool_func(**tool_args)
                            except Exception as e:
                                tool_result = f"[工具执行错误] {e}"
                        else:
                            tool_result = f"[工具未找到] {tool_name}"
                        result.ai_messages.append({
                            "tool": tool_name,
                            "tool_input": tool_args,
                            "observation": str(tool_result)[:500],
                        })
                        messages.append({
                            "role": "tool",
                            "tool_call_id": tc.id,
                            "content": str(tool_result),
                        })
                else:
                    # Final answer
                    result.result = msg.content or ""
                    break

            result.status = SubagentStatus.COMPLETED

        except Exception as e:
            logger.exception(f"Subagent task {task_id} failed: {e}")
            result.status = SubagentStatus.FAILED
            result.error = str(e)

        finally:
            result.completed_at = datetime.now()
            with _bg_lock:
                _background_tasks[task_id] = result

        return result

    def execute_async(
        self,
        task_description: str,
        task_prompt: str,
        subagent_type: str = "general-purpose",
        task_id: str = "",
    ) -> str:
        """异步执行子Agent任务（通过线程池）

        Returns:
            task_id 用于后续查询结果
        """
        if not task_id:
            task_id = str(uuid.uuid4())

        # Register pending result
        result = SubagentResult(task_id=task_id, status=SubagentStatus.PENDING)
        with _bg_lock:
            _background_tasks[task_id] = result

        # Submit to thread pool
        _executor_pool.submit(
            self.execute,
            task_description,
            task_prompt,
            subagent_type,
            task_id,
        )

        return task_id

    @staticmethod
    def get_result(task_id: str) -> Optional[SubagentResult]:
        """获取任务执行结果"""
        with _bg_lock:
            return _background_tasks.get(task_id)

    @staticmethod
    def poll_until_complete(
        task_id: str,
        poll_interval: float = 0.5,
        timeout: float = 300.0,
    ) -> SubagentResult:
        """轮询等待任务完成

        Args:
            task_id: 任务ID
            poll_interval: 轮询间隔（秒）
            timeout: 超时时间（秒）

        Returns:
            SubagentResult 最终结果

        Raises:
            TimeoutError: 超时未完成
        """
        import time
        start = time.time()
        while True:
            result = SubagentEngine.get_result(task_id)
            if result is None:
                raise ValueError(f"Task {task_id} not found")
            if result.status in (SubagentStatus.COMPLETED, SubagentStatus.FAILED, SubagentStatus.TIMED_OUT):
                return result
            if time.time() - start > timeout:
                result.status = SubagentStatus.TIMED_OUT
                result.error = f"Task timed out after {timeout}s"
                with _bg_lock:
                    _background_tasks[task_id] = result
                return result
            time.sleep(poll_interval)


# ==================== task 工具 ====================

# Module-level reference to the engine, set by StandaloneAgentEngine
_engine_ref: Optional[SubagentEngine] = None


def set_subagent_engine(engine: SubagentEngine) -> None:
    """设置全局 SubagentEngine 引用（由 StandaloneAgentEngine 调用）"""
    global _engine_ref
    _engine_ref = engine


@tool
def task(task_description: str, subagent_type: str = "general-purpose") -> str:
    """委派子任务给子Agent执行。用于将复杂任务分解为并行子任务。

    Args:
        task_description: 要委派的任务描述，需清晰完整地说明任务目标、要求和约束
        subagent_type: 子Agent类型，默认为 general-purpose
    """
    if _engine_ref is None:
        return "[错误] 子Agent引擎未初始化，无法执行委派任务"

    task_id = str(uuid.uuid4())

    try:
        # Synchronous execution (agent tool calls are blocking)
        result = _engine_ref.execute(
            task_description=task_description,
            task_prompt=task_description,
            subagent_type=subagent_type,
            task_id=task_id,
        )

        if result.status == SubagentStatus.COMPLETED:
            return result.result
        elif result.status == SubagentStatus.FAILED:
            return f"[子Agent执行失败] {result.error}"
        elif result.status == SubagentStatus.TIMED_OUT:
            return f"[子Agent执行超时] 任务在 {_engine_ref.parent_model_name} 上超时"
        else:
            return f"[子Agent状态异常] {result.status}"

    except Exception as e:
        logger.exception(f"task tool error: {e}")
        return f"[子Agent调用异常] {str(e)}"