# Guys please use the same setup code and put into the imports for consistency
script_names = ["untitled14.py", "untitled13.py"]  # This is a list of the files of the mission runs
content = ""
imports = """
from pybricks.hubs import PrimeHub
from pybricks.pupdevices import Motor, ColorSensor
from pybricks.parameters import Port, Stop, Color, Direction
from pybricks.robotics import DriveBase
from pybricks.tools import wait, StopWatch, multitask, run_task
import asyncio

hub = PrimeHub()
left_motor = Motor(Port.A, Direction.COUNTERCLOCKWISE)
right_motor = Motor(Port.B)
atarm1 = Motor(Port.E, Direction.COUNTERCLOCKWISE)
atarm2 = Motor(Port.F)
drive_base = DriveBase(left_motor, right_motor, wheel_diameter=56, axle_track=112)
color_sensor = ColorSensor(Port.C)

drive_base.settings(300, 500, 300, 200)
Color.ORANGE = Color(10, 100, 100)
Color.MAGENTA = Color(321, 100, 86)

"""

def extract_main_function(content):
    lines = content.split('\n')
    main_content = []
    in_main_function = False
    main_indent = 0
    is_async = False
    
    for line in lines:
        stripped_line = line.strip()
        
        # Find the start of main function
        if stripped_line.startswith('def main') or stripped_line.startswith('async def main'):
            in_main_function = True
            is_async = stripped_line.startswith('async def main')
            continue
        
        if in_main_function:
            # If we hit another function or class definition at the same level, we're done
            if stripped_line and not line.startswith(' ') and not line.startswith('\t'):
                if stripped_line.startswith('def ') or stripped_line.startswith('class '):
                    break
            
            # Skip the first line after def main() if it's empty
            if not stripped_line and not main_content:
                continue
                
            # If this is the first content line, determine the indent level
            if main_content == [] and stripped_line:
                main_indent = len(line) - len(line.lstrip())
            
            # Remove the main function's indentation
            if line.strip():  # Don't process empty lines
                if len(line) - len(line.lstrip()) >= main_indent:
                    main_content.append(line[main_indent:])
                else:
                    main_content.append(line)
            else:
                main_content.append('')  # Keep empty lines
    
    return '\n'.join(main_content), is_async

# Clear the main.py file and write the required imports
with open("main.py", 'w') as required_imports:
    required_imports.write(imports)

function_calls = []

# Define colors properly - one per script
colors = [
    'Color.ORANGE', 'Color.GREEN', 'Color.BLACK', 'Color.WHITE', 
    'Color.YELLOW', 'Color.BLUE', 'Color.MAGENTA', 'Color.RED', 'Color.BROWN'
]

# Process each script file and create individual functions
for i, f_name in enumerate(script_names):
    try:
        with open(f_name, 'r') as f:
            content = f.read()
            # Extract only the main function content
            main_function_content, is_async = extract_main_function(content)
            
            if main_function_content.strip():  # Only proceed if it found main function content
                func_name = f_name.replace('.py', '').replace('-', '_')
                
                func_def = f"\n{'async ' if is_async else ''}def {func_name}():\n"
                
                indented_content = '\n'.join(['    ' + line if line.strip() else line for line in main_function_content.split('\n')])
                func_def += indented_content + "\n"
                
                with open("main.py", 'a') as m:
                    m.write(func_def)
                
                # Assign one color per script
                color_condition = colors[i % len(colors)]
                function_calls.append({
                    'name': func_name,
                    'is_async': is_async,
                    'color': color_condition,
                    'filename': f_name
                })
                
            else:
                print(f"Warning: No main() function found in {f_name}")
    except FileNotFoundError:
        print(f"Warning: File {f_name} not found")

# Write the main function that checks colors and calls appropriate functions
with open("main.py", 'a') as m:
    m.write("\nasync def main():\n")
    
    for func_info in function_calls:
        m.write(f"    if color_sensor.color() == {func_info['color']}:\n")
        
        if func_info['is_async']:
            m.write(f"        await {func_info['name']}()\n")
        else:
            m.write(f"        {func_info['name']}()\n")
        m.write("        return  # Exit after running one function\n")
    
    # Add a default case
    m.write("    # Default case - no matching color detected\n")
    m.write("    print(f'Detected color: {color_sensor.color()}')\n")

# Write the main loop
with open("main.py", 'a') as m:
    m.write("\n# Main execution loop\n")
    m.write("while True:\n")
    m.write("    run_task(main())\n")
    m.write("    wait(100)\n")

print("Script merger completed successfully!")
print("Functions created:")
for func_info in function_calls:
    print(f"  - {func_info['name']}() triggered by {func_info['color']} (from {func_info['filename']})")