#!/usr/bin/env python3
"""Push a checkpoint to HuggingFace."""
import os
import sys
import glob

def main():
    if len(sys.argv) < 3:
        print("Usage: push_to_hf.py <local_filename> <hf_filename>")
        sys.exit(1)

    local_name = sys.argv[1]
    hf_name = sys.argv[2]

    from huggingface_hub import HfApi
    api = HfApi(token=os.environ.get("HF_TOKEN"))

    # Search for the file
    matches = glob.glob(f"checkpoints/{local_name}")
    if not matches:
        matches = glob.glob(f"checkpoints/*{local_name}*")

    if not matches:
        print(f"No checkpoint matching '{local_name}' found in checkpoints/")
        sys.exit(1)

    for f in matches:
        size_mb = os.path.getsize(f) / (1024 * 1024)
        if size_mb < 1:
            print(f"Skipping {f} — too small ({size_mb:.1f} MB), likely corrupted")
            continue
        print(f"Uploading {f} ({size_mb:.0f} MB) -> {hf_name}")
        api.upload_file(
            path_or_fileobj=f,
            path_in_repo=f"checkpoints/{hf_name}",
            repo_id="hardiksa/arcisvlm",
            repo_type="model",
        )
        print("Done.")

if __name__ == "__main__":
    main()