"""Shard downloader — fetches or stubs a model shard from HuggingFace Hub.

Cache layout: ~/.cache/meshnet/shards/<model>/layers_<start>-<end>/

For "stub-model" (no HF repo), a placeholder JSON file is written so the
test suite never touches the network.
"""

import json
from pathlib import Path

_DEFAULT_CACHE = Path.home() / ".cache" / "meshnet" / "shards"


def download_shard(
    model: str,
    shard_start: int,
    shard_end: int,
    cache_dir: Path = _DEFAULT_CACHE,
    hf_repo: str | None = None,
    progress: bool = True,
) -> Path:
    """Ensure the shard is present in *cache_dir* and return its local path.

    When *hf_repo* is None (or *model* is ``"stub-model"``), a placeholder
    weights file is created locally — no network access required.  This keeps
    the test suite hermetic while the real download path is exercised by
    passing a non-stub *hf_repo*.
    """
    shard_dir = cache_dir / model / f"layers_{shard_start}-{shard_end}"
    shard_dir.mkdir(parents=True, exist_ok=True)

    if hf_repo is None or model == "stub-model":
        stub_file = shard_dir / "weights.json"
        if not stub_file.exists():
            stub_file.write_text(json.dumps({
                "model": model,
                "shard_start": shard_start,
                "shard_end": shard_end,
                "stub": True,
            }))
            if progress:
                print(f"  [stub] shard placeholder written to {stub_file}", flush=True)
        else:
            if progress:
                print(f"  [stub] shard already cached at {shard_dir}", flush=True)
        return shard_dir

    from huggingface_hub import snapshot_download  # type: ignore[import]

    if progress:
        print(
            f"  Downloading layers {shard_start}-{shard_end} from {hf_repo} ...",
            flush=True,
        )

    local_dir = snapshot_download(
        repo_id=hf_repo,
        cache_dir=str(cache_dir),
        local_dir=str(shard_dir),
    )
    return Path(local_dir)