#!/usr/bin/env python3
"""
Lua Game Uploader for RP2350
Rapidly upload and execute Lua games via USB serial for quick iteration.

Usage:
    python upload_game.py <lua_file> [serial_port]

Example:
    python upload_game.py my_game.lua /dev/ttyACM0
    python upload_game.py my_game.lua COM3
"""

import sys
import os
import base64
import serial
import time
from pathlib import Path

def find_serial_port():
    """Auto-detect the RP2350 serial port"""
    import serial.tools.list_ports

    # Look for Pico/RP2350 devices
    ports = list(serial.tools.list_ports.comports())
    for port in ports:
        # Common VID/PID for Raspberry Pi Pico
        if 'Pico' in port.description or \
           'RP2350' in port.description or \
           'RP2040' in port.description or \
           (port.vid == 0x2E8A):  # Raspberry Pi VID
            return port.device

    # Fallback: return first available port
    if ports:
        print(f"Warning: Could not find RP2350, using first available port: {ports[0].device}")
        return ports[0].device

    return None

def upload_game(lua_file, serial_port=None):
    """Upload a Lua game file to the RP2350 and execute it"""

    # Check if file exists
    if not os.path.exists(lua_file):
        print(f"Error: File '{lua_file}' not found")
        return False

    # Read the file
    with open(lua_file, 'rb') as f:
        file_data = f.read()

    file_size = len(file_data)
    filename = os.path.basename(lua_file)

    # Ensure filename has .lua extension
    if not filename.endswith('.lua'):
        filename += '.lua'

    print(f"Uploading {filename} ({file_size} bytes)...")

    # Auto-detect serial port if not provided
    if serial_port is None:
        serial_port = find_serial_port()
        if serial_port is None:
            print("Error: No serial port found. Please specify manually.")
            return False
        print(f"Using serial port: {serial_port}")

    try:
        # Open serial connection
        ser = serial.Serial(serial_port, 115200, timeout=5)
        time.sleep(0.1)  # Wait for connection to stabilize

        # Encode file data as base64
        base64_data = base64.b64encode(file_data).decode('ascii')

        # Send upload command
        command = f"UPLOAD {filename} {file_size}\n"
        print(f"Sending command: {command.strip()}")
        ser.write(command.encode('ascii'))

        # Send base64-encoded file data
        print("Sending file data...")
        ser.write(base64_data.encode('ascii'))
        ser.write(b'\n')

        # Send END marker
        ser.write(b'END\n')
        print("Upload complete, waiting for confirmation...")

        # Read response from device
        start_time = time.time()
        response_lines = []
        while time.time() - start_time < 10:  # 10 second timeout
            if ser.in_waiting > 0:
                line = ser.readline().decode('utf-8', errors='ignore').strip()
                if line:
                    print(f"  << {line}")
                    response_lines.append(line)

                    # Check for success
                    if line.startswith('OK'):
                        print(f"✓ File written successfully!")

                    if line.startswith('LAUNCHED'):
                        game_name = line.split(' ', 1)[1] if ' ' in line else filename
                        print(f"✓ Game '{game_name}' launched!")
                        ser.close()
                        return True

                    # Check for error
                    if line.startswith('ERROR'):
                        print(f"✗ Upload failed: {line}")
                        ser.close()
                        return False

        print("Warning: No response received from device")
        ser.close()
        return False

    except serial.SerialException as e:
        print(f"Serial error: {e}")
        return False
    except Exception as e:
        print(f"Unexpected error: {e}")
        return False

def main():
    if len(sys.argv) < 2:
        print(__doc__)
        print("\nAvailable serial ports:")
        try:
            import serial.tools.list_ports
            for port in serial.tools.list_ports.comports():
                print(f"  {port.device}: {port.description}")
        except ImportError:
            print("  (install pyserial to see available ports)")
        sys.exit(1)

    lua_file = sys.argv[1]
    serial_port = sys.argv[2] if len(sys.argv) > 2 else None

    success = upload_game(lua_file, serial_port)
    sys.exit(0 if success else 1)

if __name__ == '__main__':
    main()