from huggingface_hub import login, HfApi
import json
import argparse
import os

def upload_json_to_hf(token, repo_id, file_path, file_name):
    # Login to Hugging Face
    login(token)
    
    # Initialize the API
    api = HfApi()
    
    # Upload the file
    try:
        api.upload_file(
            path_or_fileobj=file_path,
            path_in_repo=file_name,
            repo_id=repo_id,
            repo_type="dataset"
        )
        print(f"Successfully uploaded {file_name} to {repo_id}")
    except Exception as e:
        print(f"Error uploading file: {str(e)}")
        raise

def main():
    parser = argparse.ArgumentParser(description='Upload JSON file to Hugging Face')
    
    # Add arguments
    parser.add_argument(
        '--token',
        type=str,
        help='Hugging Face access token (or set HUGGINGFACE_TOKEN env variable)',
        default=os.getenv('HUGGINGFACE_TOKEN')
    )
    
    parser.add_argument(
        '--repo-id',
        type=str,
        required=True,
        help='Repository ID (format: username/repo-name)'
    )
    
    parser.add_argument(
        '--file-path',
        type=str,
        required=True,
        help='Path to the JSON file to upload'
    )
    
    parser.add_argument(
        '--file-name',
        type=str,
        help='Name to save the file as in the repository (defaults to the input filename)',
    )

    # Parse arguments
    args = parser.parse_args()

    # Validate token
    if not args.token:
        raise ValueError("Please provide a token either via --token or HUGGINGFACE_TOKEN environment variable")

    # If file_name is not provided, use the basename of file_path
    if not args.file_name:
        args.file_name = os.path.basename(args.file_path)

    # Validate file exists and is JSON
    if not os.path.exists(args.file_path):
        raise FileNotFoundError(f"File not found: {args.file_path}")
    
    try:
        with open(args.file_path, 'r') as f:
            json.load(f)  # Validate JSON format
    except json.JSONDecodeError:
        raise ValueError(f"File is not valid JSON: {args.file_path}")

    # Upload file
    upload_json_to_hf(args.token, args.repo_id, args.file_path, args.file_name)

if __name__ == "__main__":
    main()