"""
CodeSearch 工具
使用 Exa Code API 搜索编程相关的上下文信息
支持搜索 API、库、SDK 的文档和代码示例
"""

from typing import Optional
from pydantic import BaseModel, Field
import httpx

from .base import BaseTool


# ============================================================================
# API 配置
# ============================================================================

EXA_MCP_BASE_URL = "https://mcp.exa.ai"
EXA_MCP_ENDPOINT = "/mcp"


# ============================================================================
# 输入参数
# ============================================================================

class CodeSearchInput(BaseModel):
    """CodeSearch 输入参数"""
    query: str = Field(
        ...,
        description="搜索查询，用于查找 API、库、SDK 的相关上下文。例如：'React useState hook examples', 'Python pandas dataframe filtering', 'Express.js middleware'"
    )
    tokens_num: int = Field(
        default=5000,
        ge=1000,
        le=50000,
        description="返回的 token 数量（1000-50000）。默认 5000。较小值用于聚焦查询，较大值用于获取完整文档。"
    )


# ============================================================================
# CodeSearch 工具
# ============================================================================

class CodeSearchTool(BaseTool):
    """
    代码搜索工具 - 使用 Exa Code API

    提供高质量的编程相关上下文，包括：
    - API 文档和参考
    - 代码示例
    - 库和 SDK 使用指南
    - 编程模式和最佳实践
    """

    name = "codesearch"
    description = """搜索并获取编程任务的相关上下文，使用 Exa Code API。

功能：
- 提供最高质量和最新的库、SDK、API 上下文
- 返回全面的代码示例、文档和 API 参考
- 针对特定编程模式和解决方案进行优化

使用场景：
- 查找框架、库的使用方法
- 获取 API 文档和示例代码
- 搜索编程概念和最佳实践

参数说明：
- query: 搜索查询（如 'React useState hook examples'）
- tokens_num: 返回 token 数量（1000-50000，默认 5000）
  - 较小值：聚焦查询
  - 较大值：完整文档

示例查询：
- 'React useState hook examples'
- 'Python pandas dataframe filtering'
- 'Express.js middleware'
- 'Next.js partial prerendering configuration'
"""
    args_schema = CodeSearchInput

    def __init__(self, timeout: float = 30.0):
        """
        初始化 CodeSearch 工具

        Args:
            timeout: 请求超时时间（秒）
        """
        self.timeout = timeout

    async def run(self, query: str, tokens_num: int = 5000) -> str:
        """
        执行代码搜索

        Args:
            query: 搜索查询
            tokens_num: 返回的 token 数量

        Returns:
            搜索结果
        """
        # 构建 MCP 请求
        mcp_request = {
            "jsonrpc": "2.0",
            "id": 1,
            "method": "tools/call",
            "params": {
                "name": "get_code_context_exa",
                "arguments": {
                    "query": query,
                    "tokensNum": tokens_num,
                }
            }
        }

        headers = {
            "accept": "application/json, text/event-stream",
            "content-type": "application/json",
        }

        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                response = await client.post(
                    f"{EXA_MCP_BASE_URL}{EXA_MCP_ENDPOINT}",
                    headers=headers,
                    json=mcp_request,
                )

                if response.status_code != 200:
                    error_text = await response.aread()
                    return f"代码搜索错误 (HTTP {response.status_code}): {error_text.decode()}"

                response_text = response.text

                # 解析 SSE 响应
                lines = response_text.split("\n")
                for line in lines:
                    if line.startswith("data: "):
                        import json
                        try:
                            data = json.loads(line[6:])
                            if (data.get("result") and
                                data["result"].get("content") and
                                len(data["result"]["content"]) > 0):
                                return data["result"]["content"][0]["text"]
                        except json.JSONDecodeError:
                            continue

                return f"未找到关于 '{query}' 的代码示例或文档。请尝试不同的查询，更具体地描述库或编程概念，或检查框架名称的拼写。"

        except httpx.TimeoutException:
            return "代码搜索请求超时，请稍后重试"
        except Exception as e:
            return f"代码搜索错误: {str(e)}"

