#!/usr/bin/env python3
"""
Download all files from a HuggingFace model repository with retries,
resume support, and integrity verification.

Usage:
    python3 download_hf_model.py <repo_id> [repo_id ...]

Examples:
    python3 download_hf_model.py ai-sage/GigaAM-Multilingual
    python3 download_hf_model.py meta-llama/Llama-3-8B-Instruct

Each repo is downloaded into a directory named by replacing '/' with '.'
in the repo_id (e.g. ai-sage.GigaAM-Multilingual).

Adapted from download_lilith_q8.py.
"""

import argparse
import hashlib
import importlib
import os
import subprocess
import sys
import tempfile
import time
from pathlib import Path

MAX_RETRIES = 10
RETRY_DELAY_BASE = 15  # seconds, doubles each retry

BASE_DIR = Path(__file__).resolve().parent
RUNTIME_DIR = BASE_DIR / ".hf_download_runtime"
HF_HOME_DIR = RUNTIME_DIR / "hf-home"
HF_HUB_CACHE_DIR = HF_HOME_DIR / "hub"
HF_ASSETS_CACHE_DIR = HF_HOME_DIR / "assets"
XDG_CACHE_HOME_DIR = RUNTIME_DIR / "xdg-cache"
PIP_CACHE_DIR = RUNTIME_DIR / "pip-cache"
PYTHON_VENDOR_DIR = RUNTIME_DIR / "python-vendor"
TMP_DIR = RUNTIME_DIR / "tmp"


def prepare_runtime_dirs() -> None:
    for path in (
        RUNTIME_DIR,
        HF_HOME_DIR,
        HF_HUB_CACHE_DIR,
        HF_ASSETS_CACHE_DIR,
        XDG_CACHE_HOME_DIR,
        PIP_CACHE_DIR,
        PYTHON_VENDOR_DIR,
        TMP_DIR,
    ):
        path.mkdir(parents=True, exist_ok=True)

    os.environ["HF_HOME"] = str(HF_HOME_DIR)
    os.environ["HF_HUB_CACHE"] = str(HF_HUB_CACHE_DIR)
    os.environ["HF_ASSETS_CACHE"] = str(HF_ASSETS_CACHE_DIR)
    os.environ["XDG_CACHE_HOME"] = str(XDG_CACHE_HOME_DIR)
    os.environ["PIP_CACHE_DIR"] = str(PIP_CACHE_DIR)
    os.environ["HF_HUB_DISABLE_XET"] = "1"
    os.environ["HF_XET_RECONSTRUCT_WRITE_SEQUENTIALLY"] = "1"
    os.environ["TMPDIR"] = str(TMP_DIR)
    os.environ["TMP"] = str(TMP_DIR)
    os.environ["TEMP"] = str(TMP_DIR)
    tempfile.tempdir = str(TMP_DIR)

    if str(PYTHON_VENDOR_DIR) not in sys.path:
        sys.path.insert(0, str(PYTHON_VENDOR_DIR))


prepare_runtime_dirs()


def parse_major_minor(version: str) -> tuple[int, int]:
    parts = []
    for chunk in version.split("."):
        digits = "".join(ch for ch in chunk if ch.isdigit())
        if digits:
            parts.append(int(digits))
        if len(parts) == 2:
            break
    while len(parts) < 2:
        parts.append(0)
    return parts[0], parts[1]


def install_huggingface_hub() -> None:
    print("Installing newer huggingface_hub into the local runtime dir...")
    subprocess.check_call(
        [
            sys.executable,
            "-m",
            "pip",
            "install",
            "--upgrade",
            "--target",
            str(PYTHON_VENDOR_DIR),
            "--cache-dir",
            str(PIP_CACHE_DIR),
            "huggingface_hub>=1.0,<2",
        ]
    )


try:
    import huggingface_hub as _huggingface_hub

    if parse_major_minor(_huggingface_hub.__version__) < (1, 0):
        raise ImportError(f"Found old huggingface_hub {_huggingface_hub.__version__}")

    from huggingface_hub import HfApi, hf_hub_url
    from huggingface_hub.utils import RepositoryNotFoundError
except Exception:
    install_huggingface_hub()
    importlib.invalidate_caches()
    for name in list(sys.modules):
        if name == "huggingface_hub" or name.startswith("huggingface_hub."):
            del sys.modules[name]

    from huggingface_hub import HfApi, hf_hub_url
    from huggingface_hub.utils import RepositoryNotFoundError


def load_token() -> str | None:
    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    if token:
        return token

    token_files = [
        HF_HOME_DIR / "token",
        Path.home() / ".cache" / "huggingface" / "token",
    ]
    for token_file in token_files:
        if token_file.exists():
            return token_file.read_text().strip()
    return None


