from dashscope import Generation
import random
from http import HTTPStatus
from enum import Enum
from .qgi import Generator as BaseGenerator

class Model(str, Enum):
    # 通义千问超大规模语言模型，支持中文、英文等不同语言输入。
    # 模型支持 8,000 tokens上下文，为了保证正常的使用和输出，API限定用户输入为 6,000 tokens。
    QWEN_TURBO = "qwen-turbo"

    # 通义千问超大规模语言模型增强版，支持中文、英文等不同语言输入。
    # 模型支持 32,000 tokens上下文，为了保证正常的使用和输出，API限定用户输入为30,000 tokens。
    QWEN_PLUS = "qwen-plus"

    # 通义千问千亿级别超大规模语言模型，支持中文、英文等不同语言输入。
    # 随着模型的升级，qwen-max将滚动更新升级，如果希望使用固定版本，请使用下面的历史快照版本。
    # 当前qwen-max模型与qwen-max-0428快照版本等价，均为最新版本的qwen-max模型，也是当前通义千问2.5产品版本背后的API模型。
    # 模型支持 8,000 tokens上下文，为了保证正常的使用和输出，API限定用户输入为 6,000 tokens。
    QWEN_MAX = "qwen-max"

    # 通义千问千亿级别超大规模语言模型，支持中文、英文等不同语言输入。
    # 模型支持 30,000 tokens上下文，为了保证正常的使用和输出，API限定用户输入为 28,000 tokens。     
    QWEN_MAX_LONGCONTEXT = "qwen-max-longcontext"


class Generator(BaseGenerator):
    def __init__(self, model:str, messages:list[dict[str,str]]|None, apikey:str, seed:int=0) -> None:
        '''
        参数:
        - model: str - 使用的大模型名称，例如 qwen-turbo。
        - messages: list[dict[str,str]] - 提供给大模型的 messages 参数
        - apikey: str - 用于调用模型的 dashscope API密钥。
        - seed: int - 随机种子，用于确保生成结果的可复现性，如果忽略或为0，则生成一个随机数。
        '''
        self.model = model
        self.messages = messages
        self.apikey = apikey
        self.seed = seed
        
    def _generate(self, prompt:str|None) -> tuple[HTTPStatus, str]:
        # 设置随机数种子seed，如果没有设置，则随机数种子默认为1234
        seed = self.seed
        if seed == 0:
            seed = random.randint(1, 10000)
            
        ret = Generation.call(model=self.model,
                            messages=self.messages,
                            prompt=prompt,
                            api_key=self.apikey,
                            seed=seed,
                            result_format='text')
        
        if ret.status_code != HTTPStatus.OK:
            return ret.status_code, ret.message
        else:
            return ret.status_code, ret.output.text