from dataclasses import dataclass


@dataclass(frozen=True)
class ModelSpec:
    key: str
    label: str
    directory_name: str
    repository: str
    license_name: str
    description: str
    supports_gs: bool
    download_patterns: tuple[str, ...] = ()
    ready_patterns: tuple[str, ...] = ("config.json", "*.safetensors")
    enabled: bool = True
    r2_directory: str = ""
    r2_files: tuple[tuple[str, int, str], ...] = ()


MODELS = {
    "SMALL": ModelSpec(
        key="SMALL",
        label="DA3 Small",
        directory_name="da3_small",
        repository="depth-anything/DA3-SMALL",
        license_name="Apache 2.0",
        description="Fastest any-view model with Apache 2.0 weights",
        supports_gs=False,
        r2_directory="DA3Small",
        r2_files=(
            ("config.json", 1202, "a486e29e82b7ab4a7d4cefc1ea4526cfe2ae438a572c8ca98917cfbcde7447d2"),
            ("model.safetensors", 137248940, "364492e38a3a06d221ac75da7f6621ada3f2361cd24fde11ba79091e9f40efcf"),
        ),
    ),
    "BASE": ModelSpec(
        key="BASE",
        label="DA3 Base",
        directory_name="da3_base",
        repository="depth-anything/DA3-BASE",
        license_name="Apache 2.0",
        description="Balanced any-view model with Apache 2.0 weights",
        supports_gs=False,
        r2_directory="DA3Base",
        r2_files=(
            ("config.json", 1205, "5e34115ebc17bd2d8d43033c5f72e9446ac8833fd61d3fa160b7e67e0bb5b7b5"),
            ("configuration.json", 74, "b684d553257b2ca379fd5dfdf96a6e4a633eec3b969c97119c9875465a7c07a8"),
            ("model.safetensors", 541518028, "e01067dc1659613083d9145a9a2547ccdbe6ccbbf83c4fe7b3e8a4e2bdae78b5"),
        ),
    ),
    "LARGE_1_1": ModelSpec(
        key="LARGE_1_1",
        label="DA3 Large 1.1",
        directory_name="da3_large_1_1",
        repository="depth-anything/DA3-LARGE-1.1",
        license_name="CC BY-NC 4.0",
        description="Higher quality non-commercial model with larger VRAM use",
        supports_gs=False,
        enabled=False,
    ),
    "GIANT_1_1": ModelSpec(
        key="GIANT_1_1",
        label="DA3 Giant 1.1",
        directory_name="da3_giant_1_1",
        repository="depth-anything/DA3-GIANT-1.1",
        license_name="CC BY-NC 4.0",
        description="Non-commercial Giant model with Gaussian head",
        supports_gs=True,
        enabled=False,
    ),
    "NESTED_GIANT_LARGE_1_1": ModelSpec(
        key="NESTED_GIANT_LARGE_1_1",
        label="DA3 Nested Giant Large 1.1",
        directory_name="da3_nested_giant_large_1_1",
        repository="depth-anything/DA3NESTED-GIANT-LARGE-1.1",
        license_name="CC BY-NC 4.0",
        description="Non-commercial nested model with metric depth and Gaussian head",
        supports_gs=True,
        enabled=False,
    ),
}

BEN2_MODEL = ModelSpec(
    key="BEN2_BASE",
    label="BEN2 Base",
    directory_name="ben2_base",
    repository="PramaLLC/BEN2",
    license_name="MIT",
    description="Fast high-quality background removal model",
    supports_gs=False,
    download_patterns=("config.json", "model.safetensors"),
    r2_directory="BEN2Base",
    r2_files=(
        ("config.json", 124, "01ac876925577b52fa7c6af3f5b1d1420bc9b40016b6c25e2de89b38d8fff865"),
        ("model.safetensors", 380577976, "ea8b7907176a09667c86343dc7d00de6a6d871076cb90bb5f753618fd6fb3ebb"),
    ),
)
def moge2_model(key, label, directory_name, repository, description, r2_directory, size, checksum):
    return ModelSpec(
        key=key,
        label=label,
        directory_name=directory_name,
        repository=repository,
        license_name="MIT",
        description=description,
        supports_gs=False,
        download_patterns=("model.pt",),
        ready_patterns=("model.pt",),
        r2_directory=r2_directory,
        r2_files=(("model.pt", size, checksum),),
    )


