"""Fetch the model weights from the Hugging Face model repo and verify their SHA-256.

Weights are never committed to git. The Docker image downloads them at build time;
a local run downloads them on first start if models/ does not already have them.

    python -m serve.weights          # download (if needed) and verify
"""
from __future__ import annotations

import hashlib
import shutil
import tempfile
import urllib.request
from pathlib import Path

from serve.config import Settings


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


def ensure_model(path: Path, url: str, expected_sha256: str | None) -> Path:
    """Return path to a verified model file, downloading it if missing."""
    path = Path(path)
    if path.exists():
        if not expected_sha256 or sha256_of(path) == expected_sha256:
            return path
        raise RuntimeError(f"{path} exists but its SHA-256 does not match the expected model.")
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile(dir=path.parent, delete=False, suffix=".part") as tmp:
        tmp_path = Path(tmp.name)
        with urllib.request.urlopen(url, timeout=120) as resp:
            shutil.copyfileobj(resp, tmp)
    digest = sha256_of(tmp_path)
    if expected_sha256 and digest != expected_sha256:
        tmp_path.unlink(missing_ok=True)
        raise RuntimeError(f"Downloaded model SHA-256 {digest} != expected {expected_sha256} ({url})")
    tmp_path.replace(path)
    return path


if __name__ == "__main__":
    s = Settings()
    p = ensure_model(s.model_path, s.model_url, s.model_sha256)
    print(f"Model ready: {p} ({p.stat().st_size / 2**20:.1f} MB, sha256 {sha256_of(p)[:12]}...)")