"""Shard downloader — fetches model shards from peers or HuggingFace Hub.

Cache layout: ~/.cache/meshnet/shards/<model>/

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 os
import shutil
import tarfile
import tempfile
import time
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Any

_DEFAULT_CACHE = Path.home() / ".cache" / "meshnet" / "shards"
_PEER_TIMEOUT_SECONDS = 2.0
# Model-source tar streams are multi-GB; a short socket timeout must not kill
# them on a transient read stall. Peer probes keep the short timeout because
# they run sequentially before the race and may hit dead endpoints.
_MODEL_SOURCE_TIMEOUT_SECONDS = 30.0
_PROGRESS_INTERVAL_BYTES = 512 * 1024 * 1024
_FILE_RETRY_ATTEMPTS = 3


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."""
    # dereference: HF cache snapshots are symlinks into blobs/ — ship contents.
    with tarfile.open(fileobj=out_file, mode="w|", dereference=True) 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
            _merge_tree(extract_dir, shard_dir)
            return True
        except Exception:
            return False


class _TarOnlySource(Exception):
    """The server ignored ?file= and streamed a tar — no single-file support."""


def _fetch_source_file(
    file_url: str,
    dest: Path,
    timeout: float,
    on_chunk=None,
) -> int:
    """Download one file; skip if *dest* already matches the remote size.

    Returns the file's final size. Raises on any failure, including a
    short read against the server's Content-Length.
    """
    with urllib.request.urlopen(file_url, timeout=timeout) as resp:
        ctype = resp.getheader("Content-Type") if hasattr(resp, "getheader") else None
        if isinstance(ctype, str) and "x-tar" in ctype:
            raise _TarOnlySource(file_url)
        length = resp.getheader("Content-Length") if hasattr(resp, "getheader") else None
        expected = int(length) if isinstance(length, str) and length.isdigit() else None
        if expected is not None and dest.exists() and dest.stat().st_size == expected:
            if on_chunk is not None:
                on_chunk(expected)
            return expected  # complete from an earlier attempt/run
        received = 0
        with dest.open("wb") as out:
            while True:
                chunk = resp.read(1024 * 1024)
                if not chunk:
                    break
                out.write(chunk)
                received += len(chunk)
                if on_chunk is not None:
                    on_chunk(len(chunk))
        if expected is not None and received != expected:
            raise OSError(f"short read: got {received} of {expected} bytes")
        return received


class _SourceProgress:
    """tqdm bar over the whole per-file download (total / speed / ETA),
    falling back to plain per-file prints when tqdm is unavailable."""

    def __init__(self, enabled: bool, label: str, total_bytes: int | None):
        self._label = label
        self._bar = None
        self._enabled = enabled
        if enabled and total_bytes:
            try:
                from tqdm import tqdm  # type: ignore[import]

                self._bar = tqdm(
                    total=total_bytes,
                    unit="B",
                    unit_scale=True,
                    unit_divisor=1024,
                    desc=f"Downloading ({label})",
                )
            except Exception:
                self._bar = None

    def add_bytes(self, n: int) -> None:
        if self._bar is not None:
            self._bar.update(n)

    def rewind(self, n: int) -> None:
        """Roll the bar back after a failed attempt so retries don't double-count."""
        if self._bar is not None and n:
            self._bar.update(-n)

    def file_done(self, index: int, count: int, rel: str, size: int, reused: bool) -> None:
        note = "already complete, " if reused else ""
        line = f"  {self._label}: [{index}/{count}] {rel} ({note}{size / 1e9:.2f} GB)"
        if self._bar is not None:
            self._bar.set_postfix_str(f"{index}/{count} files", refresh=False)
        elif self._enabled:
            print(line, flush=True)

    def message(self, text: str) -> None:
        if self._bar is not None:
            self._bar.write(text)
        elif self._enabled:
            print(text, flush=True)

    def close(self) -> None:
        if self._bar is not None:
            self._bar.close()


def _reuse_local_file(expected: int, dest: Path, final: Path) -> bool:
    """Reuse an already-complete copy of a file instead of re-downloading."""
    if dest.exists() and dest.stat().st_size == expected:
        return True
    if final.exists() and final.stat().st_size == expected:
        try:
            os.link(final, dest)
        except OSError:
            shutil.copy2(final, dest)
        return True
    return False


