"""
Upload trained model to Hugging Face Hub
"""
import argparse
import sys
import os
from huggingface_hub import HfApi, create_repo
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# Add parent directory to path (if needed)
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))


def upload_to_hub(
    model_path: str,
    repo_name: str,
    organization: str = None,
    private: bool = False
):
    """
    Upload model to Hugging Face Hub.
    
    Args:
        model_path: Path to the trained model
        repo_name: Name for the repository on Hugging Face Hub
        organization: Organization name (optional)
        private: Whether to make the repository private
    """
    print("=" * 60)
    print("Uploading Model to Hugging Face Hub")
    print("=" * 60)
    
    # Create full repo ID
    if organization:
        repo_id = f"{organization}/{repo_name}"
    else:
        repo_id = repo_name
    
    print(f"\nRepository: {repo_id}")
    print(f"Private: {private}")
    
    # Load model and tokenizer
    print("\n[1/3] Loading model...")
    try:
        tokenizer = AutoTokenizer.from_pretrained(model_path)
        model = AutoModelForSequenceClassification.from_pretrained(model_path)
        print("✓ Model loaded successfully")
    except Exception as e:
        print(f"✗ Error loading model: {e}")
        return
    
    # Create repository
    print("\n[2/3] Creating repository...")
    try:
        create_repo(
            repo_id=repo_id,
            repo_type="model",
            exist_ok=True,
            private=private
        )
        print(f"✓ Repository created/verified: {repo_id}")
    except Exception as e:
        print(f"✗ Error creating repository: {e}")
        print("\nMake sure you're logged in:")
        print("  huggingface-cli login")
        return
    
    # Push to hub
    print("\n[3/3] Uploading model and tokenizer...")
    try:
        model.push_to_hub(repo_id)
        tokenizer.push_to_hub(repo_id)
        print("✓ Upload complete!")
    except Exception as e:
        print(f"✗ Error uploading: {e}")
        return
    
    print("\n" + "=" * 60)
    print("Success! 🎉")
    print("=" * 60)
    print(f"\nYour model is now available at:")
    print(f"https://huggingface.co/{repo_id}")
    
    print("\nTo use your model:")
    print(f"""
from transformers import AutoTokenizer, AutoModelForSequenceClassification

tokenizer = AutoTokenizer.from_pretrained("{repo_id}")
model = AutoModelForSequenceClassification.from_pretrained("{repo_id}")
""")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Upload model to Hugging Face Hub")
    parser.add_argument(
        "--model-path",
        type=str,
        default="./results/final_model",
        help="Path to the trained model"
    )
    parser.add_argument(
        "--repo-name",
        type=str,
        required=True,
        help="Name for the repository on Hugging Face Hub"
    )
    parser.add_argument(
        "--organization",
        type=str,
        default=None,
        help="Organization name (optional)"
    )
    parser.add_argument(
        "--private",
        action="store_true",
        help="Make the repository private"
    )
    args = parser.parse_args()
    
    print("\nBefore uploading, make sure you:")
    print("1. Have a Hugging Face account")
    print("2. Are logged in: huggingface-cli login")
    print("3. Have reviewed the model card (MODEL_CARD.md)")
    
    response = input("\nProceed with upload? (yes/no): ")
    if response.lower() in ['yes', 'y']:
        upload_to_hub(
            args.model_path,
            args.repo_name,
            args.organization,
            args.private
        )
    else:
        print("Upload cancelled.")