"""Upload NativeSpecZ-FM-76M to Hugging Face Hub.

Usage:
  HF_TOKEN=<your_write_token> python upload_to_hf.py
or:
  python upload_to_hf.py --token <your_write_token>

The script:
  1. Creates the repo `tempAstro/NativeSpecZ-FM-76M` (if doesn't exist)
  2. Uploads the contents of this submission folder
  3. Sets the README.md as the model card

Token rotation: REVOKE the token after upload via https://huggingface.co/settings/tokens
"""
from __future__ import annotations

import argparse
import os
from pathlib import Path


HF_USER = "ManmohanSharma"
REPO_NAME = "NativeSpecZ-FM-76M"
REPO_ID = f"{HF_USER}/{REPO_NAME}"

# Files to upload from this folder
SUBMISSION_FILES = [
    "README.md",
    "NativeSpecZ-FM-76M.ipynb",
    "weights/best.pt",
    "weights/training_args.json",
    "weights/best_metrics.json",
    "eval_results/desi_2500_metrics.json",
    "code/hybrid_redshift.py",
    "code/data.py",
    "code/metrics.py",
    "code/model.py",
    "code/plots.py",
    "code/__init__.py",
    "plots/scatter_redshift.png",
    "plots/spectrum_reconstruction.png",
    "plots/comparison_vs_aion.png",
    "plots/foundation_evidence.png",
    "plots/dashboard.png",
    "plots/multi_mask_reconstruction.png",
    "plots/stress_curve.png",
]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--token", default=os.environ.get("HF_TOKEN", ""))
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    if not args.token:
        raise SystemExit("ERROR: provide --token or set HF_TOKEN env var")

    try:
        from huggingface_hub import HfApi, create_repo
    except ImportError:
        raise SystemExit("ERROR: pip install huggingface_hub")

    here = Path(__file__).parent

    print(f"=== Uploading to {REPO_ID} ===")

    if not args.dry_run:
        try:
            create_repo(REPO_ID, repo_type="model", token=args.token, private=False, exist_ok=True)
            print(f"Repo {REPO_ID} ready")
        except Exception as e:
            print(f"create_repo: {e}")

    api = HfApi(token=args.token)
    for rel in SUBMISSION_FILES:
        local = here / rel
        if not local.exists():
            print(f"SKIP (missing): {rel}")
            continue
        size_mb = local.stat().st_size / 1e6
        print(f"Upload {rel}  ({size_mb:.1f} MB)")
        if args.dry_run:
            continue
        api.upload_file(
            path_or_fileobj=str(local),
            path_in_repo=rel,
            repo_id=REPO_ID,
            repo_type="model",
            token=args.token,
            commit_message=f"Upload {rel}",
        )
    print(f"\n=== Done. View at https://huggingface.co/{REPO_ID} ===")
    print("REMINDER: revoke the upload token at https://huggingface.co/settings/tokens")


if __name__ == "__main__":
    main()