#!/usr/bin/env python3
"""Collect an HF model: diskless S3 streaming or explicitly opted-in local cache.

Remote reserve (default when --s3-bucket is supplied; no local weight files):
  python3 scripts/collect_model.py \
    --repo-id deepseek-ai/DeepSeek-V4-Flash-0731 \
    --local-name deepseek-ai__DeepSeek-V4-Flash-0731 --tier s \
    --s3-bucket railiance-fi-open-weight-reserve

Use --dry-run to inspect the immutable source revision and transfer size.
Use --local-download without --s3-bucket only for an intentional local cache.
See docs/streaming-reserve.md for the resource-limited remote worker.
"""
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")

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",
        help="local staging root (VAULT HD)",
    )
    ap.add_argument(
        "--tier",
        choices=("r", "s", "w"),
        default="r",
        help="r/w → models/ (One Zone IA); s → strategic/ (Glacier)",
    )
    ap.add_argument("--s3-bucket", default=os.environ.get("FI_S3_BUCKET") or "")
    ap.add_argument(
        "--s3-endpoint",
        default=os.environ.get("FI_S3_ENDPOINT", "https://s3.nl-ams.scw.cloud"),
    )
    ap.add_argument(
        "--s3-region",
        default=os.environ.get("FI_S3_REGION", "nl-ams"),
    )
    ap.add_argument("--token", default=os.environ.get("HF_TOKEN") or None)
    ap.add_argument("--dry-run", action="store_true", help="list pinned S3 transfer without uploading")
    ap.add_argument("--part-mib", type=int, default=64, help="streaming S3 part buffer (5–128 MiB)")
    ap.add_argument("--local-download", action="store_true", help="explicitly permit a full local snapshot (no S3)")
    args = ap.parse_args()

    if args.s3_bucket:
        if args.local_download:
            ap.error("--local-download cannot be combined with --s3-bucket")
        from stream_model import collect
        return collect(args, DEFAULT_ALLOW, IGNORE)
    if not args.local_download:
        ap.error("choose --s3-bucket (diskless streaming) or explicitly --local-download")
    if args.dry_run:
        ap.error("--dry-run currently requires --s3-bucket")
    from huggingface_hub import snapshot_download

    base = Path(args.base)
    dest_kind = "strategic" if args.tier == "s" else "models"
    stage = base / "staging" / args.local_name / args.revision
    final = base / dest_kind / args.local_name / args.revision

    already = final.exists() and (
        any(final.rglob("*.safetensors")) or any(final.rglob("pytorch_model*.bin"))
    )
    if already:
        print(f"already present with weights: {final}", file=sys.stderr)
        if not args.s3_bucket:
            return 0
    else:
        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,
        )

        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")
    result = {"ok": True, "total_bytes": total, "path": str(final), "tier": args.tier}

    print(json.dumps(result, indent=2))
    return 0


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