"""V-8: first-run model download, checksum-verified, then fully offline.

The bundled product ships without weights. On first run the launcher fetches the Gemma 4
E4B QAT GGUF over HTTPS, verifies its SHA-256, and stores it in the model cache; every run
after that is offline. This is the ONLY network event the product initiates, and it
transfers model weights only, never user content (V-8, V-13).

Standalone module: it imports nothing from the rest of the package, so it can run from the
launcher before anything else is wired up, and `model.py` can import it without a cycle.
"""
from __future__ import annotations

import hashlib
import os
import sys
import urllib.request
from pathlib import Path

# Official Google QAT q4_0 (chosen in the baseline: equal quality to Q4_K_M, smaller).
MODEL_FILENAME = "gemma-4-E4B-qat-q4_0.gguf"
MODEL_URL = ("https://huggingface.co/google/gemma-4-E4B-it-qat-q4_0-gguf/"
             "resolve/main/gemma-4-E4B_q4_0-it.gguf")
MODEL_SHA256 = "e8b6a059ba86947a44ace84d6e5679795bc41862c25c30513142588f0e9dba1d"
MODEL_SIZE = 5154939136
# Documented fallback source (llama.cpp org, ungated), same family, if the primary moves.
MODEL_URL_FALLBACK = ("https://huggingface.co/ggml-org/gemma-4-E4B-it-GGUF/"
                      "resolve/main/gemma-4-E4B-it-Q4_K_M.gguf")


def model_cache_dir() -> Path:
    env = os.environ.get("PRIVSWEEP_MODEL_DIR")
    if env:
        return Path(env)
    rt = os.environ.get("PRIVSWEEP_RUNTIME_DIR")
    if rt:
        return Path(rt).parent / "model cache"  # delivered/<model cache> beside runtime/
    home = os.environ.get("PRIVSWEEP_HOME") or str(Path.home())
    return Path(home) / ".privsweep" / "models"


def cached_model_path() -> Path:
    return model_cache_dir() / MODEL_FILENAME


def is_present(path: "Path | None" = None) -> bool:
    p = path or cached_model_path()
    return p.exists() and p.stat().st_size == MODEL_SIZE


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


def verify(path: "Path | None" = None, *, expected: str = MODEL_SHA256) -> bool:
    p = path or cached_model_path()
    return p.exists() and sha256_of(p) == expected


def _download(url: str, dest: Path, progress: bool) -> None:
    req = urllib.request.Request(url, headers={"User-Agent": "privsweep"})
    with urllib.request.urlopen(req) as resp, open(dest, "wb") as out:  # HTTPS, follows redirects
        total = int(resp.headers.get("Content-Length", 0))
        done = 0
        while True:
            block = resp.read(1 << 20)
            if not block:
                break
            out.write(block)
            done += len(block)
            if progress and total:
                pct = done * 100 // total
                print(f"\rDownloading model weights (not your content): "
                      f"{pct}% ({done >> 20}/{total >> 20} MB)", end="", file=sys.stderr)
        if progress:
            print("", file=sys.stderr)


def ensure_model(progress: bool = True, *, url: str = MODEL_URL,
                 expected_sha: str = MODEL_SHA256) -> Path:
    """Download + checksum-verify the model iff absent. Idempotent and offline-safe.

    Returns the cached model path. Raises RuntimeError on a checksum mismatch (the bad
    download is discarded so a retry re-fetches).
    """
    dest = cached_model_path()
    if is_present(dest):
        return dest
    dest.parent.mkdir(parents=True, exist_ok=True)
    tmp = dest.with_name(dest.name + ".part")
    _download(url, tmp, progress)
    got = sha256_of(tmp)
    if got != expected_sha:
        try:
            tmp.unlink()
        except OSError:
            pass
        raise RuntimeError(f"model checksum mismatch (got {got[:12]}...); download discarded")
    tmp.replace(dest)
    return dest


if __name__ == "__main__":  # allow: python -m privsweep.model_download
    path = ensure_model()
    print(f"Model ready: {path}")
