"""
LLM 模块的协议定义
"""

from abc import ABC
from typing import TYPE_CHECKING, Any, Protocol, Union

from pydantic import BaseModel

from .models import ToolDefinition, ToolResult

if TYPE_CHECKING:
    from .models import LLMMessage, LLMResponse, LLMToolCall


class ToolCallData(BaseModel):
    """传递给 on_tool_start 的数据模型"""

    tool_name: str
    tool_args: dict[str, Any]


class ToolCallCompleteData(BaseModel):
    """传递给 on_tool_call_complete 的数据模型"""

    id: str
    name: str
    arguments: str
    result: "ToolResult"


class BaseCallbackHandler(ABC):
    """
    Agent/LLM 生命周期回调处理器的基类。
    下沉至 LLM 层以允许 ToolInvoker 直接调用。
    """

    async def on_agent_start(self, messages: list["LLMMessage"], **kwargs: Any) -> None:
        """在 AgentExecutor 开始运行时调用。"""
        pass

    async def on_model_start(
        self, model_name: str, messages: list["LLMMessage"], **kwargs: Any
    ) -> None:
        """在向LLM发起请求之前调用。"""
        pass

    async def on_model_end(
        self, response: "LLMResponse", duration: float, **kwargs: Any
    ) -> None:
        """在收到LLM响应之后调用。"""
        pass

    async def on_tool_start(
        self, tool_call: "LLMToolCall", data: ToolCallData, **kwargs: Any
    ) -> Union[ToolCallData, "ToolResult", None]:
        """
        在单个工具即将被执行时调用。

        返回:
            ToolCallData: 修改参数并继续执行
            ToolResult: 拦截执行并直接返回给模型
            None: 正常继续
        """
        pass

    async def on_tool_end(
        self,
        result: Union["ToolResult", None],
        error: Exception | None,
        tool_call: "LLMToolCall",
        duration: float,
        **kwargs: Any,
    ) -> None:
        """在单个工具执行完毕后调用，无论成功或失败。"""
        pass

    async def on_tool_call_complete(
        self, data: ToolCallCompleteData, **kwargs: Any
    ) -> None:
        """在工具调用完成并准备创建响应消息时调用。"""
        pass

    async def on_human_input_request(self, query: str, **kwargs: Any) -> str | None:
        """
        当 Agent 需要人类输入时调用。
        """
        return None

    async def on_agent_end(
        self, final_history: list["LLMMessage"], duration: float, **kwargs: Any
    ) -> None:
        """在 AgentExecutor 运行结束时调用。"""
        pass


class ToolExecutable(Protocol):
    """
    一个协议，定义了所有可被LLM调用的工具必须实现的行为。
    它将工具的"定义"（给LLM看）和"执行"（由框架调用）封装在一起。
    """

    async def get_definition(self) -> ToolDefinition:
        """
        异步地获取一个结构化的工具定义。
        """
        ...

    async def execute(self, context: Any | None = None, **kwargs: Any) -> ToolResult:
        """
        异步执行工具并返回一个结构化的结果。
        参数由LLM根据工具定义生成。

        Args:
            context: 运行时上下文 (RunContext)，可选注入
            **kwargs: 工具参数
        """
        ...


class ToolProvider(Protocol):
    """
    一个协议，定义了"工具提供者"的行为。
    工具提供者负责发现或实例化具体的 ToolExecutable 对象。
    """

    async def initialize(self) -> None:
        """
        异步初始化提供者。
        此方法应是幂等的，即多次调用只会执行一次初始化逻辑。
        用于执行耗时的I/O操作，如网络请求或启动子进程。
        """
        ...

    async def discover_tools(
        self,
        allowed_servers: list[str] | None = None,
        excluded_servers: list[str] | None = None,
    ) -> dict[str, ToolExecutable]:
        """
        异步发现此提供者提供的所有工具。
        在 `initialize` 成功调用后才应被调用。

        返回:
            一个从工具名称到 ToolExecutable 实例的字典。
        """
        ...

    async def get_tool_executable(
        self, name: str, config: dict[str, Any]
    ) -> ToolExecutable | None:
        """
        【保留】如果此提供者能处理名为 'name' 的工具，则返回一个可执行实例。
        此方法主要用于按需解析 ad-hoc 工具。
        """
        ...
