"""DataHub SDK - 统一的SDK接口。

提供便捷的上下文管理器接口，集成所有 DataHub API 功能。
"""

import logging
from pathlib import Path
from typing import TYPE_CHECKING, Any

from hezor_common.transfer.datahub_sdk.base.constants import (
    DEFAULT_ARGS_FILL_TIPS,
    DEFAULT_DATA_API_BASE_URL,
    DEFAULT_DATA_API_KEY,
)

if TYPE_CHECKING:
    from hezor_common.utilities.agentic.deps import Model

from hezor_common.data_model.searching.data.api_models import (
    ExecuteResponse,
    SearchResponse,
)
from hezor_common.transfer.datahub_sdk.base.data_api_client import (
    DataAPIClient,
)
from hezor_common.transfer.datahub_sdk.base.meta_info import MetaInfo
from hezor_common.transfer.datahub_sdk.build_args import (
    build_args,
    build_args_for_all_tools,
)
from hezor_common.transfer.datahub_sdk.execute_tool import (
    execute_tool,
    execute_tool_from_json,
)
from hezor_common.transfer.datahub_sdk.search_tools import search_tools
from hezor_common.transfer.datahub_sdk.tool_builder import build_tool_functions

logger = logging.getLogger(__name__)


class DatahubSDK:
    """DataHub SDK 统一接口。

    提供便捷的上下文管理器接口，集成搜索、参数构建、工具执行等功能。

    Parameters
    ----------
    base_url : str
        API 基础URL，默认使用 DEFAULT_API_BASE_URL
    timeout : float
        请求超时时间（秒）
    api_key : str | None
        API 密钥，用于认证
    meta_info : MetaInfo | None
        元信息对象，用于签名认证
    private_key_path : str | Path | None
        私钥文件路径，用于签名认证（与 meta_info 配合使用）
    password : bytes | None
        私钥密码
    meta_info_expires_in : int
        MetaInfo JWT token过期时间（秒）
    model : Model | None
        AI 模型，用于参数构建，默认为 OpenAIChat(id="hezor2")

    Examples
    --------
    >>> # 基本使用
    >>> async with DatahubSDK() as sdk:
    ...     # 搜索工具
    ...     search_resp = await sdk.search_tools("查询天气")
    ...     # 构建参数
    ...     args = await sdk.build_args(search_resp, "查询北京的天气")
    ...     # 执行工具
    ...     result = await sdk.execute_tool(search_resp.tools[0].name, args)

    >>> # 使用认证
    >>> async with DatahubSDK(api_key="your-key") as sdk:
    ...     search_resp = await sdk.search_tools("查询数据", top_k=5)

    >>> # 完整工作流
    >>> async with DatahubSDK() as sdk:
    ...     # 搜索并执行
    ...     result = await sdk.search_build_and_execute(
    ...         query="查询天气",
    ...         args_fill_tips="查询北京今天的天气"
    ...     )
    """

    def __init__(
        self,
        base_url: str = DEFAULT_DATA_API_BASE_URL,
        timeout: float = 30.0,
        api_key: str | None = DEFAULT_DATA_API_KEY,
        meta_info: MetaInfo | None = None,
        private_key_path: str | Path | None = None,
        password: bytes | None = None,
        meta_info_expires_in: int = 3600,
        model: "Model | None" = None,
    ):
        """初始化 DataHub SDK。

        Parameters
        ----------
        base_url : str
            API 基础URL
        timeout : float
            请求超时时间（秒）
        api_key : str | None
            API 密钥
        meta_info : MetaInfo | None
            元信息对象
        private_key_path : str | Path | None
            私钥文件路径
        password : bytes | None
            私钥密码
        meta_info_expires_in : int
            MetaInfo JWT token过期时间（秒）
        model : Model | None
            AI 模型，用于参数构建，默认为 OpenAIChat(id="hezor2")
        """
        self._client = DataAPIClient(
            base_url=base_url,
            timeout=timeout,
            api_key=api_key,
            meta_info=meta_info,
            private_key_path=private_key_path,
            password=password,
            meta_info_expires_in=meta_info_expires_in,
        )
        self._model = model

    async def __aenter__(self) -> "DatahubSDK":
        """进入异步上下文管理器。

        Returns
        -------
        DatahubSDK
            SDK 实例
        """
        await self._client.__aenter__()
        logger.debug("DatahubSDK context manager entered")
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        """退出异步上下文管理器。

        Parameters
        ----------
        exc_type : type | None
            异常类型
        exc_val : Exception | None
            异常值
        exc_tb : TracebackType | None
            异常traceback
        """
        await self._client.__aexit__(exc_type, exc_val, exc_tb)
        logger.debug("DatahubSDK context manager exited")

    async def search_tools(
        self,
        query: str,
        top_k: int = 3,
    ) -> SearchResponse:
        """搜索工具。

        Parameters
        ----------
        query : str
            搜索查询
        top_k : int
            返回结果数量，默认为3

        Returns
        -------
        SearchResponse
            搜索结果

        Examples
        --------
        >>> async with DatahubSDK() as sdk:
        ...     resp = await sdk.search_tools("查询天气", top_k=5)
        """
        return await search_tools(
            query=query,
            top_k=top_k,
            base_url=self._client.base_url,
            api_key=self._client.api_key,
            meta_info=self._client.meta_info,
            private_key_path=self._client.private_key_path,
            password=self._client.password,
            meta_info_expires_in=self._client.meta_info_expires_in,
        )

    async def build_args(
        self,
        search_response: SearchResponse,
        args_fill_tips: str = "请根据工具的参数说明填写合适的参数值",
        max_retries: int = 3,
        model_kwargs: dict[str, Any] | None = None,
        agent_kwargs: dict[str, Any] | None = None,
        tool_index: int = 0,
    ) -> dict[str, Any]:
        """为单个工具构建参数。

        Parameters
        ----------
        search_response : SearchResponse
            工具搜索响应
        args_fill_tips : str
            参数填写提示
        max_retries : int
            最大重试次数
        model_kwargs : dict[str, Any] | None
            模型参数
        agent_kwargs : dict[str, Any] | None
            Agent 参数
        tool_index : int
            工具索引

        Returns
        -------
        dict[str, Any]
            工具参数

        Examples
        --------
        >>> async with DatahubSDK() as sdk:
        ...     search_resp = await sdk.search_tools("查询天气")
        ...     args = await sdk.build_args(search_resp, "查询北京的天气")
        """
        return await build_args(
            search_response=search_response,
            args_fill_tips=args_fill_tips,
            model=self._model,
            max_retries=max_retries,
            model_kwargs=model_kwargs,
            agent_kwargs=agent_kwargs,
            tool_index=tool_index,
        )

    async def build_args_for_all_tools(
        self,
        search_response: SearchResponse,
        args_fill_tips: str = "请根据工具的参数说明填写合适的参数值",
        max_retries: int = 3,
        model_kwargs: dict[str, Any] | None = None,
        agent_kwargs: dict[str, Any] | None = None,
    ) -> dict[str, dict[str, Any]]:
        """为所有工具构建参数。

        Parameters
        ----------
        search_response : SearchResponse
            工具搜索响应
        args_fill_tips : str
            参数填写提示
        max_retries : int
            最大重试次数
        model_kwargs : dict[str, Any] | None
            模型参数

        Returns
        -------
        dict[str, dict[str, Any]]
            所有工具的参数，键为工具名称

        Examples
        --------
        >>> async with DatahubSDK() as sdk:
        ...     search_resp = await sdk.search_tools("查询天气")
        ...     all_args = await sdk.build_args_for_all_tools(search_resp, "查询北京的天气")
        """
        return await build_args_for_all_tools(
            search_response=search_response,
            args_fill_tips=args_fill_tips,
            model=self._model,
            max_retries=max_retries,
            model_kwargs=model_kwargs,
            agent_kwargs=agent_kwargs,
        )

    async def execute_tool(
        self,
        tool_name: str,
        args: dict[str, Any],
    ) -> ExecuteResponse:
        """执行工具。

        Parameters
        ----------
        tool_name : str
            工具名称
        args : dict[str, Any]
            工具参数

        Returns
        -------
        ExecuteResponse
            执行结果

        Examples
        --------
        >>> async with DatahubSDK() as sdk:
        ...     result = await sdk.execute_tool("get_weather", {"city": "Beijing"})
        """
        return await execute_tool(
            tool_name=tool_name,
            args=args,
            base_url=self._client.base_url,
            api_key=self._client.api_key,
            meta_info=self._client.meta_info,
            private_key_path=self._client.private_key_path,
            password=self._client.password,
            meta_info_expires_in=self._client.meta_info_expires_in,
        )

    async def execute_tool_from_json(
        self,
        request_json: str,
    ) -> ExecuteResponse:
        """从 JSON 字符串执行工具。

        Parameters
        ----------
        request_json : str
            包含工具名称和参数的 JSON 字符串

        Returns
        -------
        ExecuteResponse
            执行结果

        Examples
        --------
        >>> async with DatahubSDK() as sdk:
        ...     result = await sdk.execute_tool_from_json(
        ...         '{"tool": "get_weather", "args": {"city": "Beijing"}}'
        ...     )
        """
        return await execute_tool_from_json(
            request_json=request_json,
            base_url=self._client.base_url,
            api_key=self._client.api_key,
            meta_info=self._client.meta_info,
            private_key_path=self._client.private_key_path,
            password=self._client.password,
            meta_info_expires_in=self._client.meta_info_expires_in,
        )

    def build_tool_functions(
        self,
        tools: SearchResponse | list[Any],
    ) -> dict[str, Any]:
        """构建工具函数。

        Parameters
        ----------
        tools : SearchResponse | list[Any]
            搜索响应或工具列表

        Returns
        -------
        dict[str, Any]
            工具函数字典

        Examples
        --------
        >>> async with DatahubSDK() as sdk:
        ...     search_resp = await sdk.search_tools("查询天气")
        ...     funcs = sdk.build_tool_functions(search_resp)
        ...     # 使用生成的函数
        ...     result = await funcs["get_weather"](city="Beijing")
        """
        return build_tool_functions(
            tools=tools,
            base_url=self._client.base_url,
            api_key=self._client.api_key,
            meta_info=self._client.meta_info,
            private_key_path=self._client.private_key_path,
            password=self._client.password,
            meta_info_expires_in=self._client.meta_info_expires_in,
        )

    async def search_build_and_execute(
        self,
        query: str,
        args_fill_tips: str = DEFAULT_ARGS_FILL_TIPS,
        max_retries: int = 3,
        model_kwargs: dict[str, Any] | None = None,
        agent_kwargs: dict[str, Any] | None = None,
    ) -> ExecuteResponse:
        """一站式工作流：搜索、构建参数、执行工具。

        自动搜索最匹配的工具（top_k=1），构建参数并执行。

        Parameters
        ----------
        query : str
            搜索查询
        args_fill_tips : str
            参数填写提示
        max_retries : int
            最大重试次数
        model_kwargs : dict[str, Any] | None
            模型参数
        agent_kwargs : dict[str, Any] | None
            Agent 参数

        Returns
        -------
        ExecuteResponse
            执行结果

        Examples
        --------
        >>> async with DatahubSDK() as sdk:
        ...     result = await sdk.search_build_and_execute(
        ...         query="查询天气",
        ...         args_fill_tips="查询北京今天的天气，需要温度和湿度"
        ...     )
        """
        # 搜索工具（只获取最匹配的1个）
        logger.info(f"搜索工具: {query}")
        search_resp = await self.search_tools(query=query, top_k=1)

        if not search_resp.tools:
            raise ValueError(f"未找到匹配的工具: {query}")

        # 构建参数（使用第一个工具）
        tool_name = search_resp.tools[0].name
        logger.info(f"为工具 '{tool_name}' 构建参数")
        args = await self.build_args(
            search_response=search_resp,
            args_fill_tips=args_fill_tips,
            max_retries=max_retries,
            model_kwargs=model_kwargs,
            agent_kwargs=agent_kwargs,
            tool_index=0,
        )

        # 执行工具
        logger.info(f"执行工具 '{tool_name}' 参数: {args}")
        return await self.execute_tool(tool_name=tool_name, args=args)

    async def search_build_and_execute_all(
        self,
        query: str,
        args_fill_tips: str = DEFAULT_ARGS_FILL_TIPS,
        top_k: int = 3,
        max_retries: int = 3,
        model_kwargs: dict[str, Any] | None = None,
        agent_kwargs: dict[str, Any] | None = None,
    ) -> dict[str, ExecuteResponse]:
        """一站式工作流：搜索、为所有工具构建参数、批量执行。

        Parameters
        ----------
        query : str
            搜索查询
        args_fill_tips : str
            参数填写提示
        top_k : int
            搜索结果数量
        max_retries : int
            最大重试次数
        model_kwargs : dict[str, Any] | None
            模型参数
        agent_kwargs : dict[str, Any] | None
            Agent 参数

        Returns
        -------
        dict[str, ExecuteResponse]
            所有工具的执行结果，键为工具名称

        Examples
        --------
        >>> async with DatahubSDK() as sdk:
        ...     results = await sdk.search_build_and_execute_all(
        ...         query="查询天气",
        ...         args_fill_tips="查询北京今天的天气"
        ...     )
        ...     for tool_name, result in results.items():
        ...         print(f"{tool_name}: {result}")
        """
        # 搜索工具
        logger.info(f"搜索工具: {query}")
        search_resp = await self.search_tools(query=query, top_k=top_k)

        if not search_resp.tools:
            raise ValueError(f"未找到匹配的工具: {query}")

        # 为所有工具构建参数
        logger.info(f"为 {len(search_resp.tools)} 个工具构建参数")
        all_args = await self.build_args_for_all_tools(
            search_response=search_resp,
            args_fill_tips=args_fill_tips,
            max_retries=max_retries,
            model_kwargs=model_kwargs,
            agent_kwargs=agent_kwargs,
        )

        # 批量执行所有工具
        results: dict[str, ExecuteResponse] = {}
        for tool_name, args in all_args.items():
            logger.info(f"执行工具 '{tool_name}' 参数: {args}")
            result = await self.execute_tool(tool_name=tool_name, args=args)
            results[tool_name] = result

        return results

    async def make_call(
        self,
        query: str,
        context: str = "",
        max_retries: int = 3,
        model_kwargs: dict[str, Any] | None = None,
        agent_kwargs: dict[str, Any] | None = None,
    ) -> ExecuteResponse:
        """快捷调用：自动匹配工具并执行。

        语义化的快捷方法，自动搜索最匹配的工具、构建参数并执行。
        相比 search_build_and_execute，提供更简洁的语义。

        Parameters
        ----------
        query : str
            工具调用意图（用于搜索工具）
        context : str
            调用上下文或参数提示，默认为空
        max_retries : int
            调用失败时的最大重试次数，默认为 3
        model_kwargs : dict[str, Any] | None
            模型参数，可选
        agent_kwargs : dict[str, Any] | None
            Agent 参数，可选

        Returns
        -------
        ExecuteResponse
            执行结果

        Raises
        ------
        Exception
            当所有重试都失败时，抛出最后一次的异常

        Examples
        --------
        >>> async with DatahubSDK() as sdk:
        ...     # 简洁调用
        ...     result = await sdk.make_call("查询天气")
        ...
        ...     # 带上下文和重试
        ...     result = await sdk.make_call(
        ...         "查询天气",
        ...         context="查询北京今天的天气，需要温度和湿度",
        ...         max_retries=5
        ...     )
        """
        args_fill_tips = context if context else DEFAULT_ARGS_FILL_TIPS

        last_error: Exception | None = None
        for attempt in range(max_retries):
            try:
                logger.debug(f"make_call 尝试 {attempt + 1}/{max_retries}: query={query}")
                return await self.search_build_and_execute(
                    query=query,
                    args_fill_tips=args_fill_tips,
                    model_kwargs=model_kwargs,
                    agent_kwargs=agent_kwargs,
                )
            except Exception as e:
                last_error = e
                logger.warning(f"make_call 尝试 {attempt + 1}/{max_retries} 失败: {e}")
                if attempt < max_retries - 1:
                    logger.info(f"准备第 {attempt + 2} 次重试...")
                continue

        # 所有重试都失败
        logger.error(f"make_call 所有 {max_retries} 次尝试均失败")
        if last_error:
            raise last_error
        raise RuntimeError("make_call 失败且无异常信息")

    async def make_call_multi(
        self,
        query: str,
        context: str = "",
        top_k: int = 3,
        max_retries: int = 3,
        model_kwargs: dict[str, Any] | None = None,
        agent_kwargs: dict[str, Any] | None = None,
    ) -> dict[str, ExecuteResponse]:
        """快捷调用（多工具）：自动匹配多个工具并批量执行。

        search_build_and_execute_all 的语义化别名。
        搜索 top_k 个最匹配的工具，为每个工具构建参数并批量执行。

        Parameters
        ----------
        query : str
            工具调用意图（用于搜索工具）
        context : str
            调用上下文或参数提示，默认为空
        top_k : int
            返回工具数量，默认为 3
        max_retries : int
            参数构建失败时的最大重试次数，默认为 3
        model_kwargs : dict[str, Any] | None
            模型参数，可选
        agent_kwargs : dict[str, Any] | None
            Agent 参数，可选

        Returns
        -------
        dict[str, ExecuteResponse]
            所有工具的执行结果，键为工具名称

        Examples
        --------
        >>> async with DatahubSDK() as sdk:
        ...     # 批量调用多个工具
        ...     results = await sdk.make_call_multi("查询天气", top_k=3)
        ...     for tool_name, result in results.items():
        ...         print(f"{tool_name}: {result}")
        ...
        ...     # 带上下文
        ...     results = await sdk.make_call_multi(
        ...         "查询天气",
        ...         context="查询北京今天的天气，需要温度和湿度",
        ...         top_k=5
        ...     )
        """
        args_fill_tips = context if context else DEFAULT_ARGS_FILL_TIPS

        return await self.search_build_and_execute_all(
            query=query,
            args_fill_tips=args_fill_tips,
            top_k=top_k,
            max_retries=max_retries,
            model_kwargs=model_kwargs,
            agent_kwargs=agent_kwargs,
        )
