#!/usr/bin/env python3
"""
Push Zenith-7B model to Hugging Face Hub.

Usage:
    python push_to_hf.py --repo_id Matrix-Corp/Zenith-7b-V1 --token YOUR_TOKEN
"""

import argparse
import os
import sys
from pathlib import Path
from huggingface_hub import HfApi, login, create_repo, whoami
from huggingface_hub.utils import RepositoryNotFoundError, HfHubHTTPError


def push_model(repo_id: str, token: str = None, folder_path: str = ".", private: bool = False):
    """Push model files to Hugging Face Hub with robust error handling."""
    folder_path = Path(folder_path).resolve()

    if not folder_path.exists():
        raise ValueError(f"Folder not found: {folder_path}")

    # Check required files
    required_files = [
        "modeling_zenith.py",
        "hf_model_card.md",
        "README.md",
        "requirements.txt",
        "train.py",
        "inference.py",
        "test_model.py",
        "finetune_qwen.py",
        "Modelfile"
    ]

    missing = [f for f in required_files if not (folder_path / f).exists()]
    if missing:
        print(f"⚠️  Warning: Missing files: {missing}")
        response = input("Continue anyway? (y/N): ")
        if response.lower() != 'y':
            return

    # Authenticate
    try:
        if token:
            login(token=token)
            print("✓ Logged in with provided token")
        else:
            # Check if already logged in
            try:
                user = whoami()
                print(f"✓ Already logged in as: {user['name']}")
            except:
                print("Please login to Hugging Face:")
                login()
    except Exception as e:
        print(f"❌ Authentication failed: {e}")
        print("\nTo get a token:")
        print("1. Go to https://huggingface.co/settings/tokens")
        print("2. Create a new token with 'write' permissions")
        print("3. Run: python push_to_hf.py --token YOUR_TOKEN")
        return

    # Create API client
    api = HfApi()

    # Check if repo exists, create if not
    try:
        repo_info = api.repo_info(repo_id=repo_id, repo_type="model")
        print(f"✓ Repository exists: {repo_id}")
    except RepositoryNotFoundError:
        print(f"📝 Repository not found. Creating: {repo_id}")
        try:
            create_repo(
                repo_id=repo_id,
                token=token,
                repo_type="model",
                private=private,
                exist_ok=True
            )
            print(f"✓ Repository created")
        except Exception as e:
            print(f"❌ Failed to create repository: {e}")
            return
    except Exception as e:
        print(f"⚠️  Warning: Could not check repository: {e}")

    # Upload
    print(f"\n📤 Uploading {folder_path} to {repo_id}...")
    print("This may take a while depending on file sizes...\n")

    try:
        api.upload_folder(
            folder_path=str(folder_path),
            repo_id=repo_id,
            repo_type="model",
            commit_message=f"Upload Zenith-7B model"
        )
        print(f"\n✅ Successfully uploaded to https://huggingface.co/{repo_id}")
        print("\nNext steps:")
        print("1. Visit your model page")
        print("2. Add a model card if needed")
        print("3. Test: from transformers import AutoModel; AutoModel.from_pretrained('your-repo-id')")
    except HfHubHTTPError as e:
        if e.response.status_code == 401:
            print(f"\n❌ Unauthorized: Invalid token or no write access")
            print("   Make sure you:")
            print("   - Have a valid token with 'write' permissions")
            print("   - Own the organization/repository or have collaborator rights")
        elif e.response.status_code == 403:
            print(f"\n❌ Forbidden: You don't have permission to push to this repository")
            print("   Make sure you're a member of the organization with write access")
        elif e.response.status_code == 404:
            print(f"\n❌ Repository not found: {repo_id}")
            print("   Check the repository ID is correct")
        else:
            print(f"\n❌ HTTP Error {e.response.status_code}: {e}")
    except Exception as e:
        print(f"\n❌ Upload failed: {e}")
        print("\nTroubleshooting:")
        print("1. Check your internet connection")
        print("2. Verify you have enough disk space")
        print("3. Try logging in again: huggingface-cli login")
        print("4. Check Hugging Face status: https://status.huggingface.co")


def main():
    parser = argparse.ArgumentParser(description="Push Zenith-7B to Hugging Face")
    parser.add_argument(
        "--repo_id",
        type=str,
        default="Matrix-Corp/Zenith-7b-V1",
        help="Hugging Face repository ID (username/model-name)"
    )
    parser.add_argument(
        "--token",
        type=str,
        help="Hugging Face access token (optional if already logged in)"
    )
    parser.add_argument(
        "--folder",
        type=str,
        default=".",
        help="Folder containing model files (default: current directory)"
    )
    parser.add_argument(
        "--private",
        action="store_true",
        help="Create repository as private (default: public)"
    )

    args = parser.parse_args()
    push_model(args.repo_id, args.token, args.folder, args.private)


if __name__ == "__main__":
    main()