TOKEN = load_token()
if not TOKEN:
    print("HF token not found. Proceeding without authentication (repo may be public).")
else:
    print(f"Using HF token: {TOKEN[:8]}...")


def repo_to_dirname(repo_id: str) -> str:
    return repo_id.replace("/", ".")


def get_dir_size(path: Path) -> int:
    return sum(f.stat().st_size for f in path.rglob("*") if f.is_file())


def expected_size_for(sibling) -> int | None:
    size = getattr(sibling, "size", None)
    if size is not None:
        return size
    if sibling.lfs and sibling.lfs.get("size") is not None:
        return sibling.lfs["size"]
    return None


def format_size(n: int | None) -> str:
    if n is None:
        return "unknown size"
    if n >= 1e9:
        return f"{n / 1e9:.2f} GB"
    if n >= 1e6:
        return f"{n / 1e6:.2f} MB"
    return f"{n / 1e3:.2f} KB"


def quarantine_corrupt_files(paths: list[Path], download_dir: Path) -> None:
    if not paths:
        return
    quarantine_dir = download_dir / ".corrupt"
    quarantine_dir.mkdir(exist_ok=True)
    stamp = int(time.time())
    for path in paths:
        target = quarantine_dir / f"{path.name}.{stamp}.corrupt"
        suffix = 1
        while target.exists():
            target = quarantine_dir / f"{path.name}.{stamp}.{suffix}.corrupt"
            suffix += 1
        path.replace(target)
        print(f"Moved corrupt file aside: {target}")


def download_file(rel_path: str, dest: Path, repo_id: str) -> None:
    url = hf_hub_url(repo_id=repo_id, filename=rel_path)
    cmd = [
        "curl",
        "-L",
        "--fail",
        "--show-error",
        "--connect-timeout",
        "30",
        "--progress-bar",
        "-C",
        "-",
        "-o",
        str(dest),
        url,
    ]
    if TOKEN:
        cmd[1:1] = ["-H", f"Authorization: Bearer {TOKEN}"]
    subprocess.check_call(cmd)


def download_repo_with_retries(repo_id: str, download_dir: Path) -> Path:
    """Download all files from a repo with curl, resume, and exponential backoff."""
    for attempt in range(1, MAX_RETRIES + 1):
        try:
            print(f"\n{'=' * 60}")
            print(f"Download attempt {attempt}/{MAX_RETRIES}")
            print(f"{'=' * 60}")

            api = HfApi(token=TOKEN)
            repo_info = api.repo_info(repo_id=repo_id, token=TOKEN, files_metadata=True)
            siblings = [s for s in (repo_info.siblings or []) if s.rfilename]

            if attempt == 1:
                # Disk space check
                needed = 0
                for s in siblings:
                    exp = expected_size_for(s)
                    dest = download_dir / s.rfilename
                    have = dest.stat().st_size if dest.exists() else 0
                    if exp is not None:
                        needed += max(exp - have, 0)
                try:
                    stat = os.statvfs(download_dir)
                    free = stat.f_bavail * stat.f_frsize
                    print(f"Disk check: need ~{format_size(needed)} more, free ~{format_size(free)} on destination.")
                    if needed > free:
                        print("WARNING: not enough free space for the remaining download.")
                except OSError:
                    pass

            siblings = sorted(
                siblings,
                key=lambda s: expected_size_for(s) or 0,
                reverse=True,
            )

            size_before = get_dir_size(download_dir) if download_dir.exists() else 0
            t0 = time.time()

            for index, sibling in enumerate(siblings, start=1):
                rel_path = sibling.rfilename
                dest = download_dir / rel_path
                dest.parent.mkdir(parents=True, exist_ok=True)

                expected_size = expected_size_for(sibling)
                current_size = dest.stat().st_size if dest.exists() else 0

                if expected_size is not None and dest.exists() and current_size == expected_size:
                    print(
                        f"[{index}/{len(siblings)}] Already present: {rel_path} "
                        f"({format_size(expected_size)})"
                    )
                    continue

                if expected_size is not None and dest.exists() and current_size > expected_size:
                    quarantine_corrupt_files([dest], download_dir)
                    current_size = 0

                if current_size > 0:
                    print(
                        f"[{index}/{len(siblings)}] Resuming: {rel_path} "
                        f"from {format_size(current_size)} / {format_size(expected_size)}"
                    )
                else:
                    print(f"[{index}/{len(siblings)}] Downloading: {rel_path} ({format_size(expected_size)})")

                download_file(rel_path, dest, repo_id)

                if expected_size is not None:
                    actual_size = dest.stat().st_size
                    if actual_size != expected_size:
                        raise ValueError(
                            f"Size mismatch after download for {rel_path}: "
                            f"expected {expected_size}, got {actual_size}"
                        )

            elapsed = time.time() - t0
            size_after = get_dir_size(download_dir)
            downloaded = max(size_after - size_before, 0)
            speed = downloaded / elapsed if elapsed > 0 else 0

            print(f"\nDownload complete: {download_dir}")
            print(f"  Downloaded: {format_size(downloaded)} in {elapsed:.0f}s")
            print(f"  Avg speed:  {speed / 1e6:.1f} MB/s")
            return download_dir

        except KeyboardInterrupt:
            print("\nInterrupted by user.")
            sys.exit(1)

        except (RepositoryNotFoundError, FileNotFoundError) as exc:
            print(f"\nFatal error: {exc}")
            sys.exit(1)

        except Exception as exc:
            delay = RETRY_DELAY_BASE * (2 ** (attempt - 1))
            print(f"\nError on attempt {attempt}: {exc}")
            if attempt < MAX_RETRIES:
                print(f"Retrying in {delay}s...")
                time.sleep(delay)
            else:
                print("All retries exhausted.")
                sys.exit(1)


