#!/usr/bin/env python3
"""Collect one HF model into the VAULT reserve tree (weights-first).

Usage:
  /tmp/fi-hf/bin/python scripts/collect_model.py \\
    --repo-id Qwen/Qwen3-8B --local-name Qwen__Qwen3-8B --revision main

Layout:
  {base}/models/{local_name}/{revision}/   # HF files (no full onnx dump)
  {base}/staging/...                       # incomplete
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
import sys
from pathlib import Path

# Prefer sequential downloads on WSL+drvfs (xet parallel can OOM)
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")

from huggingface_hub import snapshot_download  # noqa: E402

DEFAULT_ALLOW = (
    "*.safetensors",
    "*.bin",
    "*.json",
    "*.txt",
    "*.model",
    "*.tiktoken",
    "tokenizer*",
    "vocab*",
    "merges.txt",
    "special_tokens_map.json",
    "generation_config.json",
    "configuration*.py",
    "modeling*.py",
    "README.md",
    "LICENSE*",
    "*.md",
    # sentence-transformers / embed extras
    "1_Pooling/*",
    "modules.json",
    "sentence_*",
    "config_sentence_transformers.json",
)

IGNORE = (
    "onnx/*",
    "openvino/*",
    "flax_model*",
    "tf_model*",
    "rust_model*",
    "*.ot",
    "*.h5",
    "imgs/*",
    ".gitattributes",
)


def sha256_file(path: Path, chunk: int = 8 * 1024 * 1024) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        while True:
            b = f.read(chunk)
            if not b:
                break
            h.update(b)
    return h.hexdigest()


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo-id", required=True)
    ap.add_argument("--local-name", required=True)
    ap.add_argument("--revision", default="main")
    ap.add_argument(
        "--base",
        default="/mnt/d/vault/coulomb/freedom-intelligence",
    )
    ap.add_argument("--token", default=os.environ.get("HF_TOKEN") or None)
    args = ap.parse_args()

    base = Path(args.base)
    stage = base / "staging" / args.local_name / args.revision
    final = base / "models" / args.local_name / args.revision

    if final.exists() and any(final.rglob("*.safetensors")) or (
        final.exists() and any(final.rglob("pytorch_model*.bin"))
    ):
        print(f"already present with weights: {final}")
        return 0

    if stage.exists():
        shutil.rmtree(stage, ignore_errors=True)
    stage.mkdir(parents=True, exist_ok=True)

    print(f"downloading {args.repo_id}@{args.revision} -> {stage}")
    snapshot_download(
        repo_id=args.repo_id,
        revision=args.revision,
        local_dir=str(stage),
        allow_patterns=list(DEFAULT_ALLOW),
        ignore_patterns=list(IGNORE),
        token=args.token,
        max_workers=1,
    )

    # strip hub cache dir if present
    cache = stage / ".cache"
    if cache.exists():
        shutil.rmtree(cache, ignore_errors=True)

    final.parent.mkdir(parents=True, exist_ok=True)
    if final.exists():
        shutil.rmtree(final)
    shutil.move(str(stage), str(final))

    artifacts = []
    total = 0
    for p in sorted(final.rglob("*")):
        if not p.is_file():
            continue
        if p.name.startswith("."):
            continue
        rel = str(p.relative_to(final))
        # hash only weight-like files (cheap for small; large still OK sequential)
        if p.suffix in {".safetensors", ".bin", ".gguf", ".pt"} or p.stat().st_size > 1_000_000:
            digest = sha256_file(p)
        else:
            digest = ""
        size = p.stat().st_size
        total += size
        artifacts.append({"path": rel, "sha256": digest, "bytes": size})

    manifest = {
        "repo_id": args.repo_id,
        "revision": args.revision,
        "local_name": args.local_name,
        "storage_path": str(final),
        "total_bytes": total,
        "artifacts": artifacts,
    }
    (final / "MANIFEST.json").write_text(json.dumps(manifest, indent=2) + "\n")
    print(json.dumps({"ok": True, "total_bytes": total, "path": str(final)}, indent=2))
    return 0


if __name__ == "__main__":
    sys.exit(main())