import argparse
import os
from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer, AutoConfig
from huggingface_hub import HfApi, login
import sys

def parse_arguments():
    parser = argparse.ArgumentParser(description='Upload model to Hugging Face Hub')
    
    parser.add_argument(
        '--model_path',
        type=str,
        required=True,
        help='Local path to the model'
    )
    
    parser.add_argument(
        '--repo_name',
        type=str,
        required=True,
        help='Destination repo name (format: username/repo_name)'
    )
    
    parser.add_argument(
        '--token',
        type=str,
        required=True,
        help='Hugging Face API token'
    )
    
    parser.add_argument(
        '--model_card',
        type=str,
        default=None,
        help='Path to model card README.md file (optional)'
    )
    
    parser.add_argument(
        '--private',
        action='store_true',
        help='Make the repository private'
    )

    return parser.parse_args()

def create_default_model_card(repo_name):
    return f"""---
language: en
tags:
- pytorch
- transformer
license: apache-2.0
---

# Model Card for {repo_name}

## Model Description

[Provide a brief description of your model here]

## Intended Use

[Describe the intended use of your model]

## Training Procedure

[Describe how the model was trained]

## Evaluation Results

[Provide evaluation metrics and results]

## Limitations

[Describe any known limitations]

## Additional Information

For more information, please visit the model repository.
"""

def upload_model(model_path, repo_name, token, model_card=None, private=False):
    try:
        # Login to Hugging Face
        login(token)
        print("Successfully logged in to Hugging Face")

        # Load model, tokenizer, and config
        print("Loading model, tokenizer, and config...")
        model = AutoModelForCausalLM.from_pretrained(model_path)
        tokenizer = AutoTokenizer.from_pretrained(model_path)
        config = AutoConfig.from_pretrained(model_path)

        # Create model card if not provided
        if model_card is None:
            print("No model card provided. Creating default model card...")
            model_card_content = create_default_model_card(repo_name)
            with open("README.md", "w", encoding="utf-8") as f:
                f.write(model_card_content)
        else:
            # Copy provided model card
            with open(model_card, "r", encoding="utf-8") as f:
                model_card_content = f.read()
            with open("README.md", "w", encoding="utf-8") as f:
                f.write(model_card_content)

        # Push everything to hub
        print(f"Uploading model to {repo_name}...")
        model.push_to_hub(repo_name, private=private)
        tokenizer.push_to_hub(repo_name, private=private)
        config.push_to_hub(repo_name, private=private)

        # Clean up
        if os.path.exists("README.md"):
            os.remove("README.md")

        print(f"Successfully uploaded model to https://huggingface.co/{repo_name}")
        return True

    except Exception as e:
        print(f"Error occurred: {str(e)}")
        return False

def main():
    args = parse_arguments()
    
    success = upload_model(
        model_path=args.model_path,
        repo_name=args.repo_name,
        token=args.token,
        model_card=args.model_card,
        private=args.private
    )
    
    if success:
        sys.exit(0)
    else:
        sys.exit(1)

if __name__ == "__main__":
    main()