"""Model registry: which backend serves each operation, its license, and its files.

Weights are never part of the image or the repo. They download lazily on first
use into MEDIA_MODELS_DIR/<backend id>/ and are verified against the pinned
SHA256 and size before anything loads them. A mismatch deletes the file and
the operation reports `unsupported` with the reason.

Default backends are commercially usable. Anything else must be listed in
MEDIA_OPTIONAL_BACKENDS by the operator before it runs.
"""

from __future__ import annotations

import hashlib
import importlib.util
import logging
import os
import threading
import urllib.request
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any

from .errors import Unsupported

log = logging.getLogger("media-worker.models")

OPERATIONS = (
    "detectScenes",
    "transcribe",
    "detectFaces",
    "trackFaces",
    "faceLandmarks",
    "embedFrames",
    "classifyContent",
    "segmentPrompted",
    "inpaintVideo",
    "redact",
    "synthesizeIdentity",
    "audioFeatures",
    "saliency",
)

HF = "https://huggingface.co"
YUNET_REV = "3cc26e7f1014a5ee5d74a42acee58bafc9d0a310"
SFACE_REV = "3d7082438a6e4551e840c9b2bb60b71e8da4b524"
WHISPER_REV = "536b0662742c02347bc0e980a01041f333bce120"
SIGLIP_REV = "4649052661e53c7000355844105f8a1792088239"
LAMA_REV = "c3c0c9e468934d62e79c329e35d82dd09ff8c444"
SAM2_REV = "de431c4043854a71d8101e17995dfe596bf101a5"


@dataclass(frozen=True)
class ModelFile:
    name: str
    url: str
    sha256: str
    sizeBytes: int


@dataclass(frozen=True)
class BackendSpec:
    """One way to serve an operation. `id` doubles as the model cache directory."""

    id: str
    model: str | None
    license: str | None
    commercialUse: bool | None
    experimental: bool
    source: str | None
    files: tuple[ModelFile, ...] = ()
    # Python modules that must be importable.
    requires: tuple[str, ...] = ()
    # Set when the backend can never run (e.g. no licensed model exists).
    unavailable_reason: str | None = None


@dataclass
class Backend:
    spec: BackendSpec
    # Builds the runtime object from {file name: verified local path} and the settings.
    load: Callable[[dict[str, str], Any], Any] | None = None


def _yunet_file() -> ModelFile:
    return ModelFile(
        "face_detection_yunet_2023mar.onnx",
        f"{HF}/opencv/face_detection_yunet/resolve/{YUNET_REV}/face_detection_yunet_2023mar.onnx",
        "8f2383e4dd3cfbb4553ea8718107fc0423210dc964f9f4280604804ed2552fa4",
        232589,
    )


def _siglip(name: str, sha256: str, size: int) -> ModelFile:
    path = name if name == "tokenizer.json" else f"onnx/{name}"
    return ModelFile(name, f"{HF}/Xenova/siglip-base-patch16-224/resolve/{SIGLIP_REV}/{path}", sha256, size)


def _whisper(name: str, sha256: str, size: int) -> ModelFile:
    return ModelFile(name, f"{HF}/Systran/faster-whisper-small/resolve/{WHISPER_REV}/{name}", sha256, size)


def _lazy(module: str, attr: str) -> Callable[[dict[str, str], Any], Any]:
    def load(files: dict[str, str], settings: Any) -> Any:
        mod = importlib.import_module(f"app.ops.{module}")
        return getattr(mod, attr)(files, settings)

    return load


YUNET = BackendSpec(
    id="yunet-2023mar",
    model="opencv-yunet-2023mar",
    license="MIT",
    commercialUse=True,
    experimental=False,
    source="https://github.com/opencv/opencv_zoo/tree/main/models/face_detection_yunet",
    files=(_yunet_file(),),
    requires=("cv2",),
)

SIGLIP_SOURCE = "https://huggingface.co/google/siglip-base-patch16-224"
SIGLIP_VISION = _siglip(
    "vision_model_quantized.onnx",
    "ef14a954f3d57e1806666432bd9785004c1dc27100aa260eee0cb0f10a5de058",
    99499129,
)
SIGLIP_FULL = _siglip(
    "model_quantized.onnx",
    "9171eb00c38b9ec82f924877356d008b79e3285dbac7cd10965827bee30c9a99",
    210977441,
)
SIGLIP_TOKENIZER = _siglip(
    "tokenizer.json",
    "4a17c975210be5ab4c36b47d8dae4eefb866dbfb1e676e394aad85dc30a3ae08",
    2398744,
)


