import os
import re
import json
import sys

def clean_json(content):
    # Step 1: Remove comments (lines starting with //)
    content = re.sub(r'^//.*\n', '', content, flags=re.MULTILINE)

    # Step 2: Replace single quotes with double quotes for JSON compatibility
    content = content.replace("'", '"')

    # Step 3: Remove trailing commas in both objects and arrays
    # Remove trailing commas in arrays
    content = re.sub(r',\s*([\]}])', r'\1', content)
    # Remove trailing commas in objects
    content = re.sub(r',\s*(\n|\})', r'\1', content)

    # Step 4: Escape backslashes in LaTeX snippets
    content = content.replace("\\", "\\\\")  # Double backslashes for JSON

    # Step 5: Parse and re-encode to ensure valid JSON (removes trailing commas)
    try:
        data = json.loads(content)  # Parse the content into a Python dictionary
        cleaned_json = json.dumps(data, indent=4)  # Re-encode into JSON, pretty-printed
    except json.JSONDecodeError as e:
        print(f"Error parsing JSON: {e}")
        return None

    return cleaned_json

def process_directory(directory):
    # Loop through files in the specified directory (non-recursively)
    for filename in os.listdir(directory):
        file_path = os.path.join(directory, filename)

        # Only process files that end with .json
        if os.path.isfile(file_path) and filename.endswith('.json'):
            print(f"Processing {file_path}")

            # Read the file content
            with open(file_path, 'r', encoding='utf-8') as file:
                content = file.read()

            # Clean the JSON content
            cleaned_content = clean_json(content)

            if cleaned_content is not None:
                # Write the cleaned content back to the file
                with open(file_path, 'w', encoding='utf-8') as file:
                    file.write(cleaned_content)
                print(f"Fixed formatting in {file_path}")
            else:
                print(f"Failed to fix {file_path} due to JSON parsing error.")

def main():
    # Check if a directory was passed as an argument, otherwise read from stdin
    if len(sys.argv) > 1:
        directory = sys.argv[1]
    else:
        print("Please provide a directory path as an argument.")
        sys.exit(1)

    # Check if the directory exists
    if not os.path.isdir(directory):
        print(f"The provided path '{directory}' is not a valid directory.")
        sys.exit(1)

    # Process all .json files in the directory
    process_directory(directory)

if __name__ == '__main__':
    main()