def _valid_source_rel_files(source: dict) -> list[str]:
    return [
        rel for rel in (source.get("files") or [])
        if isinstance(rel, str) and rel and not rel.startswith("/") and ".." not in Path(rel).parts
    ]


def _source_files_cached(source: dict, shard_dir: Path) -> bool:
    rel_files = _valid_source_rel_files(source)
    if not rel_files:
        return False
    sizes = source.get("file_sizes")
    if not isinstance(sizes, dict):
        return False
    for rel in rel_files:
        expected = sizes.get(rel)
        path = shard_dir / rel
        if not isinstance(expected, int) or not path.exists() or path.stat().st_size != expected:
            return False
    return True


def _merge_tree(src: Path, dest: Path) -> None:
    dest.mkdir(parents=True, exist_ok=True)
    for path in sorted(p for p in src.rglob("*") if p.is_file()):
        target = dest / path.relative_to(src)
        target.parent.mkdir(parents=True, exist_ok=True)
        if target.exists():
            target.unlink()
        shutil.move(str(path), str(target))


def _download_source_files(
    source: dict,
    shard_dir: Path,
    timeout: float,
    progress: bool,
    label: str,
) -> Path | None:
    """Per-file download from a model source — retries and resumes per file.

    Far more robust than one multi-GB tar stream on flaky links: a dropped
    connection costs at most one file, and completed files are skipped when
    the download is retried or the node is restarted.
    """
    url = source.get("url")
    rel_files = _valid_source_rel_files(source)
    if not isinstance(url, str) or not url or not rel_files:
        return None
    sizes = source.get("file_sizes")
    if not isinstance(sizes, dict):
        sizes = {}
    partial_dir = shard_dir.parent / f"{shard_dir.name}.partial"
    partial_dir.mkdir(parents=True, exist_ok=True)
    sep = "&" if "?" in url else "?"
    known_sizes = [sizes.get(rel) for rel in rel_files]
    total_bytes = sum(s for s in known_sizes if isinstance(s, int)) if all(
        isinstance(s, int) for s in known_sizes
    ) else None
    tracker_bar = _SourceProgress(progress, label, total_bytes)
    try:
        for index, rel in enumerate(rel_files, start=1):
            dest = partial_dir / rel
            dest.parent.mkdir(parents=True, exist_ok=True)
            known_size = sizes.get(rel)
            if isinstance(known_size, int) and _reuse_local_file(known_size, dest, shard_dir / rel):
                tracker_bar.add_bytes(known_size)
                tracker_bar.file_done(index, len(rel_files), rel, known_size, reused=True)
                continue
            file_url = f"{url}{sep}{urllib.parse.urlencode({'file': rel})}"
            for attempt in range(1, _FILE_RETRY_ATTEMPTS + 1):
                counted = 0

                def _on_chunk(n: int) -> None:
                    nonlocal counted
                    counted += n
                    tracker_bar.add_bytes(n)

                try:
                    size = _fetch_source_file(file_url, dest, timeout, on_chunk=_on_chunk)
                    tracker_bar.file_done(index, len(rel_files), rel, size, reused=False)
                    break
                except _TarOnlySource:
                    tracker_bar.message(f"  {label}: no single-file support — using tar stream")
                    return None
                except Exception as exc:
                    tracker_bar.rewind(counted)
                    tracker_bar.message(
                        f"  {label}: {rel} attempt {attempt}/{_FILE_RETRY_ATTEMPTS} failed: {exc!r}"
                    )
                    if attempt == _FILE_RETRY_ATTEMPTS:
                        return None
                    time.sleep(1.0 * attempt)
    finally:
        tracker_bar.close()
    _merge_tree(partial_dir, shard_dir)
    try:
        partial_dir.rmdir()
    except OSError:
        pass
    return shard_dir