def default_backends() -> dict[str, Backend]:
    return {
        "detectScenes": Backend(
            BackendSpec(
                id="pyscenedetect-adaptive",
                model="pyscenedetect-adaptive",
                license="BSD-3-Clause",
                commercialUse=True,
                experimental=False,
                source="https://github.com/Breakthrough/PySceneDetect",
                requires=("scenedetect", "cv2"),
            ),
            _lazy("scenes", "load_scene_detector"),
        ),
        "transcribe": Backend(
            BackendSpec(
                id="faster-whisper-small",
                model="faster-whisper-small",
                license="MIT",
                commercialUse=True,
                experimental=False,
                source="https://huggingface.co/Systran/faster-whisper-small",
                files=(
                    _whisper(
                        "config.json",
                        "b55496ac7940a7ae47d2c01eab40edfd8701feec1229d9cce3b40014383fb828",
                        2370,
                    ),
                    _whisper(
                        "model.bin",
                        "3e305921506d8872816023e4c273e75d2419fb89b24da97b4fe7bce14170d671",
                        483546902,
                    ),
                    _whisper(
                        "tokenizer.json",
                        "fb7b63191e9bb045082c79fd742a3106a12c99513ab30df4a0d47fa6cb6fd0ab",
                        2203239,
                    ),
                    _whisper(
                        "vocabulary.txt",
                        "34ce3fe1c5041027b3f8d42912270993f986dbc4bb34cf27f951e34a1e453913",
                        459861,
                    ),
                ),
                requires=("faster_whisper",),
            ),
            _lazy("transcribe", "load_transcriber"),
        ),
        "detectFaces": Backend(YUNET, _lazy("faces", "load_face_detector")),
        "trackFaces": Backend(
            BackendSpec(
                id=YUNET.id,
                model="opencv-yunet-2023mar+iou-tracker",
                license="MIT",
                commercialUse=True,
                experimental=False,
                source=YUNET.source,
                files=YUNET.files,
                requires=YUNET.requires,
            ),
            _lazy("faces", "load_face_detector"),
        ),
        "faceLandmarks": Backend(
            BackendSpec(
                id=YUNET.id,
                model="opencv-yunet-2023mar-5pt",
                license="MIT",
                commercialUse=True,
                experimental=False,
                source=YUNET.source,
                files=YUNET.files,
                requires=YUNET.requires,
            ),
            _lazy("faces", "load_face_detector"),
        ),
        "embedFrames": Backend(
            BackendSpec(
                id="siglip-base-patch16-224",
                model="siglip-base-patch16-224",
                license="Apache-2.0",
                commercialUse=True,
                experimental=False,
                source=SIGLIP_SOURCE,
                files=(SIGLIP_VISION,),
                requires=("onnxruntime",),
            ),
            _lazy("embeddings", "load_frame_embedder"),
        ),
        "classifyContent": Backend(
            BackendSpec(
                id="siglip-base-patch16-224",
                model="siglip-base-patch16-224-zero-shot",
                license="Apache-2.0",
                commercialUse=True,
                # Zero-shot prompts, not a trained moderation head.
                experimental=True,
                source=SIGLIP_SOURCE,
                files=(SIGLIP_FULL, SIGLIP_TOKENIZER),
                requires=("onnxruntime", "tokenizers"),
            ),
            _lazy("embeddings", "load_content_classifier"),
        ),
        "segmentPrompted": Backend(
            BackendSpec(
                id="sam2.1-hiera-tiny",
                model="sam2.1-hiera-tiny",
                license="Apache-2.0",
                commercialUse=True,
                experimental=True,
                source="https://github.com/facebookresearch/sam2",
                files=(
                    ModelFile(
                        "sam2.1_hiera_tiny.pt",
                        f"{HF}/facebook/sam2.1-hiera-tiny/resolve/{SAM2_REV}/sam2.1_hiera_tiny.pt",
                        "7402e0d864fa82708a20fbd15bc84245c2f26dff0eb43a4b5b93452deb34be69",
                        156008466,
                    ),
                ),
                requires=("torch", "sam2"),
            ),
            _lazy("segment", "load_segmenter"),
        ),
        "inpaintVideo": Backend(
            BackendSpec(
                id="lama-big-onnx",
                model="big-lama+dis-flow",
                license="Apache-2.0",
                commercialUse=True,
                experimental=True,
                source="https://huggingface.co/Carve/LaMa-ONNX",
                files=(
                    ModelFile(
                        "lama_fp32.onnx",
                        f"{HF}/Carve/LaMa-ONNX/resolve/{LAMA_REV}/lama_fp32.onnx",
                        "1faef5301d78db7dda502fe59966957ec4b79dd64e16f03ed96913c7a4eb68d6",
                        208044816,
                    ),
                ),
                requires=("onnxruntime", "cv2"),
            ),
            _lazy("inpaint", "load_inpainter"),
        ),
        "redact": Backend(
            BackendSpec(
                id="opencv-redact",
                model="opencv-redact",
                license="Apache-2.0",
                commercialUse=True,
                experimental=False,
                source="https://github.com/opencv/opencv",
                requires=("cv2",),
            ),
        ),
        "synthesizeIdentity": Backend(
            BackendSpec(
                id="synthetic-identity",
                model=None,
                license=None,
                commercialUse=None,
                experimental=True,
                source=None,
                unavailable_reason="no commercially licensed model",
            ),
        ),
        "audioFeatures": Backend(
            BackendSpec(
                id="ffmpeg-ebur128-librosa",
                model="ffmpeg-ebur128+librosa-beat",
                license="ISC",
                commercialUse=True,
                experimental=False,
                source="https://github.com/librosa/librosa",
                requires=("librosa",),
            ),
        ),
        "saliency": Backend(
            BackendSpec(
                id="spectral-residual",
                model="spectral-residual",
                license="Apache-2.0",
                commercialUse=True,
                experimental=False,
                source="https://doi.org/10.1109/CVPR.2007.383267",
                requires=("cv2",),
            ),
        ),
    }


