"""Shard downloader — fetches model shards from peers or 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 hashlib
import json
import shutil
import tarfile
import tempfile
import urllib.parse
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Any

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


def compute_shard_checksum(shard_dir: Path) -> str:
    """Return a stable SHA256 checksum for all regular files in a shard."""
    digest = hashlib.sha256()
    for path in sorted(p for p in shard_dir.rglob("*") if p.is_file()):
        rel_path = path.relative_to(shard_dir).as_posix()
        digest.update(rel_path.encode())
        digest.update(b"\0")
        with path.open("rb") as f:
            for chunk in iter(lambda: f.read(1024 * 1024), b""):
                digest.update(chunk)
        digest.update(b"\0")
    return digest.hexdigest()


def write_shard_archive(shard_dir: Path, out_file: Any) -> None:
    """Write a tar archive for *shard_dir* to a binary file-like object."""
    with tarfile.open(fileobj=out_file, mode="w|") as archive:
        for path in sorted(p for p in shard_dir.rglob("*") if p.is_file()):
            archive.add(path, arcname=path.relative_to(shard_dir).as_posix())


def _safe_extract_shard(archive_path: Path, target_dir: Path) -> None:
    target_root = target_dir.resolve()
    with tarfile.open(archive_path, mode="r") as archive:
        for member in archive.getmembers():
            dest = (target_dir / member.name).resolve()
            if target_root != dest and target_root not in dest.parents:
                raise ValueError("peer shard archive contains an unsafe path")
        archive.extractall(target_dir)


def _peer_download_url(
    endpoint: str,
    model: str,
    shard_start: int,
    shard_end: int,
) -> str:
    query = urllib.parse.urlencode({
        "model": model,
        "shard_start": shard_start,
        "shard_end": shard_end,
    })
    return f"{endpoint.rstrip('/')}/v1/shards/download?{query}"


def _download_shard_from_peer(
    peer: dict,
    model: str,
    shard_start: int,
    shard_end: int,
    shard_dir: Path,
    timeout: float,
) -> bool:
    endpoint = peer.get("endpoint")
    checksum = peer.get("checksum")
    if not isinstance(endpoint, str) or not isinstance(checksum, str):
        return False

    shard_dir.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(prefix="meshnet-peer-", dir=shard_dir.parent) as tmp:
        tmp_root = Path(tmp)
        archive_path = tmp_root / "shard.tar"
        extract_dir = tmp_root / "extract"
        extract_dir.mkdir()
        try:
            with urllib.request.urlopen(
                _peer_download_url(endpoint, model, shard_start, shard_end),
                timeout=timeout,
            ) as resp, archive_path.open("wb") as out:
                while True:
                    chunk = resp.read(1024 * 1024)
                    if not chunk:
                        break
                    out.write(chunk)
            _safe_extract_shard(archive_path, extract_dir)
            if compute_shard_checksum(extract_dir) != checksum:
                return False
            if shard_dir.exists():
                shutil.rmtree(shard_dir)
            shutil.move(str(extract_dir), str(shard_dir))
            return True
        except Exception:
            return False


def _download_model_source(
    source: dict,
    shard_dir: Path,
    timeout: float,
) -> Path | None:
    url = source.get("url")
    if not isinstance(url, str) or not url:
        endpoint = source.get("endpoint")
        if not isinstance(endpoint, str):
            return None
        url = f"{endpoint.rstrip('/')}/v1/model-files/download"
    shard_dir.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(prefix="meshnet-model-source-", dir=shard_dir.parent) as tmp:
        tmp_root = Path(tmp)
        archive_path = tmp_root / "model-files.tar"
        extract_dir = tmp_root / "extract"
        extract_dir.mkdir()
        try:
            with urllib.request.urlopen(url, timeout=timeout) as resp, archive_path.open("wb") as out:
                while True:
                    chunk = resp.read(1024 * 1024)
                    if not chunk:
                        break
                    out.write(chunk)
            _safe_extract_shard(archive_path, extract_dir)
            shutil.move(str(extract_dir), str(shard_dir))
            return shard_dir
        except Exception:
            return None


def _download_huggingface_subset(
    hf_repo: str,
    cache_dir: Path,
    shard_dir: Path,
    allow_patterns: list[str] | None,
) -> Path:
    from huggingface_hub import snapshot_download  # type: ignore[import]

    kwargs = {
        "repo_id": hf_repo,
        "cache_dir": str(cache_dir),
        "local_dir": str(shard_dir),
    }
    if allow_patterns:
        kwargs["allow_patterns"] = allow_patterns
    try:
        return Path(snapshot_download(**kwargs))
    except TypeError:
        kwargs.pop("allow_patterns", None)
        return Path(snapshot_download(**kwargs))


def _download_from_fastest_source(
    *,
    model_sources: list[dict],
    hf_repo: str,
    cache_dir: Path,
    shard_dir: Path,
    progress: bool,
    timeout: float,
) -> tuple[str, Path] | None:
    shard_dir.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(prefix="meshnet-race-", dir=shard_dir.parent) as tmp:
        tmp_root = Path(tmp)
        jobs: dict[Any, tuple[str, Path]] = {}
        pool = ThreadPoolExecutor(max_workers=min(4, len(model_sources) + 1))
        try:
            for index, source in enumerate(model_sources):
                label = str(source.get("type") or "model-source")
                candidate = tmp_root / f"source-{index}"
                jobs[pool.submit(_download_model_source, source, candidate, timeout)] = (label, candidate)
            allow_patterns = _allow_patterns_from_sources(model_sources)
            hf_candidate = tmp_root / "huggingface"
            jobs[pool.submit(_download_huggingface_subset, hf_repo, cache_dir, hf_candidate, allow_patterns)] = (
                "HuggingFace",
                hf_candidate,
            )
            for future in as_completed(jobs):
                label, candidate = jobs[future]
                try:
                    result = future.result()
                except Exception:
                    continue
                if result is None:
                    continue
                if shard_dir.exists():
                    shutil.rmtree(shard_dir)
                shutil.move(str(candidate), str(shard_dir))
                if progress:
                    print(f"  download source: {label}", flush=True)
                pool.shutdown(wait=False, cancel_futures=True)
                return label, shard_dir
        finally:
            pool.shutdown(wait=False, cancel_futures=True)
    return None


def _allow_patterns_from_sources(model_sources: list[dict]) -> list[str] | None:
    patterns: set[str] = set()
    for source in model_sources:
        for rel in source.get("files") or []:
            if isinstance(rel, str) and rel and not rel.startswith("/") and ".." not in Path(rel).parts:
                patterns.add(rel)
    return sorted(patterns) if patterns else None


def download_shard(
    model: str,
    shard_start: int,
    shard_end: int,
    cache_dir: Path = _DEFAULT_CACHE,
    hf_repo: str | None = None,
    progress: bool = True,
    peers: list[dict] | None = None,
    model_sources: list[dict] | None = None,
    peer_timeout: float = _PEER_TIMEOUT_SECONDS,
) -> 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}"

    for peer in peers or []:
        if progress:
            print(f"  Trying peer shard download from {peer.get('endpoint')} ...", flush=True)
        if _download_shard_from_peer(
            peer,
            model,
            shard_start,
            shard_end,
            shard_dir,
            timeout=peer_timeout,
        ):
            if progress:
                print("  download source: peer", flush=True)
            return shard_dir

    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

    if progress:
        print(
            f"  Downloading layers {shard_start}-{shard_end} from {hf_repo} ...",
            flush=True,
        )
    if model_sources:
        if progress:
            print("  Racing tracker model source against HuggingFace ...", flush=True)
        raced = _download_from_fastest_source(
            model_sources=model_sources,
            hf_repo=hf_repo,
            cache_dir=cache_dir,
            shard_dir=shard_dir,
            progress=progress,
            timeout=peer_timeout,
        )
        if raced is not None:
            return raced[1]

    if progress:
        print("  download source: HuggingFace", flush=True)

    return _download_huggingface_subset(hf_repo, cache_dir, shard_dir, None)