# chat_completion/models/message.py
from dataclasses import dataclass
from typing import List, Optional, Literal
from .base import BaseModel
from .tools import ToolCall
from .tokens import LogProbs
from .audio import AudioResponse


@dataclass
class Message(BaseModel):
    """Represents a chat completion message generated by the model."""

    role: Literal["assistant", "user", "system"]
    content: Optional[str]
    refusal: Optional[str] = None
    tool_calls: Optional[List[ToolCall]] = None

    @classmethod
    def from_dict(cls, data: dict) -> "Message":
        return cls(
            role=data["role"],
            content=data.get("content"),
            refusal=data.get("refusal"),
            tool_calls=(
                [ToolCall.from_dict(tc) for tc in data["tool_calls"]]
                if data.get("tool_calls")
                else None
            ),
        )


@dataclass
class Choice(BaseModel):
    """Represents a single chat completion choice."""

    index: int
    message: Message
    finish_reason: Literal["stop", "length", "content_filter", "tool_calls"]
    logprobs: Optional[LogProbs] = None
    audio: Optional[AudioResponse] = None

    @classmethod
    def from_dict(cls, data: dict) -> "Choice":
        return cls(
            index=data["index"],
            message=Message.from_dict(data["message"]),
            finish_reason=data["finish_reason"],
            logprobs=(
                LogProbs.from_dict(data["logprobs"]) if data.get("logprobs") else None
            ),
            audio=AudioResponse.from_dict(data["audio"]) if data.get("audio") else None,
        )
