# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""The model's files from the Hugging Face Hub, checked against this package's own list.

    from anuvani_diar.hub import bundle_dir, golden_audio_dir
    bundle_dir()          # model/ of CIIL/Anuvani-Diarization, downloaded once into the Hugging Face cache
    golden_audio_dir()    # golden/audio/: the 12 recordings anuvani-diar-check diarises

Every file's sha256 must equal the one in `assets/bundle_files.json` (or `assets/golden12.json` for the recordings),
which ship inside the package, so the weights a Diarizer loads are the ones this package was tested with whatever
revision the Hub serves. Downloads come from the release tag this package was made for (`REVISION`), or from the
main branch if that tag is missing; `ANUVANI_DIAR_REVISION` (or `revision=`) names another branch, tag or commit. No
Hugging Face account is needed; huggingface_hub's own settings apply (`HF_HUB_CACHE`, `HF_HUB_OFFLINE`,
`HF_ENDPOINT`). With `model/` comes the repository's `config.json` (the release's description): the Hub counts
requests for it as the model's downloads; nothing else is sent.
"""
import hashlib
import json
import os
from pathlib import Path
from typing import Dict, Optional

REPO = "CIIL/Anuvani-Diarization"
# the repository's tag for this release
REVISION = "v1.0.0"
_HERE = Path(__file__).resolve().parent
_CHECKED: Dict[tuple, str] = {}


def assets_dir() -> Path:
    """The package's assets: `$ANUVANI_DIAR_ASSETS` (set by the package to its own copy), else the source tree's."""
    env = os.environ.get("ANUVANI_DIAR_ASSETS")
    if env:
        return Path(env)
    return _HERE / "assets" if (_HERE / "assets").is_dir() else _HERE.parents[1] / "assets"


def sha256(path: Path) -> str:
    """The file's sha256, computed once per process for each (file, size, modification time)."""
    st = path.stat()
    key = (str(path.resolve()), st.st_size, st.st_mtime_ns)
    if key not in _CHECKED:
        h = hashlib.sha256()
        with open(path, "rb") as f:
            for block in iter(lambda: f.read(1 << 24), b""):
                h.update(block)
        _CHECKED[key] = h.hexdigest()
    return _CHECKED[key]


def _download(patterns, revision: Optional[str], token) -> Path:
    try:
        from huggingface_hub import snapshot_download
    except ImportError as e:
        raise ImportError("the model's files come from the Hugging Face Hub: pip install huggingface_hub, or pass "
                          "bundle= (or set ANUVANI_DIAR_BUNDLE) to a downloaded copy") from e
    chosen = revision or os.environ.get("ANUVANI_DIAR_REVISION")
    try:
        try:
            return Path(snapshot_download(REPO, revision=chosen or REVISION, allow_patterns=patterns, token=token))
        # the release tag is missing: the main branch, still sha256-checked
        except Exception as e:
            if chosen or type(e).__name__ != "RevisionNotFoundError":
                raise
            return Path(snapshot_download(REPO, revision=None, allow_patterns=patterns, token=token))
    except Exception as e:
        raise RuntimeError(f"could not download {patterns[0].split('/')[0]}/ from {REPO} ({type(e).__name__}: {e}). "
                           "Check the network connection; without one, download the repository elsewhere "
                           f"(hf download {REPO}) and set ANUVANI_DIAR_BUNDLE to its model/ folder.") from e


def _check(root: Path, want: Dict[str, str], what: str) -> None:
    bad = [rel for rel, h in want.items() if not (root / rel).is_file() or sha256(root / rel) != h]
    if bad:
        raise RuntimeError(f"{what}: {', '.join(bad)} under {root} do not match this package's sha256 list "
                           "(another revision of the repository?). Pin the revision this package was made for with "
                           "ANUVANI_DIAR_REVISION, or update the package.")


def bundle_dir(revision: Optional[str] = None, token=None) -> Path:
    """`model/` of the Hub repository (MANIFEST, decoder, both networks' safetensors), every file sha256-checked."""
    spec = json.loads((assets_dir() / "bundle_files.json").read_text())
    sub, want = spec["subdir"], spec["files"]
    # config.json last: the Hub counts its requests as the model's downloads
    root = _download([f"{sub}/{rel}" for rel in sorted(want)] + ["config.json"], revision, token) / sub
    _check(root, want, f"{REPO}/{sub}")
    return root


def golden_audio_dir(revision: Optional[str] = None, token=None) -> Path:
    """`golden/audio/` of the Hub repository: the release's 12 install-check recordings, sha256-checked."""
    files = json.loads((assets_dir() / "golden12.json").read_text())["files"]
    want = {f"{sid}.flac": v["flac_sha256"] for sid, v in files.items()}
    root = _download([f"golden/audio/{f}" for f in sorted(want)], revision, token) / "golden" / "audio"
    _check(root, want, f"{REPO}/golden/audio")
    return root