import json
import os
from pathlib import Path
import argparse

def process_submissions(input_file, output_dir='submissions'):
    """
    Process a JSON file containing multiple model submissions and split it into
    individual JSON files in the submissions directory.
    
    Args:
        input_file (str): Path to the input JSON file
        output_dir (str): Directory where individual submission files will be stored
    """
    # Create submissions directory if it doesn't exist
    Path(output_dir).mkdir(exist_ok=True)
    
    # Read the input file
    try:
        with open(input_file, 'r') as f:
            submissions = json.load(f)
        
        if not isinstance(submissions, list):
            print(f"Error: Input file {input_file} must contain a JSON array of submissions")
            return
        
        # Process each submission
        for submission in submissions:
            if 'model_name' not in submission:
                print(f"Warning: Skipping submission without model_name field")
                continue
                
            model_name = submission['model_name']
            # Create a safe filename from the model name
            safe_filename = f"{model_name.replace('/', '_')}.json"
            output_path = os.path.join(output_dir, safe_filename)
            
            # Write individual submission file
            with open(output_path, 'w') as f:
                json.dump(submission, f, indent=4)
            
            print(f"Created submission file: {output_path}")
        
        print(f"\nProcessed {len(submissions)} submissions successfully!")
        
    except FileNotFoundError:
        print(f"Error: Input file '{input_file}' not found")
    except json.JSONDecodeError:
        print(f"Error: Input file '{input_file}' is not valid JSON")
    except Exception as e:
        print(f"Error processing submissions: {str(e)}")

def main():
    # Set up argument parser
    parser = argparse.ArgumentParser(
        description='Process a JSON file containing model submissions and split into individual files.'
    )
    parser.add_argument(
        'input_file',
        help='Path to the input JSON file containing model submissions'
    )
    parser.add_argument(
        '--output-dir',
        '-o',
        default='submissions',
        help='Directory where individual submission files will be stored (default: submissions)'
    )
    
    # Parse arguments
    args = parser.parse_args()
    
    # Process submissions
    process_submissions(args.input_file, args.output_dir)

if __name__ == "__main__":
    main()