#!/usr/bin/env python3
"""Upload model artifact to Hugging Face Model Hub"""

import os
import sys
from pathlib import Path
from huggingface_hub import HfApi, login

def main():
    print("Uploading model artifact to Hugging Face Model Hub...")
    print("Repository: solarevat/cv200-artifact")
    print("")
    
    # Get the paths
    script_dir = Path(__file__).parent
    cv_dir = script_dir.parent / "cv"
    artifact_dir = cv_dir / "runs" / "SMOKE-T2-FULL" / "artifact"
    
    files_to_upload = [
        artifact_dir / "model.ts",
        artifact_dir / "labels.json",
        artifact_dir / "preprocess.json",
    ]
    
    # Check files exist
    for f in files_to_upload:
        if not f.exists():
            print(f"❌ Error: File not found: {f}")
            sys.exit(1)
    
    # Login - check for token in environment or command line
    token = os.environ.get("HF_TOKEN") or (sys.argv[1] if len(sys.argv) > 1 else None)
    
    if token:
        login(token=token)
        print("✅ Logged in using provided token")
    else:
        # Try to check if already logged in
        try:
            api = HfApi()
            api.whoami()
            print("✅ Already logged in to Hugging Face")
        except Exception:
            print("❌ Not logged in. Please provide a token:")
            print("   Option 1: Set environment variable: export HF_TOKEN=your_token_here")
            print("   Option 2: Pass as argument: python3 upload_model.py your_token_here")
            print("   Option 3: Login interactively: huggingface-cli login")
            print("\n   Get your token from: https://huggingface.co/settings/tokens")
            sys.exit(1)
    
    api = HfApi()
    
    # Upload files
    repo_id = "solarevat/cv200-artifact"
    print(f"\nUploading to {repo_id}...")
    
    for file_path in files_to_upload:
        print(f"  Uploading {file_path.name}...")
        api.upload_file(
            path_or_fileobj=str(file_path),
            path_in_repo=file_path.name,
            repo_id=repo_id,
            repo_type="model",
        )
    
    print("\n✅ Upload complete! The model is now available at:")
    print(f"   https://huggingface.co/{repo_id}")
    print("\nThe Dockerfile will automatically download it during build.")

if __name__ == "__main__":
    main()