"""Registry of supported models and their capabilities."""

from __future__ import annotations

from dataclasses import dataclass
from enum import StrEnum
from typing import Literal


class Capability(StrEnum):
    """Capabilities a model may support."""

    TEXT_GENERATION = "text generation"
    VISION = "vision"
    TEXT_TO_IMAGE = "text-to-image"
    REFERENCE_IMAGES = "reference images"
    CONTROL_IMAGES = "control images"


@dataclass(frozen=True)
class ModelInfo:
    """Describes a supported model and its capabilities."""

    name: str
    provider: str
    type: Literal["text", "image"]
    capabilities: list[Capability]


TEXT_MODELS: list[ModelInfo] = [
    ModelInfo(
        name="mistral-large-latest",
        provider="Mistral",
        type="text",
        capabilities=[Capability.TEXT_GENERATION],
    ),
    ModelInfo(
        name="mistral-small-latest",
        provider="Mistral",
        type="text",
        capabilities=[Capability.TEXT_GENERATION],
    ),
    ModelInfo(
        name="pixtral-large-latest",
        provider="Mistral",
        type="text",
        capabilities=[Capability.TEXT_GENERATION, Capability.VISION],
    ),
    ModelInfo(
        name="pixtral-12b-latest",
        provider="Mistral",
        type="text",
        capabilities=[Capability.TEXT_GENERATION, Capability.VISION],
    ),
]

IMAGE_MODELS: list[ModelInfo] = [
    ModelInfo(
        name="flux-dev",
        provider="BlackForestLabs",
        type="image",
        capabilities=[Capability.TEXT_TO_IMAGE],
    ),
    ModelInfo(
        name="flux-pro",
        provider="BlackForestLabs",
        type="image",
        capabilities=[Capability.TEXT_TO_IMAGE],
    ),
    ModelInfo(
        name="flux-pro-1.1",
        provider="BlackForestLabs",
        type="image",
        capabilities=[Capability.TEXT_TO_IMAGE],
    ),
    ModelInfo(
        name="flux-pro-1.1-ultra",
        provider="BlackForestLabs",
        type="image",
        capabilities=[Capability.TEXT_TO_IMAGE],
    ),
    ModelInfo(
        name="flux-2-pro",
        provider="BlackForestLabs",
        type="image",
        capabilities=[Capability.TEXT_TO_IMAGE, Capability.REFERENCE_IMAGES],
    ),
    ModelInfo(
        name="flux-kontext-pro",
        provider="BlackForestLabs",
        type="image",
        capabilities=[Capability.TEXT_TO_IMAGE, Capability.REFERENCE_IMAGES],
    ),
    ModelInfo(
        name="flux-pro-1.0-canny",
        provider="BlackForestLabs",
        type="image",
        capabilities=[
            Capability.TEXT_TO_IMAGE,
            Capability.CONTROL_IMAGES,
        ],
    ),
    ModelInfo(
        name="flux-pro-1.0-depth",
        provider="BlackForestLabs",
        type="image",
        capabilities=[
            Capability.TEXT_TO_IMAGE,
            Capability.CONTROL_IMAGES,
        ],
    ),
    ModelInfo(
        name="flux-pro-1.0-fill",
        provider="BlackForestLabs",
        type="image",
        capabilities=[
            Capability.TEXT_TO_IMAGE,
        ],
    ),
    ModelInfo(
        name="flux-pro-1.0-expand",
        provider="BlackForestLabs",
        type="image",
        capabilities=[
            Capability.TEXT_TO_IMAGE,
        ],
    ),
]

ALL_MODELS: list[ModelInfo] = TEXT_MODELS + IMAGE_MODELS