import os
from pathlib import Path
from huggingface_hub import HfApi, create_repo
from dotenv import load_dotenv
import yaml

load_dotenv()

def upload_models():
    hf_token = os.getenv("HF_TOKEN")
    if not hf_token:
        print("❌ HF_TOKEN not found in .env file")
        print("Set it first: export HF_TOKEN=your_token_here")
        return False
    
    hf_repo = os.getenv("HF_MODEL_REPO", "junaid17/damagelens-models")
    username = hf_repo.split("/")[0]
    repo_name = hf_repo.split("/")[1]
    
    # Load config
    with open("model_config.yaml", "r") as f:
        config = yaml.safe_load(f)
    
    api = HfApi()
    
    # Create repo if it doesn't exist
    print(f"\n📦 Creating repository: {hf_repo}")
    try:
        api.create_repo(
            repo_id=hf_repo,
            repo_type="model",
            private=config["repository"].get("private", False),
            exist_ok=True,
            token=hf_token
        )
        print(f"✓ Repository ready: {hf_repo}")
    except Exception as e:
        print(f"⚠️  Could not create repo (may already exist): {str(e)}")
    
    # Upload each model
    checkpoints_dir = Path("checkpoints")
    for model_type, model_info in config["models"].items():
        filename = model_info["filename"]
        local_path = checkpoints_dir / filename
        
        if not local_path.exists():
            print(f"⚠️  Skipping {filename} - file not found at {local_path}")
            continue
        
        print(f"\n📤 Uploading {filename} ({model_info['size_mb']}MB)...")
        try:
            api.upload_file(
                path_or_fileobj=str(local_path),
                path_in_repo=filename,
                repo_id=hf_repo,
                repo_type="model",
                token=hf_token,
                commit_message=f"Upload {model_type} model checkpoint"
            )
            print(f"✓ Successfully uploaded {filename}")
        except Exception as e:
            print(f"❌ Failed to upload {filename}: {str(e)}")
            return False
    
    print("\n" + "="*50)
    print("✅ All models uploaded successfully!")
    print(f"📍 View at: https://huggingface.co/{hf_repo}")
    print("="*50 + "\n")
    return True


if __name__ == "__main__":
    import sys
    
    print("🚀 HuggingFace Model Uploader")
    print("="*50)
    
    success = upload_models()
    sys.exit(0 if success else 1)