MOGE2_MODELS = {
    "MOGE2_VITS_NORMAL": moge2_model(
        "MOGE2_VITS_NORMAL", "MoGe-2 ViT-S Normal", "moge2_vits_normal",
        "Ruicheng/moge-2-vits-normal", "Fastest MoGe-2 model", "MoGe2ViTSNormal",
        140550416, "79a16621928c2bf0ed04659218c55c01075e950507f40bb3332fb4c873d3e1dc",
    ),
    "MOGE2_VITB_NORMAL": moge2_model(
        "MOGE2_VITB_NORMAL", "MoGe-2 ViT-B Normal", "moge2_vitb_normal",
        "Ruicheng/moge-2-vitb-normal", "Balanced MoGe-2 model", "MoGe2ViTBNormal",
        419110160, "16b8110e86d5dc5a849db120ca96ef3a223fd30b0c9146d1d81db504073da5f6",
    ),
    "MOGE2_VITL_NORMAL": moge2_model(
        "MOGE2_VITL_NORMAL", "MoGe-2 ViT-L Normal", "moge2_vitl_normal",
        "Ruicheng/moge-2-vitl-normal", "Highest-quality MoGe-2 model", "MoGe2ViTLNormal",
        1323815904, "280741fd09bc3f403ccff9967784c2a391b52d2c0742ae3efdb21d9f90cc1a01",
    ),
}
MOGE2_MODEL = MOGE2_MODELS["MOGE2_VITB_NORMAL"]
DOWNLOADABLE_MODELS = {
    **MODELS,
    BEN2_MODEL.key: BEN2_MODEL,
    **MOGE2_MODELS,
}
MODEL_NUMBERS = {
    model_key: number for number, model_key in enumerate(MODELS)
}


_ENUM_ITEMS_CACHE = {}
_DEPTH_PLANE_ENUM_ITEMS_CACHE = {}
_SURFACE_ENUM_ITEMS_CACHE = {}


def enum_items(supports_gs=None, show_download_status=False):
    model_ready = None
    if show_download_status:
        from .client import model_ready

    models = [
        model
        for model in MODELS.values()
        if model.enabled
        and (supports_gs is None or model.supports_gs == supports_gs)
    ]
    download_status = (
        tuple(model_ready(model.key) for model in models)
        if model_ready is not None
        else ()
    )
    cache_key = (
        supports_gs,
        show_download_status,
        download_status,
    )
    cached_items = _ENUM_ITEMS_CACHE.get(cache_key)
    if cached_items is not None:
        return cached_items

    items = []
    for index, model in enumerate(models):
        label = model.label
        if download_status and not download_status[index]:
            label = f"{label} (Missing)"
        items.append(
            (
                model.key,
                label,
                f"{model.description}; license: {model.license_name}",
                MODEL_NUMBERS[model.key],
            )
        )

    _ENUM_ITEMS_CACHE[cache_key] = items
    return items


def model_enum_items(_owner, _context):
    return enum_items(show_download_status=True)


def gs_model_enum_items(_owner, _context):
    return enum_items(supports_gs=True, show_download_status=True)


def moge2_model_enum_items(_owner, _context):
    from .client import model_ready

    return [
        (
            model.key,
            model.label if model_ready(model.key) else f"{model.label} (Missing)",
            f"{model.description}; license: {model.license_name}",
            index,
        )
        for index, model in enumerate(MOGE2_MODELS.values())
    ]


def depth_plane_model_enum_items(_owner, _context):
    from .client import model_ready

    models = [
        *MOGE2_MODELS.values(),
        *(model for model in MODELS.values() if model.enabled),
    ]
    ready = tuple(model_ready(model.key) for model in models)
    cached = _DEPTH_PLANE_ENUM_ITEMS_CACHE.get(ready)
    if cached is not None:
        return cached
    items = []
    for index, model in enumerate(models):
        label = model.label if ready[index] else f"{model.label} (Missing)"
        items.append(
            (
                model.key,
                label,
                f"{model.description}; license: {model.license_name}",
                index,
            )
        )
    items.append(("ORIGINAL", "Original", "Use the original image as depth", len(items)))
    _DEPTH_PLANE_ENUM_ITEMS_CACHE[ready] = items
    return items


def surface_model_enum_items(owner, context):
    depth_items = depth_plane_model_enum_items(owner, context)
    cache_key = id(depth_items)
    cached = _SURFACE_ENUM_ITEMS_CACHE.get(cache_key)
    if cached is not None:
        return cached
    items = [
        item
        for item in depth_items
        if item[0] != "ORIGINAL"
    ]
    _SURFACE_ENUM_ITEMS_CACHE[cache_key] = items
    return items


def selected_model_key(settings, export_format=None):
    selected_format = export_format or settings.export_format
    if selected_format == "gs_ply":
        return settings.gs_model
    return settings.model


def get_model(model_key):
    try:
        return MODELS[model_key]
    except KeyError as error:
        raise ValueError(f"Unknown DA3 model: {model_key}") from error


def get_downloadable_model(model_key):
    try:
        return DOWNLOADABLE_MODELS[model_key]
    except KeyError as error:
        raise ValueError(f"Unknown model: {model_key}") from error


def model_directory_name(model_key):
    """Return the model's folder name inside the shared models directory."""
    return get_downloadable_model(model_key).directory_name