#  Copyright 2026 zhaoxi826
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

"""Tavily Web Search Tool Plugin for KiloStar.

Provides intelligent web search capabilities via Tavily API.
API key 取值优先级：调用参数 > GlobalStateMachine 中 ``tavily_search`` 工具配置 >
环境变量 ``TAVILY_API_KEY``。
"""

import os
from typing import List, Literal, Dict, Optional

from kilostar.plugin.tool_plugin.base_tool import BaseToolData
from tavily import AsyncTavilyClient


class TavilySearchToolData(BaseToolData):
    """Tavily 搜索工具的元数据：面向所有节点开放。"""

    is_system: bool = False
    action_scope: List[
        Literal[
            "control_node",
            "consciousness_node",
            "regulatory_node",
            "growth_node",
        ]
    ] = ["control_node", "consciousness_node", "regulatory_node"]
    config_args: Dict[str, str] = {
        "api_key": "",
        "max_results": "5",
        "search_depth": "basic",
        "include_answer": "true",
    }
    category: str = "search"


async def _resolve_api_key(explicit: Optional[str]) -> Optional[str]:
    """按优先级解析 Tavily API key：显式参数 > GSM 配置 > 环境变量。"""
    if explicit:
        return explicit
    try:
        from kilostar.core.global_state_machine.gsm_snapshot import fetch_snapshot

        # 工具调用是高频热路径，走 Object Store 快照而不是 actor RPC
        snapshot = await fetch_snapshot()
        cfg = snapshot.tool_configs.get("tavily_search") or {}
        if isinstance(cfg, dict) and cfg.get("api_key"):
            return cfg["api_key"]
    except Exception:
        pass
    return os.environ.get("TAVILY_API_KEY")


async def tavily_search(
    query: str,
    max_results: int = 5,
    search_depth: str = "basic",
    include_answer: bool = True,
    api_key: Optional[str] = None,
) -> str:
    """使用 Tavily 进行网络搜索，获取高质量的网络搜索结果。

    Args:
        query: 搜索查询内容
        max_results: 返回的最大结果数量（1-10）
        search_depth: 搜索深度，"basic" 或 "advanced"
        include_answer: 是否包含 AI 生成的答案摘要
        api_key: 可选；不传则按 GSM 配置 → 环境变量顺序解析

    Returns:
        格式化的搜索结果文本，包含标题、URL、摘要和可选的 AI 答案
    """
    resolved_key = await _resolve_api_key(api_key)
    if not resolved_key:
        return (
            "[Error] Tavily API key 未配置。"
            "请在 ``/api/v1/resource/tool/config`` 写入或设置环境变量 ``TAVILY_API_KEY``。"
        )

    try:
        client = AsyncTavilyClient(api_key=resolved_key)
        result = await client.search(
            query=query,
            max_results=min(max_results, 10),
            search_depth=search_depth,
            include_answer=include_answer,
        )

        lines = []
        if include_answer and result.get("answer"):
            lines.append(f"【AI 摘要】{result['answer']}\n")

        results = result.get("results", [])
        if not results:
            return "No results found for the query."

        lines.append("【搜索结果】")
        for i, item in enumerate(results, 1):
            title = item.get("title", "Untitled")
            url = item.get("url", "")
            content = item.get("content", "").strip()
            lines.append(f"\n{i}. {title}")
            lines.append(f"   URL: {url}")
            if content:
                lines.append(f"   {content[:300]}{'...' if len(content) > 300 else ''}")

        return "\n".join(lines)
    except Exception as e:
        return f"[Error] Tavily search failed: {str(e)}"