def verify_integrity(repo_id: str, download_dir: Path) -> tuple[bool, list[Path]]:
    """Verify all downloaded files against repo metadata."""
    print(f"\n{'=' * 60}")
    print("Verifying file integrity...")
    print(f"{'=' * 60}")

    api = HfApi(token=TOKEN)
    try:
        repo_info = api.repo_info(repo_id=repo_id, token=TOKEN, files_metadata=True)
    except Exception as exc:
        print(f"Could not fetch repo metadata for verification: {exc}")
        return False, []

    siblings = [s for s in (repo_info.siblings or []) if s.rfilename]
    errors = []
    checked = 0
    corrupt_files: list[Path] = []

    for sibling in siblings:
        if not sibling.rfilename:
            continue

        file_path = download_dir / sibling.rfilename
        if not file_path.exists():
            errors.append(f"MISSING: {sibling.rfilename}")
            continue

        expected_sha = None
        hash_type = None
        if sibling.lfs and sibling.lfs.get("sha256"):
            expected_sha = sibling.lfs["sha256"]
            hash_type = "sha256"
        elif hasattr(sibling, "blob_id") and sibling.blob_id:
            expected_sha = sibling.blob_id
            hash_type = "git-sha1"
        else:
            continue

        checked += 1
        size = file_path.stat().st_size
        print(f"  Checking {sibling.rfilename} ({format_size(size)})...", end=" ", flush=True)

        if hash_type == "sha256":
            digest = hashlib.sha256()
            with open(file_path, "rb") as handle:
                while chunk := handle.read(8 * 1024 * 1024):
                    digest.update(chunk)
            actual = digest.hexdigest()
        else:
            digest = hashlib.sha1()
            digest.update(f"blob {size}\0".encode())
            with open(file_path, "rb") as handle:
                while chunk := handle.read(8 * 1024 * 1024):
                    digest.update(chunk)
            actual = digest.hexdigest()

        if actual == expected_sha:
            print("OK")
        else:
            print("MISMATCH")
            corrupt_files.append(file_path)
            errors.append(
                f"HASH MISMATCH: {sibling.rfilename}\n"
                f"  expected: {expected_sha}\n"
                f"  actual:   {actual}"
            )

    print(f"\nChecked {checked} file(s).")

    if errors:
        print("\nIntegrity errors:")
        for err in errors:
            print(f"  - {err}")
        return False, corrupt_files

    print("All files verified successfully.")
    return True, []


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Download files from HuggingFace model repositories with retry and verification."
    )
    parser.add_argument(
        "repos",
        nargs="+",
        help="HuggingFace repo IDs to download (e.g. ai-sage/GigaAM-Multilingual)",
    )
    args = parser.parse_args()

    for repo_id in args.repos:
        dirname = repo_to_dirname(repo_id)
        download_dir = BASE_DIR / dirname
        download_dir.mkdir(parents=True, exist_ok=True)

        print(f"\n{'#' * 60}")
        print(f"# Repo:    {repo_id}")
        print(f"# Target:  {download_dir}")
        print(f"{'#' * 60}")

        download_repo_with_retries(repo_id, download_dir)

        ok, corrupt_files = verify_integrity(repo_id, download_dir)
        if not ok:
            quarantine_corrupt_files(corrupt_files, download_dir)
            print(f"\nIntegrity check FAILED for {repo_id}. Re-run the script to resume.")
            sys.exit(1)

        print(f"\nDone. Model saved to: {download_dir}")


if __name__ == "__main__":
    main()