"""
PydanticAI LLM客户端模块
使用 pydantic_ai 库，通过 OpenAIProvider 支持任意 OpenAI 兼容端点。
自动从环境变量读取配置: OPENAI_API_KEY, OPENAI_BASE_URL, OPENAI_MODEL

参考: https://ai.pydantic.dev/models/openai/
"""

from collections.abc import AsyncGenerator
import logging
import os
from typing import Any

from .llm_interface import (
    ChatMessage,
    LLMFactory,
    LLMInterface,
    LLMResponse,
    MessageRole,
)


logger = logging.getLogger(__name__)


class PydanticAIClient(LLMInterface):
    """
    PydanticAI LLM客户端

    使用 pydantic_ai 的 Agent 进行对话，支持任意 OpenAI 兼容端点。
    自动从环境变量读取配置（无需指定 provider）：
    - OPENAI_API_KEY: API密钥
    - OPENAI_BASE_URL: API基础URL（默认: https://ai.pydantic.dev）
    - OPENAI_MODEL: 模型名称

    使用示例:
        export OPENAI_API_KEY="your-api-key"
        export OPENAI_BASE_URL="https://ai.pydantic.dev"
        export OPENAI_MODEL="MiniMax-M2.7"

        # 或直接使用
        client = PydanticAIClient(
            model="MiniMax-M2.7",
            api_key="your-api-key",
            base_url="https://ai.pydantic.dev"
        )
    """

    def __init__(
        self,
        model: str,
        api_key: str | None = None,
        base_url: str | None = None,
        **kwargs: Any,
    ) -> None:
        """
        初始化PydanticAI客户端
        Args:
            model: 模型名称（默认: 环境变量OPENAI_MODEL）
            api_key: API密钥（默认: 环境变量OPENAI_API_KEY）
            base_url: API基础URL（默认: 环境变量OPENAI_BASE_URL 或 https://ai.pydantic.dev）
            **kwargs: 其他配置参数
        """
        super().__init__(model=model, **kwargs)

        # 导入pydantic_ai
        try:
            from pydantic_ai import Agent
            from pydantic_ai.models.openai import OpenAIChatModel
            from pydantic_ai.providers.openai import OpenAIProvider
        except ImportError:
            raise ImportError(
                "pydantic_ai is not installed. Install it with: pip install pydantic-ai"
            )

        # 获取配置（环境变量优先）
        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
        self.base_url = (
            base_url or os.getenv("OPENAI_BASE_URL") or "https://ai.pydantic.dev"
        )
        self.model_name = model or os.getenv("OPENAI_MODEL")

        if not self.api_key:
            raise ValueError(
                "API密钥未提供。请设置 OPENAI_API_KEY 环境变量，"
                "或在 https://ai.pydantic.dev 获取API密钥。"
            )

        if not self.model_name:
            raise ValueError(
                "模型名称未提供。请设置 OPENAI_MODEL 环境变量，或通过参数传入模型名称。"
            )

        # 创建OpenAI兼容provider
        self._provider = OpenAIProvider(
            base_url=self.base_url,
            api_key=self.api_key,
        )

        # 创建模型实例
        self._model = OpenAIChatModel(
            self.model_name,
            provider=self._provider,
        )

        # 创建Agent
        self._agent = Agent(
            self._model,
            **kwargs,
        )

        logger.info(
            f"PydanticAI客户端初始化成功，模型: {self.model_name}，"
            f"base_url: {self.base_url}"
        )

    def _convert_message_role(self, role: MessageRole) -> str:
        """转换消息角色"""
        role_mapping = {
            MessageRole.SYSTEM: "system",
            MessageRole.USER: "user",
            MessageRole.ASSISTANT: "assistant",
        }
        return role_mapping.get(role, "user")

    async def chat(self, messages: list[ChatMessage]) -> LLMResponse:
        """
        发送聊天请求
        Args:
            messages: 消息列表
        Returns:
            LLM响应
        """
        try:
            # 构建对话历史
            conversation = []
            for msg in messages:
                role = self._convert_message_role(msg.role)
                conversation.append(f"{role}: {msg.content}")

            user_message = "\n".join(conversation)

            # 使用Agent进行对话
            result = await self._agent.run(user_message)

            # 获取结果
            content = ""
            if hasattr(result, "output"):
                content = str(result.output)
            else:
                content = str(result)

            return LLMResponse(
                content=content,
                model=self.model_name,
            )

        except Exception as e:
            logger.error(f"PydanticAI聊天请求失败: {e}")
            raise

    async def complete(self, prompt: str, **kwargs: Any) -> LLMResponse:
        """
        文本补全接口
        Args:
            prompt: 输入提示
            **kwargs: 其他参数
        Returns:
            LLM响应
        """
        try:
            # 使用Agent进行补全
            result = await self._agent.run(prompt)

            content = str(result.output) if hasattr(result, "output") else str(result)

            return LLMResponse(
                content=content,
                model=self.model_name,
            )

        except Exception as e:
            logger.error(f"PydanticAI补全请求失败: {e}")
            raise

    async def extract_structured(
        self, content: str, schema: dict[str, Any], **kwargs: Any
    ) -> dict[str, Any]:
        """
        结构化数据提取接口
        Args:
            content: 输入内容
            schema: 输出模式定义
            **kwargs: 其他参数
        Returns:
            提取的结构化数据
        """
        # 使用chat接口，传入schema作为提示
        prompt = f"""请从以下内容中提取结构化数据：

内容：
{content}

请按照以下schema格式返回JSON：
{schema}

只返回JSON，不要其他内容。"""

        messages = [ChatMessage(role=MessageRole.USER, content=prompt)]
        response = await self.chat(messages)

        # 解析JSON响应
        import json

        try:
            return json.loads(response.content)
        except json.JSONDecodeError:
            return {"error": "JSON解析失败", "content": response.content}

    async def chat_stream(
        self, messages: list[ChatMessage]
    ) -> AsyncGenerator[str, None]:
        """
        流式聊天请求
        Args:
            messages: 消息列表
        Yields:
            响应内容片段
        """
        try:
            # 构建对话历史
            conversation = []
            for msg in messages:
                role = self._convert_message_role(msg.role)
                conversation.append(f"{role}: {msg.content}")

            user_message = "\n".join(conversation)

            # 流式运行 - run_stream 返回一个可以用作异步迭代器的结果
            result = await self._agent.run(user_message)
            # 对于非流式结果，yield 完整输出
            if hasattr(result, "output"):
                yield str(result.output)
            else:
                yield str(result)

        except Exception as e:
            logger.error(f"PydanticAI流式聊天请求失败: {e}")
            raise


# 注册为默认provider（当设置OPENAI_API_KEY时自动使用）
LLMFactory.register("pydantic", PydanticAIClient)
LLMFactory.register("pydantic_ai", PydanticAIClient)
