"""WenxinProvider - 百度文心 ERNIE Provider

支持 ERNIE 4.5/5.0 系列模型。
鉴权：AK/SK → access_token（缓存 29 天）
API：百度千帆平台 OpenAI 兼容接口
"""

from __future__ import annotations

import logging
import time
from typing import Any

from agentkit.llm.providers.openai import OpenAICompatibleProvider
from agentkit.llm.protocol import LLMRequest, LLMResponse

logger = logging.getLogger(__name__)

# 文心模型到端点的映射
WENXIN_MODEL_MAP = {
    "ernie-4.5-turbo-128k": "ernie-4.5-turbo-128k",
    "ernie-5.0": "ernie-5.0",
    "ernie-x1.1": "ernie-x1.1",
    "ernie-4.0-8k": "ernie-4.0-8k",
    "ernie-3.5-8k": "ernie-3.5-8k",
}

# 默认 base URL（千帆 v2 OpenAI 兼容接口）
WENXIN_DEFAULT_BASE_URL = "https://qianfan.baidubce.com/v2"


class WenxinProvider(OpenAICompatibleProvider):
    """百度文心 ERNIE Provider

    通过千帆平台 v2 OpenAI 兼容接口调用文心模型。

    鉴权方式：
    - 方式1（推荐）：直接使用 API Key，走 OpenAI 兼容接口
    - 方式2（传统）：AK/SK 换取 access_token

    使用方式：
        provider = WenxinProvider(api_key="your_api_key")
        # 或使用 AK/SK
        provider = WenxinProvider(api_key="", access_key="ak", secret_key="sk")
    """

    def __init__(
        self,
        api_key: str = "",
        access_key: str | None = None,
        secret_key: str | None = None,
        base_url: str = WENXIN_DEFAULT_BASE_URL,
        default_model: str = "ernie-4.5-turbo-128k",
        **kwargs: Any,
    ):
        # If AK/SK provided, use token-based auth
        self._access_key = access_key
        self._secret_key = secret_key
        self._access_token: str | None = None
        self._token_expires_at: float = 0.0

        # Resolve API key
        effective_api_key = api_key
        if not api_key and access_key and secret_key:
            effective_api_key = "pending_token"  # Will be resolved on first request

        super().__init__(
            api_key=effective_api_key,
            base_url=base_url,
            default_model=default_model,
            **kwargs,
        )

    async def chat(self, request: LLMRequest) -> LLMResponse:
        """发送 chat 请求，处理文心特殊鉴权"""
        # Resolve access token if using AK/SK
        if self._access_key and self._secret_key and not self._api_key.startswith("pkf"):
            await self._ensure_access_token()
            if self._access_token:
                self._api_key = self._access_token

        # Map model name
        request.model = WENXIN_MODEL_MAP.get(request.model, request.model)

        return await super().chat(request)

    async def _ensure_access_token(self) -> None:
        """确保 access_token 有效（缓存 29 天）"""
        if self._access_token and time.time() < self._token_expires_at:
            return

        try:
            import httpx

            url = (
                f"https://aip.baidubce.com/oauth/2.0/token?"
                f"grant_type=client_credentials&client_id={self._access_key}"
                f"&client_secret={self._secret_key}"
            )

            async with httpx.AsyncClient(timeout=30.0) as client:
                response = await client.post(url)
                data = response.json()

                if "access_token" in data:
                    self._access_token = data["access_token"]
                    # Cache for 29 days (token valid for 30 days)
                    self._token_expires_at = time.time() + 29 * 86400
                    logger.info("Wenxin access token refreshed")
                else:
                    logger.error(f"Failed to get Wenxin access token: {data}")

        except Exception as e:
            logger.error(f"Wenxin token refresh failed: {e}")
