"""Downloads the Whisper model files at build time and verifies each one.

Used only in the Dockerfile (``model`` stage). The revision is a fixed commit
of the Hugging Face repository and every file has a pinned SHA-256: a changed
or tampered file stops the build. At runtime the container has no network and
``HF_HUB_OFFLINE=1``.

Model: Systran/faster-whisper-small (CTranslate2 conversion of OpenAI Whisper
``small``, MIT licence). The weights are stored as float16; faster-whisper
quantises them to int8 on load (``compute_type='int8'``).
"""

from __future__ import annotations

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

REPOSITORY = "Systran/faster-whisper-small"
REVISION = "536b0662742c02347bc0e980a01041f333bce120"
BASE_URL = f"https://huggingface.co/{REPOSITORY}/resolve/{REVISION}"
FILES: dict[str, str] = {
    "config.json": "b55496ac7940a7ae47d2c01eab40edfd8701feec1229d9cce3b40014383fb828",
    "model.bin": "3e305921506d8872816023e4c273e75d2419fb89b24da97b4fe7bce14170d671",
    "tokenizer.json": "fb7b63191e9bb045082c79fd742a3106a12c99513ab30df4a0d47fa6cb6fd0ab",
    "vocabulary.txt": "34ce3fe1c5041027b3f8d42912270993f986dbc4bb34cf27f951e34a1e453913",
}
CHUNK_BYTES = 1024 * 1024
TIMEOUT_SEC = 120


class ChecksumMismatch(Exception):
    """A downloaded file does not have the pinned SHA-256."""


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


def verify(path: Path, expected: str) -> None:
    actual = sha256_of(path)
    if actual != expected:
        raise ChecksumMismatch(f"{path.name}: expected {expected}, got {actual}")


def download(url: str, target: Path) -> None:
    if not url.startswith("https://"):
        raise ValueError("only https downloads are allowed")
    with urllib.request.urlopen(url, timeout=TIMEOUT_SEC) as response, target.open("wb") as out:  # noqa: S310
        for chunk in iter(lambda: response.read(CHUNK_BYTES), b""):
            out.write(chunk)


def fetch_all(target_dir: Path, files: dict[str, str] = FILES, base_url: str = BASE_URL) -> None:
    target_dir.mkdir(parents=True, exist_ok=True)
    for name, expected in files.items():
        target = target_dir / name
        download(f"{base_url}/{name}", target)
        verify(target, expected)
        print(f"ok {name}")


if __name__ == "__main__":
    fetch_all(Path(sys.argv[1] if len(sys.argv) > 1 else "/models/small"))