import serial
import time
import sys
import serial.tools.list_ports

# --- CONFIGURATION ---
# Replace this with your actual port name!
# Windows: 'COM3', 'COM4', etc.
# Mac/Linux: '/dev/tty.usbmodem...' or '/dev/ttyUSB0'
PORT_NAME = "COM14" 
BAUD_RATE = 921600

def find_serial_port():
    """Helper to list ports if you don't know the name."""
    ports = serial.tools.list_ports.comports()
    print("Available Ports:")
    for p in ports:
        print(f" - {p.device}: {p.description}")
    return ports

def connect():
    """Establishes the serial connection."""
    try:
        ser = serial.Serial(PORT_NAME, BAUD_RATE, timeout=2)
        print(f"Connected to {PORT_NAME} at {BAUD_RATE} baud.")
        
        # DTR/RTS can sometimes reset ESP32 boards. 
        # We wait 2 seconds to let the Master Node boot up.
        time.sleep(2) 
        return ser
    except serial.SerialException as e:
        print(f"Error connecting to port: {e}")
        print("Did you close the Arduino Serial Monitor? (Only one app can use the port at a time)")
        sys.exit(1)

def discover_nodes(ser):
    """Sends 'a' and waits for the Master to report the highest address."""
    print("Enumerating nodes...", end="", flush=True)
    
    # 1. Send the 'a' command (Address Assignment)
    # We send it as a raw byte (ASCII 97)
    ser.write(b'a')
    
    # 2. Wait for response (1 byte)
    response = ser.read(1)
    
    if len(response) == 0:
        print("\nNo response! Is the Master board powered and flashed?")
        return 0
    
    # 3. Decode the byte
    # The Master sends its own address (e.g., 2).
    # Since addresses are 0-indexed (0, 1, 2), the count is Address + 1.
    max_id = int.from_bytes(response, 'big')
    count = max_id + 1
    
    print(f" Done!")
    print(f"Found {count} nodes (Addresses 0 to {max_id}).")
    return count

def set_led(ser, address, bitmask):
    """
    Turns on specific LEDs on a specific node.
    Protocol: ['l'] [ADDRESS] [BITMASK]
    """
    packet = bytearray([ord('l'), address, bitmask])
    ser.write(packet)

def set_all(ser, bitmask):
    """
    Turns on specific LEDs on ALL nodes at once.
    Protocol: ['L'] [BITMASK]
    """
    packet = bytearray([ord('L'), bitmask])
    ser.write(packet)

# --- MAIN EXECUTION ---
if __name__ == "__main__":
    # Optional: Uncomment to see list of ports
    # find_serial_port()
    
    ser = connect()
    
    # Step 1: Assign Addresses
    node_count = discover_nodes(ser)
    
    if node_count == 0:
        print("No nodes found. Exiting.")
        ser.close()
        sys.exit()

    print("Starting animation loop... (Press Ctrl+C to stop)")
    
    try:
        while True:
            # Animation 1: The "Wipe" (One by one)
            for i in [0,1,2]:
                # Turn ON LED on Node i (Mask 1 = 0b001 = LED On)
                set_led(ser, i, 1) 
                time.sleep(0.2)
                # Turn OFF LED on Node i (Mask 0 = All Off)
                set_led(ser, i, 0)
            
            time.sleep(0.2)
            set_all(ser, 1)
            time.sleep(.2)
            set_all(ser, 0)
            time.sleep(0.2) 
            
            for i in [2, 1, 0]:
                set_led(ser, i, 1) 
                time.sleep(0.2)
                # Turn OFF LED on Node i (Mask 0 = All Off)
                set_led(ser, i, 0)

            time.sleep(0.2)
            set_all(ser, 1)
            time.sleep(.2)
            set_all(ser, 0)
            time.sleep(0.2) 

            # # Animation 2: Global Flash
            # for _ in range(3):
            #     set_all(ser, 1) # All ON
            #     time.sleep(0.1)
            #     set_all(ser, 0) # All OFF
            #     time.sleep(0.1)
                
    except KeyboardInterrupt:
        print("\nStopping...")
        set_all(ser, 0) # Turn everything off before quitting
        ser.close()
        print("Closed.")