#!/usr/bin/env python3

import subprocess
import serial
import sys
import glob

def upload_src(port):
    subprocess.call(["mpremote", "connect", port, "fs", "cp", "-r", ".", ":"], cwd="src")

def upload_lib(port):
    subprocess.call(["mpremote", "connect", port, "fs", "cp", "-r", "lib", ":"])

def list_files(port):
    subprocess.call(["mpremote", "connect", port, "fs", "ls", ":"])

def reset_device(port):
    with serial.Serial(port, baudrate=115200) as ser:
        ser.write(b'\x03\x03\x04')

def follow_serial(port):
    with serial.Serial(port, baudrate=115200) as ser:
        while True:
            if ser.in_waiting > 0:
                data = ser.readline().decode('utf-8').strip()
                print(data)

def clean_settings(port):
    subprocess.call(["mpremote", "connect", port, "fs", "rm", ":/settings.json"])

def flash_firmware(port):
    # Find MicroPython firmware binary
    firmware_files = glob.glob("*.bin")
    if not firmware_files:
        print("Error: No .bin firmware file found in current directory")
        print("Please download MicroPython firmware and place it in the project directory")
        sys.exit(1)
    
    firmware = firmware_files[0]
    if len(firmware_files) > 1:
        print(f"Warning: Multiple .bin files found, using: {firmware}")
    
    print(f"Flashing MicroPython firmware: {firmware}")
    print("Erasing flash...")
    subprocess.call(["esptool.py", "--port", port, "erase_flash"])
    
    print(f"Writing firmware to flash...")
    subprocess.call([
        "esptool.py",
        "--port", port,
        "--baud", "460800",
        "write_flash", "0",
        firmware
    ])
    print("Flash complete!")

def main():
    port = "/dev/ttyACM0"
    commands = []
    i = 1
    
    # Parse arguments manually to preserve order
    while i < len(sys.argv):
        arg = sys.argv[i]
        if arg in ["-p", "--port"]:
            if i + 1 < len(sys.argv):
                port = sys.argv[i + 1]
                i += 2
            else:
                print(f"Error: {arg} requires a port argument")
                sys.exit(1)
        elif arg in ["-s", "--src"]:
            commands.append(("src", upload_src))
            i += 1
        elif arg in ["-r", "--reset"]:
            commands.append(("reset", reset_device))
            i += 1
        elif arg in ["-f", "--follow"]:
            commands.append(("follow", follow_serial))
            i += 1
        elif arg == "--lib":
            commands.append(("lib", upload_lib))
            i += 1
        elif arg == "--ls":
            commands.append(("ls", list_files))
            i += 1
        elif arg == "--clean":
            commands.append(("clean", clean_settings))
            i += 1
        elif arg == "--flash":
            commands.append(("flash", flash_firmware))
            i += 1
        elif arg in ["-h", "--help"]:
            print("LED Driver development tools")
            print("\nUsage:")
            print("  ./dev.py [-p PORT] [FLAGS...]")
            print("\nFlags:")
            print("  -p, --port PORT    Serial port (default: /dev/ttyACM0)")
            print("  -s, --src          Upload src directory")
            print("  -r, --reset        Reset device")
            print("  -f, --follow       Follow serial output")
            print("  --lib              Upload lib directory")
            print("  --ls               List files on device")
            print("  --clean            Remove settings.json from device")
            print("  --flash            Flash MicroPython firmware")
            print("\nExamples:")
            print("  ./dev.py -p /dev/ttyACM0 -s -r -f")
            print("  ./dev.py --flash -s -r")
            sys.exit(0)
        else:
            print(f"Error: Unknown argument: {arg}")
            print("Use -h or --help for usage information")
            sys.exit(1)
    
    # Execute commands in the order they were given
    if not commands:
        print("No commands specified. Use -h or --help for usage information.")
        sys.exit(1)
    
    for cmd_name, cmd_func in commands:
        if cmd_name == "reset":
            print("Resetting device...")
        elif cmd_name == "follow":
            print("Following serial output (Ctrl+C to exit)...")
        elif cmd_name == "flash":
            pass  # flash_firmware prints its own messages
        else:
            print(f"{cmd_name.capitalize()}...")
        cmd_func(port)

if __name__ == "__main__":
    main()