"""YuanbaoProvider - 腾讯混元/元宝 Provider

支持 Hunyuan 2.0/T1 系列模型。
API：腾讯云 OpenAI 兼容接口
鉴权：Bearer API Key
"""

from __future__ import annotations

import logging
from typing import Any

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

logger = logging.getLogger(__name__)

# 混元模型映射
YUANBAO_MODEL_MAP = {
    "hunyuan-turbos-latest": "hunyuan-turbos-latest",
    "hunyuan-2.0": "hunyuan-2.0",
    "hunyuan-t1": "hunyuan-t1",
    "hunyuan-vision-1.5": "hunyuan-vision-1.5",
}

# 腾讯混元 API base URL
YUANBAO_DEFAULT_BASE_URL = "https://api.hunyuan.cloud.tencent.com/v1"


class YuanbaoProvider(OpenAICompatibleProvider):
    """腾讯混元/元宝 Provider

    通过腾讯云 OpenAI 兼容接口调用混元模型。

    使用方式：
        provider = YuanbaoProvider(
            api_key="your_hunyuan_api_key",
            default_model="hunyuan-turbos-latest",
        )

    特殊参数：
    - enable_enhancement: 增强模式（通过 LLMRequest._extra 传递）
    """

    def __init__(
        self,
        api_key: str,
        base_url: str = YUANBAO_DEFAULT_BASE_URL,
        default_model: str = "hunyuan-turbos-latest",
        enable_enhancement: bool = False,
        **kwargs: Any,
    ):
        self._enable_enhancement = enable_enhancement
        super().__init__(
            api_key=api_key,
            base_url=base_url,
            default_model=default_model,
            **kwargs,
        )

    async def chat(self, request: LLMRequest) -> LLMResponse:
        """发送 chat 请求，处理混元模型映射和增强模式"""
        request.model = YUANBAO_MODEL_MAP.get(request.model, request.model)

        # Add enhancement parameter if enabled
        if self._enable_enhancement:
            if not hasattr(request, "_extra") or request._extra is None:
                request._extra = {}
            request._extra["enable_enhancement"] = True

        return await super().chat(request)
