"""
Push DANTE-Mosaic-3.5B to the Hugging Face Hub.

Prerequisites:
    pip install -U huggingface_hub transformers
    huggingface-cli login          # interactive token entry, OR
    export HF_TOKEN=hf_xxxxxxxxxx  # non-interactive

Usage:
    python push_to_hub.py \\
        --model-dir ../hf_model/ \\
        --repo  YourOrg/DANTE-Mosaic-3.5B \\
        --private                # optional, default public

Note: requires ~6 GB upload bandwidth. The huggingface_hub uploader resumes on
network errors, so re-running is safe.
"""

from __future__ import annotations

import argparse
import os
import shutil
from pathlib import Path

from huggingface_hub import HfApi, create_repo, upload_folder


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--model-dir", required=True,
                   help="Path containing model.safetensors, config.json, tokenizer files")
    p.add_argument("--repo", required=True,
                   help="HuggingFace repo id, e.g. 'YourOrg/DANTE-Mosaic-3.5B'")
    p.add_argument("--private", action="store_true",
                   help="Create as private repo (default: public)")
    p.add_argument("--readme", default="README.md",
                   help="Path to model card (will be uploaded as README.md)")
    p.add_argument("--token", default=None,
                   help="HF token (overrides $HF_TOKEN). Optional.")
    p.add_argument("--commit-message", default="Initial release of DANTE-Mosaic-3.5B")
    p.add_argument("--allow-patterns", nargs="*", default=None,
                   help="Limit upload to these glob patterns (e.g. '*.safetensors').")
    return p.parse_args()


def main():
    args = parse_args()

    token = args.token or os.environ.get("HF_TOKEN")
    if not token:
        print("WARNING: no --token and no $HF_TOKEN; relying on cached HF login.")

    model_dir = Path(args.model_dir).expanduser().resolve()
    if not model_dir.exists():
        raise FileNotFoundError(f"Model dir not found: {model_dir}")
    weight_files = list(model_dir.glob("*.safetensors")) + list(model_dir.glob("*.bin"))
    if not weight_files:
        raise FileNotFoundError(f"No weight files in {model_dir}")
    print(f"\nModel directory : {model_dir}")
    print(f"Weight files    : {len(weight_files)} "
          f"(total {sum(f.stat().st_size for f in weight_files)/1e9:.1f} GB)")

    readme_path = Path(args.readme)
    if readme_path.exists():
        target = model_dir / "README.md"
        if target.exists() and target.resolve() != readme_path.resolve():
            print(f"Replacing existing model card: {target}")
        if target.resolve() != readme_path.resolve():
            shutil.copy(readme_path, target)
        print(f"Model card      : {readme_path}  ->  {target}")
    else:
        print(f"WARNING: README {readme_path} not found; the upload will not include a model card.")

    print(f"Target repo     : {args.repo}  (private={args.private})")
    print()

    print("Creating repo on the Hub (if it does not exist)...")
    create_repo(repo_id=args.repo, private=args.private, exist_ok=True, token=token)

    print(f"Uploading folder...")
    upload_folder(
        folder_path=str(model_dir),
        repo_id=args.repo,
        commit_message=args.commit_message,
        token=token,
        allow_patterns=args.allow_patterns,
    )

    print("\nDONE!")
    print(f"Model is live at: https://huggingface.co/{args.repo}")


if __name__ == "__main__":
    main()