# Opt-in backends that are not clean enough to be defaults.
SFACE = Backend(
    BackendSpec(
        id="sface-2021dec",
        model="opencv-sface-2021dec",
        license="Apache-2.0 (weights); training data provenance unverified",
        commercialUse=None,
        experimental=True,
        source="https://github.com/opencv/opencv_zoo/tree/main/models/face_recognition_sface",
        files=(
            ModelFile(
                "face_recognition_sface_2021dec.onnx",
                f"{HF}/opencv/face_recognition_sface/resolve/{SFACE_REV}/face_recognition_sface_2021dec.onnx",
                "0ba9fbfa01b5270c96627c4ef784da859931e02f04419c829e83484087c34e79",
                38696353,
            ),
        ),
        requires=("cv2",),
    ),
    _lazy("faces", "load_identity_embedder"),
)
OPTIONAL_BACKENDS = {"faceIdentity": SFACE}


def _sha256(path: str) -> str:
    digest = hashlib.sha256()
    with open(path, "rb") as handle:
        for chunk in iter(lambda: handle.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _module_available(name: str) -> bool:
    try:
        return importlib.util.find_spec(name) is not None
    except (ImportError, ValueError):
        return False


@dataclass
class Registry:
    settings: Any
    backends: dict[str, Backend] = field(default_factory=default_backends)
    optional: dict[str, Backend] = field(default_factory=lambda: dict(OPTIONAL_BACKENDS))
    _runtimes: dict[str, Any] = field(default_factory=dict)
    _verified: set[str] = field(default_factory=set)
    _failed: dict[str, str] = field(default_factory=dict)
    _lock: threading.Lock = field(default_factory=threading.Lock)
    _file_locks: dict[str, threading.Lock] = field(default_factory=dict)

    # -- availability -------------------------------------------------------

    def _license_block(self, spec: BackendSpec, key: str) -> str | None:
        if spec.commercialUse is True:
            return None
        if key in self.settings.optional_backends or spec.id in self.settings.optional_backends:
            return None
        return "backend is not licensed for commercial use and not enabled in MEDIA_OPTIONAL_BACKENDS"

    def unavailable_reason(self, backend: Backend, key: str) -> str | None:
        spec = backend.spec
        if spec.unavailable_reason:
            return spec.unavailable_reason
        blocked = self._license_block(spec, key)
        if blocked:
            return blocked
        missing = [name for name in spec.requires if not _module_available(name)]
        if missing:
            return f"python module not installed: {', '.join(missing)}"
        if spec.id in self._failed:
            return self._failed[spec.id]
        if not self.settings.allow_model_download:
            absent = [f.name for f in spec.files if not os.path.isfile(self._file_path(spec, f))]
            if absent:
                return "model files not installed and MEDIA_MODEL_DOWNLOAD=0"
        return None

    def backend(self, operation: str) -> Backend:
        return self.backends[operation]

    def optional_enabled(self, key: str) -> bool:
        backend = self.optional.get(key)
        return backend is not None and self.unavailable_reason(backend, key) is None

    # -- runtime -------------------------------------------------------------

    def runtime(self, operation: str) -> Any:
        """Return the loaded runtime for an operation, downloading weights on first use."""
        return self._runtime(self.backends[operation], operation)

    def optional_runtime(self, key: str) -> Any:
        return self._runtime(self.optional[key], key)

    def _runtime(self, backend: Backend, key: str) -> Any:
        reason = self.unavailable_reason(backend, key)
        if reason:
            raise Unsupported(reason)
        cache_key = f"{backend.spec.id}:{key}"
        with self._lock:
            if cache_key in self._runtimes:
                return self._runtimes[cache_key]
        files = {f.name: self.ensure_file(backend.spec, f) for f in backend.spec.files}
        runtime = backend.load(files, self.settings) if backend.load else None
        with self._lock:
            self._runtimes.setdefault(cache_key, runtime)
            return self._runtimes[cache_key]

    # -- files ---------------------------------------------------------------

    def _file_path(self, spec: BackendSpec, model_file: ModelFile) -> str:
        return os.path.join(self.settings.models_dir, spec.id, model_file.name)

    def _file_lock(self, path: str) -> threading.Lock:
        with self._lock:
            return self._file_locks.setdefault(path, threading.Lock())

    def ensure_file(self, spec: BackendSpec, model_file: ModelFile) -> str:
        path = self._file_path(spec, model_file)
        with self._file_lock(path):
            if path in self._verified:
                return path
            if not os.path.isfile(path):
                if not self.settings.allow_model_download:
                    raise Unsupported(f"model file {model_file.name} is not installed")
                self._download(model_file, path)
            ok = os.path.getsize(path) == model_file.sizeBytes and _sha256(path) == model_file.sha256
            if not ok:
                os.unlink(path)
                reason = f"model file {model_file.name} failed SHA256 verification"
                self._failed[spec.id] = reason
                raise Unsupported(reason)
            self._verified.add(path)
            return path

    def _download(self, model_file: ModelFile, path: str) -> None:
        os.makedirs(os.path.dirname(path), exist_ok=True)
        partial = f"{path}.part"
        log.info("downloading %s", model_file.url)
        try:
            with (
                urllib.request.urlopen(model_file.url, timeout=60) as response,
                open(partial, "wb") as handle,
            ):
                written = 0
                for chunk in iter(lambda: response.read(1 << 20), b""):
                    written += len(chunk)
                    if written > model_file.sizeBytes:
                        raise ValueError("download larger than pinned size")
                    handle.write(chunk)
            os.replace(partial, path)
        except Exception as error:
            if os.path.exists(partial):
                os.unlink(partial)
            raise Unsupported(f"model file {model_file.name} could not be downloaded: {error}") from error

    # -- capabilities ----------------------------------------------------------

    def operation_capabilities(self) -> dict[str, dict[str, Any]]:
        result: dict[str, dict[str, Any]] = {}
        for operation in OPERATIONS:
            backend = self.backends[operation]
            spec = backend.spec
            result[operation] = {
                "available": self.unavailable_reason(backend, operation) is None,
                "experimental": spec.experimental,
                "model": spec.model,
                "license": spec.license,
                "commercialUse": spec.commercialUse,
            }
        return result

    def prefetch(self) -> None:
        """Download and verify every default backend's files (operator warm-up)."""
        for operation in OPERATIONS:
            backend = self.backends[operation]
            if self.unavailable_reason(backend, operation) is None:
                for model_file in backend.spec.files:
                    self.ensure_file(backend.spec, model_file)