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

from __future__ import annotations

from dataclasses import dataclass
from typing import Literal


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

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


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

IMAGE_MODELS: list[ModelInfo] = [
    ModelInfo(
        name="flux-dev",
        provider="BlackForestLabs",
        type="image",
        capabilities=["text-to-image"],
    ),
    ModelInfo(
        name="flux-pro",
        provider="BlackForestLabs",
        type="image",
        capabilities=["text-to-image"],
    ),
    ModelInfo(
        name="flux-pro-1.1",
        provider="BlackForestLabs",
        type="image",
        capabilities=["text-to-image"],
    ),
    ModelInfo(
        name="flux-pro-1.1-ultra",
        provider="BlackForestLabs",
        type="image",
        capabilities=["text-to-image", "high resolution"],
    ),
    ModelInfo(
        name="flux-2-pro",
        provider="BlackForestLabs",
        type="image",
        capabilities=["text-to-image", "reference images"],
    ),
    ModelInfo(
        name="flux-kontext-pro",
        provider="BlackForestLabs",
        type="image",
        capabilities=["text-to-image", "reference images"],
    ),
    ModelInfo(
        name="flux-pro-1.0-canny",
        provider="BlackForestLabs",
        type="image",
        capabilities=["text-to-image", "control images", "edge detection"],
    ),
    ModelInfo(
        name="flux-pro-1.0-depth",
        provider="BlackForestLabs",
        type="image",
        capabilities=["text-to-image", "control images", "depth map"],
    ),
    ModelInfo(
        name="flux-pro-1.0-fill",
        provider="BlackForestLabs",
        type="image",
        capabilities=["text-to-image", "inpainting"],
    ),
    ModelInfo(
        name="flux-pro-1.0-expand",
        provider="BlackForestLabs",
        type="image",
        capabilities=["text-to-image", "outpainting"],
    ),
]

ALL_MODELS: list[ModelInfo] = TEXT_MODELS + IMAGE_MODELS