"""Push this project to the Hugging Face Hub.

A bare `hf upload <repo> .` would push .venv (1.4 GB), 9,688 dataset images and
__pycache__ to the Hub -- `hf upload` does not read .gitignore. This script runs
the upload with the right excludes and publishes MODEL_CARD.md as the Hub's
README.md, leaving the GitHub README.md alone.

    python push_to_hub.py                 # dry run: list what would be uploaded
    python push_to_hub.py --push          # do it
    python push_to_hub.py --push --data   # include ./data (~181 MB, 9.7k images)
    python push_to_hub.py --push --create-pr
"""

import argparse
import os
import subprocess
import sys

REPO_ID = "GAD-Research-Lab/MedicalAI-Light-Weight"
ROOT = os.path.dirname(os.path.abspath(__file__))

# `hf upload` matches these with fnmatch against repo-relative POSIX paths.
EXCLUDE = [
    ".venv/*",
    ".git/*",
    "**/__pycache__/*",
    "*.pyc",
    "build/*",
    "dist/*",
    "*.spec",
    "update_config.json",
    "results.csv",
    "analyzing_images_for_ai.md",          # local scratch notes (gitignored)
    "blip-xray-finetuned/xray_blip.pth",   # {'epoch': N} stub, no weights -- misleading on the Hub
    "MODEL_CARD.md",   # uploaded separately, as README.md
    "README.md",       # GitHub README; the model card takes its place on the Hub
]
DATA_EXCLUDE = ["data/*"]


def hf_executable():
    """Prefer the `hf` next to the running interpreter, so a venv run stays in its venv."""
    bindir = os.path.dirname(sys.executable)
    for candidate in (os.path.join(bindir, "hf.exe"), os.path.join(bindir, "hf")):
        if os.path.exists(candidate):
            return candidate
    return "hf"


def files_to_upload(exclude):
    """Mirror hf upload's own filtering so the dry run is accurate."""
    from huggingface_hub.utils import filter_repo_objects

    paths = []
    for dirpath, dirnames, filenames in os.walk(ROOT):
        dirnames[:] = [d for d in dirnames if d not in (".git", ".venv", "__pycache__")]
        for name in filenames:
            full = os.path.join(dirpath, name)
            paths.append(os.path.relpath(full, ROOT).replace(os.sep, "/"))
    return sorted(filter_repo_objects(paths, allow_patterns=None, ignore_patterns=exclude))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--push", action="store_true", help="Actually upload (default is a dry run)")
    ap.add_argument("--data", action="store_true", help="Include ./data (~181 MB of images)")
    ap.add_argument("--create-pr", action="store_true", help="Open a PR instead of committing to main")
    ap.add_argument("--private", action="store_true", help="Create the repo private if it does not exist")
    ap.add_argument("--repo-id", default=REPO_ID)
    ap.add_argument("--message", default="Upload models, ONNX exports and application code")
    args = ap.parse_args()

    exclude = list(EXCLUDE) + ([] if args.data else DATA_EXCLUDE)

    files = files_to_upload(exclude)
    total = sum(os.path.getsize(os.path.join(ROOT, f)) for f in files)
    print(f"{len(files)} files, {total / 1e9:.2f} GB -> {args.repo_id}\n")
    for f in files:
        size = os.path.getsize(os.path.join(ROOT, f))
        print(f"  {size / 1e6:>9.2f} MB  {f}" if size > 1e6 else f"  {'':>9}     {f}")

    if not args.push:
        print("\nDry run. Re-run with --push to upload.")
        return 0

    hf = hf_executable()
    common = ["--repo-type", "model"]
    if args.create_pr:
        common.append("--create-pr")
    if args.private:
        common.append("--private")

    # 1. Model card first, so the repo is never briefly published without one.
    print("\n-> Uploading MODEL_CARD.md as README.md")
    card = subprocess.run(
        [hf, "upload", args.repo_id, os.path.join(ROOT, "MODEL_CARD.md"), "README.md",
         "--commit-message", "Add model card", *common],
        cwd=ROOT,
    )
    if card.returncode != 0:
        print("Model card upload failed; stopping before the bulk upload.", file=sys.stderr)
        return card.returncode

    # 2. Everything else.
    print("\n-> Uploading project files")
    bulk = subprocess.run(
        [hf, "upload", args.repo_id, ".", "--commit-message", args.message,
         "--exclude", *exclude, *common],
        cwd=ROOT,
    )
    if bulk.returncode != 0:
        print("Bulk upload FAILED.", file=sys.stderr)
        return bulk.returncode

    # Never trust the exit code alone -- confirm against the Hub. A 403 on the LFS
    # endpoint can still leave small files committed, which looks like success.
    from huggingface_hub import HfApi

    remote = set(HfApi().list_repo_files(args.repo_id))
    missing = [f for f in files if f not in remote and f != "MODEL_CARD.md"]
    if missing:
        print(f"\n{len(missing)} file(s) did NOT land on the Hub:", file=sys.stderr)
        for f in missing:
            print(f"  {f}", file=sys.stderr)
        return 1

    print(f"\nVerified {len(files)} files: https://huggingface.co/{args.repo_id}")
    return 0


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