def _download_model_source(
    source: dict,
    shard_dir: Path,
    timeout: float,
    progress: bool = False,
    label: str = "model-source",
) -> Path | None:
    # Prefer per-file transfers whenever the source advertises its file list;
    # fall through to the tar stream if the server lacks single-file support
    # or per-file transfers keep failing.
    if source.get("files"):
        fetched = _download_source_files(source, shard_dir, timeout, progress, label)
        if fetched is not None:
            return fetched
    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:
            received = 0
            next_report = _PROGRESS_INTERVAL_BYTES
            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)
                    received += len(chunk)
                    if progress and received >= next_report:
                        print(f"  {label}: {received / 1e9:.1f} GB received ...", flush=True)
                        next_report += _PROGRESS_INTERVAL_BYTES
            if progress:
                print(f"  {label}: transfer complete ({received / 1e9:.2f} GB), extracting ...", flush=True)
            _safe_extract_shard(archive_path, extract_dir)
            _merge_tree(extract_dir, shard_dir)
            return shard_dir
        except Exception as exc:
            if progress:
                print(f"  {label}: download failed ({url}): {exc!r}", flush=True)
            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 _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 _allow_patterns_from_remote_index(
    hf_repo: str,
    cache_dir: Path,
    shard_start: int,
    shard_end: int,
) -> list[str] | None:
    """Fetch just the SafeTensors index + config (a few KB) from HF and compute
    which weight files the assigned layer range needs, so a HuggingFace fallback
    download stays layer-scoped even when the tracker has no model_sources
    (e.g. it has no local snapshot for this repo cached yet)."""
    try:
        from huggingface_hub import hf_hub_download  # type: ignore[import]

        from .safetensors_selection import (
            INDEX_FILENAME,
            METADATA_FILENAMES,
            layers_from_config_dict,
            select_files_for_layers_from_index,
        )

        index_path = hf_hub_download(repo_id=hf_repo, filename=INDEX_FILENAME, cache_dir=str(cache_dir))
        weight_map = json.loads(Path(index_path).read_text(encoding="utf-8")).get("weight_map")
    except Exception:
        return None
    if not isinstance(weight_map, dict):
        return None

    total_layers: int | None = None
    try:
        config_path = hf_hub_download(repo_id=hf_repo, filename="config.json", cache_dir=str(cache_dir))
        config = json.loads(Path(config_path).read_text(encoding="utf-8"))
        total_layers = layers_from_config_dict(config)
    except Exception:
        pass

    selected = select_files_for_layers_from_index(
        weight_map, shard_start, shard_end, total_layers=total_layers
    )
    return sorted(selected | METADATA_FILENAMES)


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
    if progress:
        print(f"  Target location: {shard_dir}", flush=True)

    for source in model_sources or []:
        label = str(source.get("type") or "model-source")
        if _source_files_cached(source, shard_dir):
            if progress:
                print(f"  [{label}] requested files already cached at {shard_dir}", flush=True)
            return shard_dir

    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,
        )
    # Tracker (or peer) model sources are preferred outright — usually LAN-fast.
    # HuggingFace is only the fallback when every advertised source fails.
    for source in model_sources or []:
        label = str(source.get("type") or "model-source")
        if progress:
            print(f"  Downloading from {label} model source (HuggingFace is the fallback) ...", flush=True)
        fetched = _download_model_source(
            source,
            shard_dir,
            timeout=max(peer_timeout, _MODEL_SOURCE_TIMEOUT_SECONDS),
            progress=progress,
            label=label,
        )
        if fetched is not None:
            if progress:
                print(f"  download source: {label}", flush=True)
            return fetched
    if model_sources and progress:
        print("  All model sources failed — falling back to HuggingFace ...", flush=True)

    allow_patterns = None
    if model_sources:
        allow_patterns = _allow_patterns_from_sources(model_sources)
    if allow_patterns is None:
        allow_patterns = _allow_patterns_from_remote_index(hf_repo, cache_dir, shard_start, shard_end)
    if progress:
        if allow_patterns:
            print("  download source: HuggingFace (layer-filtered)", flush=True)
        else:
            print(
                "  download source: HuggingFace (full snapshot — no SafeTensors index found)",
                flush=True,
            )

    return _download_huggingface_subset(hf_repo, cache_dir, shard_dir, allow_patterns)