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


'''SPARK_FREE_APIV1_KEYS'''
SPARK_FREE_APIV1_KEYS = {
    'lite': ['aXB6b3RsR2V2TnFRc2FmdldTWGk6Y29vRXhpTlJrSHRRdEhra0lxTms=']
}
    

'''IFLYTEKSparkEndpoints'''
class IFLYTEKSparkEndpoints(BaseEndpoint):
    provider, model = "iflytek", "spark"
    def __init__(self, **kwargs):
        super(IFLYTEKSparkEndpoints, self).__init__(**kwargs)
        self.registervariant("thirdparty", io_supported=[ModelIOType.fromtag("T2T")])
        self.registerapi(
            version="thirdparty", name="cggapi", io=ModelIOType.fromtag("T2T"), handler="cggapi",
            priority=10, note='thirdparty api: https://api.cenguigui.cn/api/chat/?msg=',
        )
        self.registervariant("lite", io_supported=[ModelIOType.fromtag("T2T")])
        self.registerapi(
            version="lite", name="officialapiv1", io=ModelIOType.fromtag("T2T"), handler="officialapiv1",
            priority=100, note='official api: https://spark-api-open.xf-yun.com/v1'
        )
    '''officialapiv1'''
    def officialapiv1(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://spark-api-open.xf-yun.com/v1', api_key=SecretUtils.b64decode(random.choice(SPARK_FREE_APIV1_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)
    '''cggapi'''
    def cggapi(self, req: ChatRequest, request_overrides: dict = None, version: str = None) -> ChatResponse:
        request_overrides = request_overrides or {}
        try:
            resp = self.session.get("https://api.cenguigui.cn/api/chat/", params={"msg": req.text}, **request_overrides)
            resp.raise_for_status()
        except:
            resp = self.session.get("https://api-v1.cenguigui.cn/api/chat/", params={"msg": req.text}, **request_overrides)
            resp.raise_for_status()
        api_resp_raw_data = resp2json(resp)
        return ChatResponse(text=api_resp_raw_data['data']['content'], raw=api_resp_raw_data)