"""Model registry and lazy fetch from the Hugging Face model repo.

The Space and the FraXtex QGIS plugin consume the same published artifacts, so what
you see here is what QGIS users get. Weights are not committed to this Space: they
are downloaded on first use and cached by `huggingface_hub`.

`channels` below drives the UI only — whether the DEM upload is shown before any
model has been opened. At inference time the channel count is read from the ONNX
graph itself (`inference.load_session`), which stays authoritative.
"""

from huggingface_hub import hf_hub_download

REPO_ID = "ayoubft/fraXteX"

MODELS = {
    "unet-rgb": {
        "filename": "rgb/unet-rgb.onnx",
        "label": "U-Net (RGB)",
        "channels": 3,
        "note": (
            "U-Net trained from scratch, 31 M parameters. Normalizes inside the graph "
            "with dataset statistics, and accepts any patch size that is a multiple of 16."
        ),
    },
    "sam2-rgb": {
        "filename": "rgb/sam2-rgb.onnx",
        "label": "SAM 2 (RGB) — slow",
        "channels": 3,
        "note": (
            "SAM 2 (hiera-tiny) with a frozen encoder, prompt-free. Runs at 512x512 "
            "patches, so it is noticeably slower than the U-Net on this CPU Space."
        ),
    },
    "unet-rgbdem": {
        "filename": "rgbdem/unet-rgbdem.onnx",
        "label": "U-Net (RGB+DEM)",
        "channels": 4,
        "note": (
            "Same U-Net architecture, taking a DEM as a 4th channel alongside the RGB "
            "bands. Fixed 256x256 patches, inputs scaled to 0-1."
        ),
    },
    "segformer-rgbdem": {
        "filename": "rgbdem/segformer-rgbdem.onnx",
        "label": "SegFormer (RGB+DEM)",
        "channels": 4,
        "note": (
            "SegFormer with a transformer encoder, taking a DEM as a 4th channel "
            "alongside the RGB bands. Fixed 256x256 patches, inputs scaled to 0-1."
        ),
    },
}

# Dropdown order: the RGB models first, since they need no DEM upload.
CHOICES = [(spec["label"], key) for key, spec in MODELS.items()]
DEFAULT = "unet-rgb"

_cache = {}


def needs_dem(key):
    return MODELS[key]["channels"] == 4


def note_for(key):
    return MODELS[key]["note"]


def get_model_path(key, feedback=None):
    """Download the model on first use and return its local path."""
    if key in _cache:
        return _cache[key]
    spec = MODELS[key]
    if feedback is not None:
        feedback.pushInfo(f"Fetching {spec['filename']} from {REPO_ID} (first use only)...")
    path = hf_hub_download(repo_id=REPO_ID, filename=spec["filename"])
    _cache[key] = path
    return path