"""SAM2 ONNX model variant definitions and download helpers."""

import hashlib
import json
from pathlib import Path


DEFAULT_MODELS_DIR = Path.home() / ".samgis" / "models"
DEFAULT_MODEL = "sam2.1_hiera_tiny_uint8"

# HF base URL for ONNX weights
_HF_BASE = "https://huggingface.co/aletrn/sam2-onnx-weights/resolve/main"

# Local filenames expected by Sam2OnnxSession
_ENCODER = "encoder.onnx"
_DECODER = "decoder.onnx"
_METADATA = "metadata.json"

# Remote filename → local filename
_UINT8_FILES = {
    "encoder_uint8.onnx": _ENCODER,
    "decoder_uint8.onnx": _DECODER,
    _METADATA: _METADATA,
}

_FP32_FILES = {
    _ENCODER: _ENCODER,
    _DECODER: _DECODER,
    _METADATA: _METADATA,
}

# Each HF folder contains both fp32 and uint8 files; the variant key
# determines which remote files to download and the local folder name.
_VARIANTS = ("tiny", "small", "base_plus", "large")
_HF_FOLDERS = {
    "tiny": "sam2.1_hiera_tiny",
    "small": "sam2.1_hiera_small",
    "base_plus": "sam2.1_hiera_base_plus",
    "large": "sam2.1_hiera_large",
}
_SIZE_UINT8_MB = {
    "tiny": 32,
    "small": 39,
    "base_plus": 73,
    "large": 211,
}  # pragma: no mutate (cosmetic display values)
_SIZE_FP32_MB = {
    "tiny": 120,
    "small": 148,
    "base_plus": 281,
    "large": 829,
}  # pragma: no mutate (cosmetic display values)

# SHA-256 checksums keyed by HF folder → remote filename.
_CHECKSUMS_PATH = Path(__file__).parent / "checksums.json"
with open(_CHECKSUMS_PATH) as _f:
    _CHECKSUMS_RAW = json.load(_f)
_SHA256: dict[str, dict[str, str]] = {
    k: v for k, v in _CHECKSUMS_RAW.items() if not k.startswith("_")
}
CHECKSUMS_COMMIT: str = _CHECKSUMS_RAW.get("_commit", "")


def _build_variant(name: str, precision: str, files: dict, size_mb: int) -> dict:
    hf_folder = _HF_FOLDERS[name]
    folder_hashes = _SHA256.get(hf_folder, {})
    return {
        "label": (
            f"SAM 2.1 {name.replace('_', ' ').title()} ({precision}, ~{size_mb} MB)"
        ),
        "size_mb": size_mb,
        "files": files,
        "urls": {remote: f"{_HF_BASE}/{hf_folder}/{remote}" for remote in files},
        "sha256": {
            local: folder_hashes[remote]
            for remote, local in files.items()
            if remote in folder_hashes
        },
    }


MODELS = {}
for _v in _VARIANTS:
    MODELS[f"sam2.1_hiera_{_v}_uint8"] = _build_variant(
        _v,
        "uint8",
        _UINT8_FILES,
        _SIZE_UINT8_MB[_v],
    )
    MODELS[f"sam2.1_hiera_{_v}_fp32"] = _build_variant(
        _v,
        "fp32",
        _FP32_FILES,
        _SIZE_FP32_MB[_v],
    )


def get_model_dir(variant: str = DEFAULT_MODEL) -> Path:
    """Return the local directory for a model variant."""
    return DEFAULT_MODELS_DIR / variant


def is_model_downloaded(variant: str = DEFAULT_MODEL) -> bool:
    """Check if all required files exist for a model variant."""
    model_dir = get_model_dir(variant)
    files = MODELS[variant]["files"]
    return all((model_dir / local_name).is_file() for local_name in files.values())


def _sha256_file(
    path: Path, chunk_size: int = 1 << 20
) -> str:  # pragma: no mutate (any positive chunk_size works)
    """Compute SHA-256 hex digest of a file."""
    h = hashlib.sha256()
    with open(path, "rb") as f:
        while chunk := f.read(chunk_size):
            h.update(chunk)
    return h.hexdigest()


def verify_download(variant: str = DEFAULT_MODEL) -> list[str]:
    """Verify SHA-256 of downloaded files.

    Returns list of failed filenames (empty = OK).
    """
    model_dir = get_model_dir(variant)
    expected = MODELS[variant].get("sha256", {})
    failures = []
    for local_name, expected_hash in expected.items():
        path = model_dir / local_name
        if not path.is_file():
            failures.append(local_name)
            continue
        if _sha256_file(path) != expected_hash:
            failures.append(local_name)
    return failures


def get_downloaded_variants() -> list[str]:
    """Return list of variant keys whose files are all present locally."""
    return [k for k in MODELS if is_model_downloaded(k)]


def save_active_variant(variant: str) -> None:
    """Persist the active model variant in QgsSettings."""
    from qgis.core import QgsSettings

    QgsSettings().setValue("samgis/active_variant", variant)


def load_active_variant() -> str:
    """Load the persisted model variant, falling back to DEFAULT_MODEL."""
    from qgis.core import QgsSettings

    stored = QgsSettings().value("samgis/active_variant", DEFAULT_MODEL)
    if stored in MODELS:
        return stored
    return DEFAULT_MODEL
