import pydantic
from typing import Optional, Dict, Any
from . import base_original_provider_usage


class OpenAICompletionsUsage(base_original_provider_usage.BaseOriginalProviderUsage):
    """OpenAI calls token usage data (or token usage data in OpenAI format). Updated 11.03.2025"""

    completion_tokens: int
    """Number of tokens in the generated completion."""

    prompt_tokens: int
    """Number of tokens in the prompt."""

    total_tokens: int
    """Total number of tokens used in the request (prompt + completion)."""

    completion_tokens_details: Optional["CompletionTokensDetails"] = None
    """Breakdown of tokens used in a completion."""

    prompt_tokens_details: Optional["PromptTokensDetails"] = None
    """Breakdown of tokens used in the prompt."""

    def to_backend_compatible_flat_dict(self, parent_key_prefix: str) -> Dict[str, int]:
        result = {**self.__dict__}

        if self.completion_tokens_details is not None:
            result["completion_tokens_details"] = (
                self.completion_tokens_details.model_dump()
            )

        if self.prompt_tokens_details is not None:
            result["prompt_tokens_details"] = self.prompt_tokens_details.model_dump()

        return self.flatten_result_and_add_model_extra(
            result=result, parent_key_prefix=parent_key_prefix
        )

    @classmethod
    def from_original_usage_dict(
        cls, usage_dict: Dict[str, Any]
    ) -> "OpenAICompletionsUsage":
        usage_dict = {**usage_dict}
        completion_tokens_details_raw = usage_dict.pop(
            "completion_tokens_details", None
        )
        prompt_tokens_details_raw = usage_dict.pop("prompt_tokens_details", None)

        completion_tokens_details = (
            CompletionTokensDetails(**completion_tokens_details_raw)
            if isinstance(completion_tokens_details_raw, dict)
            else None
        )

        prompt_tokens_details = (
            PromptTokensDetails(**prompt_tokens_details_raw)
            if isinstance(prompt_tokens_details_raw, dict)
            else None
        )

        return cls(
            **usage_dict,
            completion_tokens_details=completion_tokens_details,
            prompt_tokens_details=prompt_tokens_details,
        )


class CompletionTokensDetails(pydantic.BaseModel):
    model_config = pydantic.ConfigDict(extra="allow")

    accepted_prediction_tokens: Optional[int] = None
    """
    When using Predicted Outputs, the number of tokens in the prediction that
    appeared in the completion.
    """

    audio_tokens: Optional[int] = None
    """Audio input tokens generated by the model."""

    reasoning_tokens: Optional[int] = None
    """Tokens generated by the model for reasoning."""

    rejected_prediction_tokens: Optional[int] = None
    """
    When using Predicted Outputs, the number of tokens in the prediction that did
    not appear in the completion. However, like reasoning tokens, these tokens are
    still counted in the total completion tokens for purposes of billing, output,
    and context window limits.
    """


class PromptTokensDetails(pydantic.BaseModel):
    model_config = pydantic.ConfigDict(extra="allow")

    audio_tokens: Optional[int] = None
    """Audio input tokens present in the prompt."""

    cached_tokens: Optional[int] = None
    """Cached tokens present in the prompt."""