#!/usr/bin/env python3
"""
HuggingFace Checkpoint Uploader for AGILLM-3
Usage: python hf_upload.py <checkpoint.pt> [--message "commit msg"]
Requires: HF_TOKEN env var or ~/.huggingface/token
"""
import os
import sys
import argparse
from pathlib import Path
from datetime import datetime

try:
    from huggingface_hub import HfApi, upload_file
except ImportError:
    print("Installing huggingface_hub...")
    os.system("pip install huggingface_hub -q")
    from huggingface_hub import HfApi, upload_file

REPO_ID = "OpenTransformer/AGILLM-3-large"

def upload_checkpoint(ckpt_path: str, message: str = None):
    path = Path(ckpt_path)
    if not path.exists():
        print(f"Error: {ckpt_path} not found")
        return False
    
    token = os.environ.get("HF_TOKEN") or None
    api = HfApi(token=token)
    
    # Extract step from filename
    step = "unknown"
    if "step" in path.stem:
        step = path.stem.split("step")[-1].split(".")[0]
    
    ts = datetime.now().strftime("%Y-%m-%d %H:%M")
    commit_msg = message or f"Checkpoint step {step} - {ts}"
    
    print(f"Uploading {path.name} to {REPO_ID}...")
    print(f"Size: {path.stat().st_size / 1e9:.2f} GB")
    
    try:
        upload_file(
            path_or_fileobj=str(path),
            path_in_repo=f"checkpoints/{path.name}",
            repo_id=REPO_ID,
            commit_message=commit_msg,
            token=token
        )
        print(f"SUCCESS: https://huggingface.co/{REPO_ID}")
        return True
    except Exception as e:
        print(f"Upload failed: {e}")
        return False

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("checkpoint", help="Path to checkpoint .pt file")
    parser.add_argument("--message", "-m", help="Commit message")
    args = parser.parse_args()
    
    success = upload_checkpoint(args.checkpoint, args.message)
    sys.exit(0 if success else 1)