"""工具组合 - SequentialChain, ParallelFanOut, DynamicSelector

支持工具的高级组合模式：
- SequentialChain: 顺序执行，前一个输出作为后一个输入
- ParallelFanOut: 并行执行多个工具，结果合并
- DynamicSelector: 根据 LLM 判断动态选择工具
"""

import asyncio
import json
import logging
from typing import Any

from agentkit.tools.base import Tool

logger = logging.getLogger(__name__)


class SequentialChain(Tool):
    """顺序链 - 依次执行多个工具，前一个输出作为后一个输入

    示例::

        chain = SequentialChain(
            name="retrieve_and_summarize",
            tools=[retrieve_tool, summarize_tool],
        )
        result = await chain.execute(query="AI trends")
    """

    def __init__(
        self,
        name: str,
        description: str,
        tools: list[Tool],
        version: str = "1.0.0",
        tags: list[str] | None = None,
    ):
        super().__init__(
            name=name,
            description=description,
            version=version,
            tags=tags or ["composition", "sequential"],
        )
        self._chain_tools = tools

    async def execute(self, **kwargs) -> dict:
        """顺序执行所有工具"""
        result = kwargs

        for tool in self._chain_tools:
            if isinstance(result, dict):
                result = await tool.safe_execute(**result)
            else:
                raise TypeError(
                    f"SequentialChain: tool '{tool.name}' must return dict, got {type(result)}"
                )
            logger.debug(f"SequentialChain step '{tool.name}' completed")

        return result if isinstance(result, dict) else {"result": result}


class ParallelFanOut(Tool):
    """并行扇出 - 同时执行多个工具，结果合并

    示例::

        fan_out = ParallelFanOut(
            name="multi_source_search",
            tools=[web_search, db_search, cache_search],
        )
        result = await fan_out.execute(query="AI trends")
    """

    def __init__(
        self,
        name: str,
        description: str,
        tools: list[Tool],
        merge_strategy: str = "merge",
        version: str = "1.0.0",
        tags: list[str] | None = None,
    ):
        super().__init__(
            name=name,
            description=description,
            version=version,
            tags=tags or ["composition", "parallel"],
        )
        self._parallel_tools = tools
        self._merge_strategy = merge_strategy

    async def execute(self, **kwargs) -> dict:
        """并行执行所有工具并合并结果"""
        tasks = [tool.safe_execute(**kwargs) for tool in self._parallel_tools]
        results = await asyncio.gather(*tasks, return_exceptions=True)

        merged = {}
        errors = []

        for tool, result in zip(self._parallel_tools, results):
            if isinstance(result, Exception):
                logger.warning(f"ParallelFanOut tool '{tool.name}' failed: {result}")
                errors.append({"tool": tool.name, "error": str(result)})
            elif isinstance(result, dict):
                if self._merge_strategy == "namespace":
                    merged[tool.name] = result
                else:
                    merged.update(result)

        if errors:
            merged["_errors"] = errors

        return merged


class DynamicSelector(Tool):
    """动态选择器 - 根据输入动态选择合适的工具执行

    支持两种选择模式：
    - "keyword": 根据输入中的关键字匹配工具名/标签
    - "llm": 通过 LLM 判断选择最合适的工具

    示例::

        selector = DynamicSelector(
            name="smart_tool",
            tools=[search_tool, calculate_tool, translate_tool],
            mode="keyword",
        )
        result = await selector.execute(query="search for AI", _intent="search")
    """

    def __init__(
        self,
        name: str,
        description: str,
        tools: list[Tool],
        mode: str = "keyword",
        llm_client: Any = None,
        version: str = "1.0.0",
        tags: list[str] | None = None,
    ):
        super().__init__(
            name=name,
            description=description,
            version=version,
            tags=tags or ["composition", "dynamic"],
        )
        self._selector_tools = tools
        self._mode = mode
        self._llm_client = llm_client

    async def execute(self, **kwargs) -> dict:
        """动态选择并执行工具"""
        # 移除内部参数
        intent = kwargs.pop("_intent", None)

        if self._mode == "keyword":
            selected = self._select_by_keyword(kwargs, intent)
        elif self._mode == "llm":
            selected = await self._select_by_llm(kwargs)
        else:
            raise ValueError(f"Unknown selector mode: {self._mode}")

        if selected is None:
            return {
                "error": "no_matching_tool",
                "available": [t.name for t in self._selector_tools],
            }

        logger.info(f"DynamicSelector selected tool '{selected.name}'")
        return await selected.safe_execute(**kwargs)

    def _select_by_keyword(self, kwargs: dict, intent: str | None) -> Tool | None:
        """根据关键字匹配工具"""
        # 优先使用显式 intent
        if intent:
            for tool in self._selector_tools:
                if intent.lower() in tool.name.lower():
                    return tool
                if intent.lower() in " ".join(tool.tags).lower():
                    return tool

        # 从输入值中推断
        input_text = " ".join(str(v) for v in kwargs.values()).lower()
        best_match = None
        best_score = 0

        for tool in self._selector_tools:
            score = 0
            # 工具名匹配
            for word in tool.name.lower().split("_"):
                if word in input_text:
                    score += 2
            # 标签匹配
            for tag in tool.tags:
                if tag.lower() in input_text:
                    score += 1
            # 描述匹配
            for word in tool.description.lower().split():
                if word in input_text and len(word) > 3:
                    score += 1

            if score > best_score:
                best_score = score
                best_match = tool

        return best_match

    async def _select_by_llm(self, kwargs: dict) -> Tool | None:
        """通过 LLM 选择工具"""
        if self._llm_client is None:
            logger.warning("DynamicSelector: no LLM client, falling back to keyword mode")
            return self._select_by_keyword(kwargs, None)

        tool_descriptions = []
        for i, tool in enumerate(self._selector_tools):
            tool_descriptions.append(
                f"{i}. {tool.name}: {tool.description} (tags: {', '.join(tool.tags)})"
            )

        prompt = (
            f"Given the following input:\n{json.dumps(kwargs, ensure_ascii=False)}\n\n"
            f"Available tools:\n" + "\n".join(tool_descriptions) + "\n\n"
            f"Which tool number should be used? Reply with just the number."
        )

        try:
            if hasattr(self._llm_client, "chat"):
                response = await self._llm_client.chat(
                    messages=[{"role": "user", "content": prompt}],
                    model="gpt-4",
                    temperature=0,
                    max_tokens=10,
                )
            elif callable(self._llm_client):
                response = await self._llm_client(
                    messages=[{"role": "user", "content": prompt}],
                    model="gpt-4",
                    temperature=0,
                    max_tokens=10,
                )
            else:
                return self._select_by_keyword(kwargs, None)

            if isinstance(response, str):
                idx = int(response.strip())
            elif isinstance(response, dict):
                idx = int(response.get("content", "0").strip())
            else:
                idx = 0

            if 0 <= idx < len(self._selector_tools):
                return self._selector_tools[idx]

        except (ValueError, IndexError, Exception) as e:
            logger.warning(f"DynamicSelector LLM selection failed: {e}")

        return self._select_by_keyword(kwargs, None)
