#  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.

from typing import Sequence, Any

from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.models.anthropic import AnthropicModel
from pydantic_ai.models.gemini import GeminiModel
from pydantic_ai.providers.openai import OpenAIProvider
from pydantic_ai.providers.anthropic import AnthropicProvider
from pydantic_ai.providers.deepseek import DeepSeekProvider
from pydantic_ai.providers.google import GoogleProvider
from pydantic_ai.toolsets import AbstractToolset

from kilostar.core.global_state_machine.model_provider import Provider
from kilostar.utils.agent_model import ResponseModel, DepsModel
from kilostar.utils.error import ModelNotExistError


class AgentFactory:
    """模型工厂：把内部的 ``Provider`` 元数据翻译成 pydantic-ai 的 ``Agent``。

    支持 openai / claude / deepseek / gemini 四类后端，差异通过
    ``_models_mapping`` 中的 ``model_class`` + ``provider_class`` 键值对屏蔽。
    同时支持传入本地工具（tools）和外部工具集（toolsets），包括 MCP 服务器。
    """

    def __init__(self):
        self._models_mapping = {
            "openai": {
                "model_class": OpenAIChatModel,
                "provider_class": OpenAIProvider,
                "provider_kwargs": {"base_url": True, "api_key": True},
            },
            "claude": {
                "model_class": AnthropicModel,
                "provider_class": AnthropicProvider,
                "provider_kwargs": {"api_key": True},
            },
            "deepseek": {
                "model_class": OpenAIChatModel,
                "provider_class": DeepSeekProvider,
                "provider_kwargs": {"api_key": True},
            },
            "gemini": {
                "model_class": GeminiModel,
                "provider_class": GoogleProvider,
                "provider_kwargs": {"api_key": True},
            },
        }

    def create_agent(
        self,
        provider: Provider,
        model_id: str,
        output_type: ResponseModel,
        system_prompt: str,
        deps_type: DepsModel,
        agent_name: str,
        tools: list = None,
        toolsets: Sequence[AbstractToolset[Any]] = None,
    ) -> Agent:
        """将输入的 provider 对象实例化为一个 pydantic-ai 的 agent 对象。

        Args:
            provider: Provider 对象，从 global_state_machine 中获取
            model_id: 模型名
            output_type: 输出格式
            system_prompt: 系统提示词
            deps_type: 依赖类型，在 agent 运行时动态输入的格式化消息
            agent_name: agent 的名字
            tools: 本地工具函数列表
            toolsets: 外部工具集列表（包括 MCP 服务器等 AbstractToolset 实例）

        Returns:
            被实例化的 pydantic-ai 的 Agent 对象
        """
        if model_id not in provider.provider_models:
            raise ModelNotExistError("模型不存在")
        if provider.provider_type not in self._models_mapping:
            raise ValueError(f"不支持的协议类型: {provider.provider_type}")

        config = self._models_mapping[provider.provider_type]
        model_class = config["model_class"]
        provider_class = config["provider_class"]
        provider_kwargs = config["provider_kwargs"]

        # 构建 provider 实例化参数
        init_kwargs = {}
        if provider_kwargs.get("api_key"):
            init_kwargs["api_key"] = provider.provider_apikey
        if provider_kwargs.get("base_url"):
            init_kwargs["base_url"] = provider.provider_url

        model_provider = provider_class(**init_kwargs)

        # 对于 Gemini，provider 需要传递给 model
        if provider.provider_type == "gemini":
            model = model_class(
                model_name=model_id,
                provider=model_provider,
            )
        else:
            model = model_class(model_id, provider=model_provider)

        # 创建 Agent，同时传入 tools 和 toolsets
        agent = Agent(
            model=model,
            name=agent_name,
            system_prompt=system_prompt,
            output_type=output_type,
            deps_type=deps_type,
            tools=tools or [],
            toolsets=toolsets or [],
        )
        return agent