"""Parsed chat completion types for structured output."""

from typing import Any, Generic, Optional, TypeVar
from pydantic import BaseModel, Field

ContentType = TypeVar("ContentType")


class SDKRouterMeta(BaseModel):
    """Metadata from SDKRouter when using model aliases.

    This is returned when using alias models like @smart+agents.
    Contains information about model resolution and auto-upgrade.
    """

    original_alias: Optional[str] = Field(
        default=None,
        description="The alias that was requested, e.g. '@smart+agents'"
    )
    resolved_model: Optional[str] = Field(
        default=None,
        description="The actual model selected, e.g. 'google/gemini-3.1-pro-preview'"
    )
    estimated_tokens: Optional[int] = Field(
        default=None,
        description="Estimated tokens for the request (used for context length filtering)"
    )
    upgraded: bool = Field(
        default=False,
        description="True if model was upgraded due to context length requirements"
    )
    upgrade_reason: Optional[str] = Field(
        default=None,
        description="Reason for upgrade, e.g. 'context_length'"
    )


class ParsedMessage(BaseModel, Generic[ContentType]):
    """Message with parsed content from structured output."""

    role: str
    content: Optional[str] = None
    parsed: Optional[ContentType] = None
    refusal: Optional[str] = None
    tool_calls: Optional[list[Any]] = None
    function_call: Optional[Any] = None

    model_config = {"arbitrary_types_allowed": True}

    @classmethod
    def from_message(cls, message: Any, parsed_content: Optional[ContentType] = None) -> "ParsedMessage[ContentType]":
        """Create ParsedMessage from raw message object."""
        return cls(
            role=message.role,
            content=message.content if hasattr(message, "content") else None,
            parsed=parsed_content,
            refusal=message.refusal if hasattr(message, "refusal") else None,
            tool_calls=message.tool_calls if hasattr(message, "tool_calls") else None,
            function_call=message.function_call if hasattr(message, "function_call") else None,
        )


class ParsedChoice(BaseModel, Generic[ContentType]):
    """Choice with parsed message."""

    index: int
    message: ParsedMessage[ContentType]
    finish_reason: Optional[str] = None
    logprobs: Optional[Any] = None

    model_config = {"arbitrary_types_allowed": True}

    @classmethod
    def from_choice(cls, choice: Any, parsed_content: Optional[ContentType] = None) -> "ParsedChoice[ContentType]":
        """Create ParsedChoice from raw choice object."""
        return cls(
            index=choice.index,
            message=ParsedMessage.from_message(choice.message, parsed_content=parsed_content),
            finish_reason=choice.finish_reason if hasattr(choice, "finish_reason") else None,
            logprobs=choice.logprobs if hasattr(choice, "logprobs") else None,
        )


class ParsedChatCompletion(BaseModel, Generic[ContentType]):
    """ChatCompletion with parsed content from structured output."""

    id: str
    choices: list[ParsedChoice[ContentType]]
    created: int
    model: str
    object: str = "chat.completion"
    system_fingerprint: Optional[str] = None
    usage: Optional[Any] = None
    service_tier: Optional[str] = None
    sdkrouter_meta: Optional[SDKRouterMeta] = Field(
        default=None,
        alias="_meta",
        description="SDKRouter metadata when using model aliases"
    )

    model_config = {"arbitrary_types_allowed": True, "populate_by_name": True}

    @classmethod
    def from_completion(
        cls,
        completion: Any,
        parsed_choices: list[ParsedChoice[ContentType]],
    ) -> "ParsedChatCompletion[ContentType]":
        """Create ParsedChatCompletion from raw completion object."""
        return cls(
            id=completion.id,
            choices=parsed_choices,
            created=completion.created,
            model=completion.model,
            object=completion.object if hasattr(completion, "object") else "chat.completion",
            system_fingerprint=completion.system_fingerprint if hasattr(completion, "system_fingerprint") else None,
            usage=completion.usage if hasattr(completion, "usage") else None,
            service_tier=completion.service_tier if hasattr(completion, "service_tier") else None,
        )


__all__ = [
    "ContentType",
    "ParsedMessage",
    "ParsedChoice",
    "ParsedChatCompletion",
    "SDKRouterMeta",
]
