from typing import Optional

import aiohttp

from huibiao_framework.client.llm.huize_qwen32b_awq_client import HuiZeQwen32bQwqClient
from huibiao_framework.client.llm.abstrast_llm_client import AbstractLlmModelClient
from huibiao_framework.client.llm.mindie_ds_32b_awq import MindleDs32bAwqClient
from huibiao_framework.config.config import ModelConfig, LlmModelNameConstant
from huibiao_framework.execption.model import LLMException
from loguru import logger


class LlmModelClient:
    __MODEL_REGISTRY = {
        LlmModelNameConstant.HuiZeQwen32bAwq: HuiZeQwen32bQwqClient,
        LlmModelNameConstant.MindileDs32bAwq: MindleDs32bAwqClient,
    }

    def __init__(self, session: aiohttp.ClientSession, model: str = None):
        self.model_name: str = model or ModelConfig.LLM_MODEL_TYPE
        self.__model: Optional[AbstractLlmModelClient] = None

        if self.model_name in self.__MODEL_REGISTRY:
            self.__model = self.__MODEL_REGISTRY[self.model_name](session)
        else:
            logger.error(f"未注册的模型类型: {self.model_name}")
            raise LLMException("错误的llm模型类型")

    async def simple_query(self, content, session_id: str = "", **kwargs) -> str:
        return await self.__model.simple_query(content, session_id, **kwargs)

    @classmethod
    def register_model(
        cls, model_name: str, model_client: type(AbstractLlmModelClient)
    ):
        """
        注册一个自定义的模型，实现AbstractLlmModelClient类
        """
        logger.info(f"注册模型: {model_name}")
        if model_name in cls.__MODEL_REGISTRY:
            logger.warning(f"已存在模型: {model_name},请勿重新注册")
        else:
            cls.__MODEL_REGISTRY[model_name] = model_client

    @staticmethod
    def register_model_decorator(model_name):
        """
        注册LLM模型的装饰器
        @LlmModelClient.register_model_decorator({模型名})
        """

        def decorator(cls):
            # 注册模型到客户端
            LlmModelClient.register_model(model_name, cls)
            return cls

        return decorator
