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

import hashlib
import json
from pathlib import Path
from typing import TypedDict


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}
_SIZE_FP32_MB = {"tiny": 120, "small": 148, "base_plus": 281, "large": 829}

# 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", "")


class ModelInfo(TypedDict):
    label: str
    size_mb: int
    files: dict[str, str]
    urls: dict[str, str]
    sha256: dict[str, str]


def _build_variant(
    name: str, precision: str, files: dict[str, str], size_mb: int
) -> ModelInfo:
    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: dict[str, ModelInfo] = {}
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:
    """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, model_dir: Path | None = None
) -> list[str]:
    """Verify SHA-256 of downloaded files. Returns list of failed filenames (empty = OK)."""
    model_dir = model_dir or get_model_dir(variant)
    expected: dict[str, str] = MODELS[variant]["sha256"]
    failures: list[str] = []
    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
