#!/usr/bin/env python3
"""Collect one HF model into local staging, optionally upload to Scaleway.

Usage:
  python3 scripts/collect_model.py \\
    --repo-id Qwen/Qwen3-8B --local-name Qwen__Qwen3-8B --revision main

  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

Layout (local staging, still required for huggingface_hub):
  {base}/models/{local_name}/{revision}/   # R / W
  {base}/strategic/{local_name}/{revision}/  # S
  {base}/staging/...

S3 prefix (SoT once bucket is live):
  s3://{bucket}/models|strategic/{local_name}/{revision}/
  s3://{bucket}/manifests/{local_name}/{revision}/MANIFEST.json

Credentials: AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY, or SCW_ACCESS_KEY +
SCW_SECRET_KEY. Never commit keys.
"""
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",
        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)
    args = ap.parse_args()

    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}

    if args.s3_bucket:
        prefix = f"{dest_kind}/{args.local_name}/{args.revision}"
        storage_class = "GLACIER" if args.tier == "s" else "ONEZONE_IA"
        uploaded = upload_tree(
            final,
            bucket=args.s3_bucket,
            prefix=prefix,
            endpoint=args.s3_endpoint,
            region=args.s3_region,
            storage_class=storage_class,
        )
        man_key = f"manifests/{args.local_name}/{args.revision}/MANIFEST.json"
        upload_file(
            final / "MANIFEST.json",
            bucket=args.s3_bucket,
            key=man_key,
            endpoint=args.s3_endpoint,
            region=args.s3_region,
            storage_class="STANDARD",
        )
        result["s3"] = {
            "bucket": args.s3_bucket,
            "prefix": f"s3://{args.s3_bucket}/{prefix}/",
            "manifest": f"s3://{args.s3_bucket}/{man_key}",
            "storage_class": storage_class,
            "objects": uploaded,
        }

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


def _s3_client(endpoint: str, region: str):
    try:
        import boto3
    except ImportError as exc:
        raise SystemExit("boto3 required for --s3-bucket (pip install boto3)") from exc
    access = os.environ.get("AWS_ACCESS_KEY_ID") or os.environ.get("SCW_ACCESS_KEY")
    secret = os.environ.get("AWS_SECRET_ACCESS_KEY") or os.environ.get("SCW_SECRET_KEY")
    if not access or not secret:
        raise SystemExit("set AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY or SCW_ACCESS_KEY/SCW_SECRET_KEY")
    return boto3.client(
        "s3",
        region_name=region,
        endpoint_url=endpoint,
        aws_access_key_id=access,
        aws_secret_access_key=secret,
    )


def upload_file(
    path: Path,
    *,
    bucket: str,
    key: str,
    endpoint: str,
    region: str,
    storage_class: str,
) -> None:
    client = _s3_client(endpoint, region)
    extra = {"StorageClass": storage_class} if storage_class else {}
    client.upload_file(str(path), bucket, key, ExtraArgs=extra)


def upload_tree(
    root: Path,
    *,
    bucket: str,
    prefix: str,
    endpoint: str,
    region: str,
    storage_class: str,
) -> int:
    n = 0
    for p in sorted(root.rglob("*")):
        if not p.is_file() or p.name.startswith("."):
            continue
        rel = p.relative_to(root).as_posix()
        key = f"{prefix.rstrip('/')}/{rel}"
        upload_file(
            p,
            bucket=bucket,
            key=key,
            endpoint=endpoint,
            region=region,
            storage_class=storage_class,
        )
        n += 1
        print(f"uploaded {key}", file=sys.stderr)
    return n


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