from dataclasses import dataclass, field
from typing import List, Optional, Dict, Any


# ------------------- 嵌套子实体类（同样支持空初始化） -------------------
@dataclass
class ChoiceMessage:
    """choices 中 message 字段的实体类"""
    role: Optional[str] = None  # 设为可选，默认None
    content: Optional[str] = None  # 设为可选，默认None
    refusal: Optional[Any] = None
    annotations: Optional[Any] = None
    audio: Optional[Any] = None
    function_call: Optional[Any] = None
    tool_calls: List[Any] = field(default_factory=list)  # 列表默认空列表
    reasoning: Optional[Any] = None
    reasoning_content: Optional[Any] = None


@dataclass
class Choice:
    """choices 数组元素的实体类"""
    index: Optional[int] = None  # 设为可选，默认None
    message: ChoiceMessage = field(default_factory=ChoiceMessage)  # 默认空实例
    logprobs: Optional[Any] = None
    finish_reason: Optional[str] = None
    stop_reason: Optional[Any] = None
    token_ids: Optional[Any] = None


@dataclass
class Usage:
    """usage 字段的实体类"""
    prompt_tokens: Optional[int] = None  # 设为可选，默认None
    total_tokens: Optional[int] = None
    completion_tokens: Optional[int] = None
    prompt_tokens_details: Optional[Any] = None


# ------------------- 主实体类 LLMResponse（核心调整） -------------------
@dataclass
class LLMResponse:
    """AI 聊天补全响应的主实体类（支持空初始化）"""
    id: Optional[str] = None
    object: Optional[str] = None
    created: Optional[int] = None
    model: Optional[str] = None
    choices: List[Choice] = field(default_factory=list)  # 默认空列表
    usage: Usage = field(default_factory=Usage)  # 默认空Usage实例
    service_tier: Optional[Any] = None
    system_fingerprint: Optional[Any] = None
    prompt_logprobs: Optional[Any] = None
    prompt_token_ids: Optional[Any] = None
    kv_transfer_params: Optional[Any] = None

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'LLMResponse':
        """保留原有的从字典批量创建实例的方法"""
        usage_data = data.get('usage', {})
        usage = Usage(
            prompt_tokens=usage_data.get('prompt_tokens'),
            total_tokens=usage_data.get('total_tokens'),
            completion_tokens=usage_data.get('completion_tokens'),
            prompt_tokens_details=usage_data.get('prompt_tokens_details')
        )

        choices = []
        for choice_data in data.get('choices', []):
            message_data = choice_data.get('message', {})
            message = ChoiceMessage(
                role=message_data.get('role'),
                content=message_data.get('content'),
                refusal=message_data.get('refusal'),
                annotations=message_data.get('annotations'),
                audio=message_data.get('audio'),
                function_call=message_data.get('function_call'),
                tool_calls=message_data.get('tool_calls', []),
                reasoning=message_data.get('reasoning'),
                reasoning_content=message_data.get('reasoning_content')
            )

            choice = Choice(
                index=choice_data.get('index'),
                message=message,
                logprobs=choice_data.get('logprobs'),
                finish_reason=choice_data.get('finish_reason'),
                stop_reason=choice_data.get('stop_reason'),
                token_ids=choice_data.get('token_ids')
            )
            choices.append(choice)

        return cls(
            id=data.get('id'),
            object=data.get('object'),
            created=data.get('created'),
            model=data.get('model'),
            choices=choices,
            usage=usage,
            service_tier=data.get('service_tier'),
            system_fingerprint=data.get('system_fingerprint'),
            prompt_logprobs=data.get('prompt_logprobs'),
            prompt_token_ids=data.get('prompt_token_ids'),
            kv_transfer_params=data.get('kv_transfer_params')
        )


# ------------------- 你的使用方式示例（空初始化 + 逐个赋值） -------------------
if __name__ == "__main__":
    # 模拟 API 返回的 result 字典
    result = {
        'id': 'chatcmpl-aac04a6d4249d601',
        'object': 'chat.completion',
        'created': 1770704937,
        'model': '/data/models/Qwen2-7B-Instruct-AWQ/Qwen2-7B-Instruct-AWQ'
    }

    # 1. 空初始化（完全不传参）
    resp = LLMResponse()

    # 2. 逐个赋值（和你代码里的写法一致）
    resp.id = result.get("id")
    resp.object = result.get("object")
    resp.created = result.get("created")
    resp.model = result.get("model")

    # 3. 也可以给嵌套字段赋值
    choice = Choice()
    choice.index = 0
    choice.message.role = "assistant"
    choice.message.content = "机器学习是一种人工智能技术..."
    resp.choices.append(choice)

    # 4. 赋值后访问属性
    print("ID：", resp.id)  # 输出：chatcmpl-aac04a6d4249d601
    print("回复内容：", resp.choices[0].message.content)  # 输出：机器学习是一种人工智能技术...