import os
import subprocess


def modify_python_code_for_pygbag(input_code):
    # Step 3: Replace "def main():" with "async def main():"
    #modified_code = input_code.replace("def main():", "async def main():")
    modified_code = input_code.replace('if __name__ == "__main__":', "async def tmp():")

    modified_code += \
"""\n
async def main():
    game = Game()
    pygame.init()
    global event
    running = True
    while running:
        event = pygame.event.poll()
        running = game.run(event)
        pygame.display.flip()
        await asyncio.sleep(0.00)
    pygame.quit()

asyncio.run(main())
"""

    modified_code = modified_code.replace("self.PIPE_GAP = int(15)", "self.PIPE_GAP = int(100)")
        
    # Step 2: Add "await asyncio.sleep(0)" after "clock.tick(state_manager.fps)"
    #if "clock.tick(state_manager.fps)" in modified_code:
    #    modified_code = modified_code.replace(
    #        "clock.tick(state_manager.fps)",
    #        "clock.tick(state_manager.fps)\n        await asyncio.sleep(0)",
    #    )
    
    #
    # insert "await asyncio.sleep(0)" right after pygame.display.flip()
    #if "pygame.display.flip()" in modified_code:
    #    modified_code = modified_code.replace(
    #        "running = game.run(event)\n        pygame.display.flip()",
    #        "running = game.run(event)\n        pygame.display.flip()\n        await asyncio.sleep(0)",
    #    )

    return "import asyncio\n" + modified_code

def move_files_to_new_directories(old_directory, new_directory):
    # move all the files *_debug*py files under
    idx = 20
    for filename in os.listdir(old_directory):
        if "one_shot_debug" in filename or "sequential_debug" in filename:
            print(filename)
            filepath = os.path.join(old_directory, filename)
            code = open(filepath, "r").read()
            assert type(code) == str
            code = modify_python_code_for_pygbag(code)

            new_baseline_dir = os.path.join(new_directory, f"{idx}")
            idx += 1
            os.makedirs(new_baseline_dir, exist_ok=True)

            # save code to new_baseline_dir/main.py
            with open(os.path.join(new_baseline_dir, "main.py"), "w") as f:
                f.write(code)


if __name__ == "__main__":
    import sys
    game = sys.argv[1]

    dir_path = f"{game}/"
    with open(dir_path + "/final.py", "r") as f:
        code = f.read()
        code = modify_python_code_for_pygbag(code)
    os.makedirs(dir_path, exist_ok=True)
    with open(dir_path + "/main.py", "w") as f:
        f.write(code)

    # Change the current working directory to the target directory
    print(f"Compiling in {dir_path}...")

    # Execute the pygbag compile command
    result = subprocess.run(["pygbag", "--build",  dir_path], capture_output=True, text=True)

    # Check if the compile was successful
    if result.returncode == 0:
        print("Compilation successful.")
    else:
        print(f"Compilation failed with error: {result.stderr}")