"""
Upload ML models to Hugging Face Hub
This allows the models to be loaded in Hugging Face Spaces
"""

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

def upload_models():
    """Upload models to Hugging Face Hub."""
    
    # Check if models exist
    models_dir = Path("models")
    if not models_dir.exists():
        print("Error: models/ directory not found!")
        print("Please train the models first:")
        print("  python train_conflict_model.py")
        print("  python generate_embeddings.py")
        return
    
    # Check for model files
    model_files = {
        "conflict_predictor.pkl": models_dir / "conflict_predictor.pkl",
        "package_embeddings.json": models_dir / "package_embeddings.json",
        "embedding_info.json": models_dir / "embedding_info.json"
    }
    
    missing = [name for name, path in model_files.items() if not path.exists()]
    if missing:
        print(f"Error: Missing model files: {missing}")
        print("Please train the models first:")
        print("  python train_conflict_model.py")
        print("  python generate_embeddings.py")
        return
    
    # Login to Hugging Face
    print("Logging in to Hugging Face...")
    print("(You'll need to enter your HF token - get it from https://huggingface.co/settings/tokens)")
    try:
        login()
    except Exception as e:
        print(f"Login error: {e}")
        print("\nYou can also set HF_TOKEN environment variable:")
        print("  $env:HF_TOKEN='your_token_here'  # PowerShell")
        return
    
    # Initialize API
    api = HfApi()
    
    # Repository name for models
    repo_id = "ysakhale/dependency-conflict-models"
    
    # Create repository if it doesn't exist
    try:
        api.create_repo(
            repo_id=repo_id,
            repo_type="model",
            exist_ok=True,
            private=False
        )
        print(f"Repository {repo_id} is ready!")
    except Exception as e:
        print(f"Note: {e}")
    
    # Upload each model file
    print("\nUploading models...")
    for filename, filepath in model_files.items():
        print(f"Uploading {filename}...")
        try:
            api.upload_file(
                path_or_fileobj=str(filepath),
                path_in_repo=filename,
                repo_id=repo_id,
                repo_type="model"
            )
            print(f"  ✓ {filename} uploaded successfully!")
        except Exception as e:
            print(f"  ✗ Error uploading {filename}: {e}")
    
    print(f"\n✅ Models uploaded to: https://huggingface.co/{repo_id}")
    print("\nNext step: Update ml_models.py to load from this repository")

if __name__ == "__main__":
    upload_models()