"""Model-related types mirroring the Skytells API response shapes."""

from __future__ import annotations

from enum import Enum
from typing import Any, Dict, List, Optional


class ModelPrivacy(str, Enum):
    """Privacy level of a model."""

    PUBLIC = "public"
    PRIVATE = "private"


class ModelType(str, Enum):
    """The type of AI model."""

    IMAGE = "image"
    VIDEO = "video"
    AUDIO = "audio"
    MUSIC = "music"
    TEXT = "text"
    CODE = "code"
    MULTIMODAL = "multimodal"


class PricingUnit(str, Enum):
    """Billing unit for a model's pricing."""

    IMAGE = "image"
    VIDEO = "video"
    SECOND = "second"
    PREDICTION = "prediction"
    GPU = "gpu"
    IMAGE_MEGAPIXEL = "image_megapixel"
    COMPUTING_SECOND = "computing_second"
    AUDIO_SECOND = "audio_second"
    VIDEO_SECOND = "video_second"
    TOKEN = "token"
    FIVE_SECONDS = "5 seconds"
    MINUTE = "minute"


class PricingOperator(str, Enum):
    """Comparison operator used in pricing criteria."""

    EQUALS = "equals"
    DOUBLE_EQUALS = "=="


class ModelFieldsOptions:
    """
    Options for requesting additional model fields.

    Args:
        fields: List of extra fields to include in the response.
                Valid values: ``"input_schema"``, ``"output_schema"``.

    Example::

        options = ModelFieldsOptions(fields=["input_schema", "output_schema"])
    """

    def __init__(self, fields: Optional[List[str]] = None) -> None:
        self.fields: List[str] = fields or []

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "ModelFieldsOptions":
        return cls(fields=d.get("fields", []))
