'''
Function:
    Implementation of BaiduQianfanEndpoints
Author:
    Zhenchao Jin
WeChat Official Account (微信公众号):
    Charles的皮卡丘
'''
import random
from openai import OpenAI
from .base import BaseEndpoint
from ..utils import ChatRequest, ChatResponse, ModelIOType, SecretUtils


'''QIANFAN_FREE_API_V2_KEYS'''
QIANFAN_FREE_API_V2_KEYS = {
    'deepseek-v3': [
        'YmNlLXYzL0FMVEFLLTV0Zlg0MUhNUjV3UmVKakdKTDFwUC9lYzZkNjcwMWUwYTliODRkZDQ5NTFjZTk3N2YyYTliYzVjNjI0ZDUz',
        'YmNlLXYzL0FMVEFLLUlsQUdXcnBQSUZBTUozZzhrYkQ0SS9mMTdjMGE5MDliODkxYzg5YjBkY2U1M2Q5MTM0NDhkODZhODdiYWQ5',
        'YmNlLXYzL0FMVEFLLWpORnJNb3drSnQ5Smc0d2RLUEZwRC9iOTliZGE0ZDE3ZTJkMDJmZjg4ZTQxMDlhZGNkY2E3NDlmNjhkMGVl',
    ],
    'ernie-4.0-turbo-8k': [
        'YmNlLXYzL0FMVEFLLTV0Zlg0MUhNUjV3UmVKakdKTDFwUC9lYzZkNjcwMWUwYTliODRkZDQ5NTFjZTk3N2YyYTliYzVjNjI0ZDUz',
        'YmNlLXYzL0FMVEFLLUlsQUdXcnBQSUZBTUozZzhrYkQ0SS9mMTdjMGE5MDliODkxYzg5YjBkY2U1M2Q5MTM0NDhkODZhODdiYWQ5',
        'YmNlLXYzL0FMVEFLLWpORnJNb3drSnQ5Smc0d2RLUEZwRC9iOTliZGE0ZDE3ZTJkMDJmZjg4ZTQxMDlhZGNkY2E3NDlmNjhkMGVl',
    ],
}


'''BaiduQianfanEndpoints'''
class BaiduQianfanEndpoints(BaseEndpoint):
    provider, model = "baidu", "qianfan"
    def __init__(self, **kwargs):
        super(BaiduQianfanEndpoints, self).__init__(**kwargs)
        self.registervariant("deepseek-v3", io_supported=[ModelIOType.fromtag("T2T")])
        self.registerapi(
            version="deepseek-v3", name="officialapiv2", io=ModelIOType.fromtag("T2T"), handler="officialapiv2",
            priority=100, note='official api: https://qianfan.baidubce.com/v2'
        )
        self.registervariant("ernie-4.0-turbo-8k", io_supported=[ModelIOType.fromtag("T2T")])
        self.registerapi(
            version="ernie-4.0-turbo-8k", name="officialapiv2", io=ModelIOType.fromtag("T2T"), handler="officialapiv2",
            priority=100, note='official api: https://qianfan.baidubce.com/v2'
        )
    '''officialapiv2'''
    def officialapiv2(self, req: ChatRequest, request_overrides: dict = None, version: str = None) -> ChatResponse:
        # init
        request_overrides = request_overrides or {}
        openai_cfg: dict = req.openai
        # create client
        client = OpenAI(
            base_url='https://qianfan.baidubce.com/v2', api_key=SecretUtils.b64decode(random.choice(QIANFAN_FREE_API_V2_KEYS[version])), **openai_cfg['client']
        )
        # construct message
        messages_override = openai_cfg['client.chat.completions.create'].pop('messages', None)
        messages = messages_override or [{"role": "user", "content": req.text}]
        # specify model
        model_override = openai_cfg['client.chat.completions.create'].pop('model', None)
        model = model_override or version
        # get response
        resp = client.chat.completions.create(model=model, messages=messages, **openai_cfg['client.chat.completions.create'])
        text = self._extracttextfromchatcompletionsobj(resp)
        # return
        return ChatResponse(text=text, raw=resp.to_json() or resp)