import serial
import time

def write_radio_address(com_port, new_address):
    # The baudrate configured in your initUsart0() is 57600
    BAUDRATE = 57600

    try:
        # Open the serial connection
        with serial.Serial(com_port, BAUDRATE, timeout=2) as ser:
            print(f"[*] Connected to {com_port} at {BAUDRATE} baud.")

            # --- DTR / RTS HARDWARE RESET SEQUENCE ---
            # Replicating "mode comX: dtr=on" behavior to force a clean hardware reset
            print("[*] Pulsing DTR/RTS to reset the robot...")
            ser.dtr = False
            ser.rts = False
            time.sleep(0.1)
            ser.dtr = True
            ser.rts = True
            
            # CRITICAL: After reset, the ATmega2560 stays in the bootloader for ~1.5 to 2 seconds.
            # We must wait for the bootloader to timeout so your main.c firmware starts running!
            print("[*] Waiting 2.0 seconds for main application (main.c) to start up...")
            time.sleep(2.0)
            
            # Flush any bootloader noise or residual bytes from the serial buffers
            ser.reset_input_buffer()
            ser.reset_output_buffer()
            # ----------------------------------------

            # Check if the address is valid (16-bit integer: 0 to 65535)
            if not (0 <= new_address <= 65535):
                print("[-] Error: the address must be between 0 and 65535.")
                return

            # LSB and MSB extraction (little-endian, as required by usart.c)
            lsb = new_address & 0xFF
            msb = (new_address >> 8) & 0xFF

            print(f"[*] Preparing to write address: {new_address} (Hex: 0x{new_address:04X})")
            print(f"[*] Sending LSB: 0x{lsb:02X}, MSB: 0x{msb:02X}")

            # 1. Send 0x02 to select menuChoice = 2
            ser.write(bytes([0x02]))
            time.sleep(0.1)

            # 2. Send LSB and MSB sequentially
            ser.write(bytes([lsb, msb]))

            # 3. Wait for the confirmation (0xAA) from the robot
            print("[*] Data sent. Waiting for write confirmation from the robot...")
            response = ser.read(1)

            if response:
                if response[0] == 0xAA:
                    print("[+] SUCCESS: The robot confirmed the EEPROM write!")
                else:
                    print(f"[-] ERROR: Unexpected response received from the robot: 0x{response[0]:02X}")
            else:
                print("[-] TIMEOUT ERROR: No response. Verify that:")
                print("    1. The robot is turned on and connected to the correct port.")
                print("    2. The robot's HARDWARE SELECTOR is set to number 6.")

    except serial.SerialException as e:
        print(f"[-] Serial connection error: {e}")
    except Exception as e:
        print(f"[-] Unexpected error: {e}")

if __name__ == '__main__':
    # You can change these values according to your needs
    PORT = '/dev/ttyUSB0'          # Your COM port on Windows
    NEW_ADDRESS = 5026     # Insert the new address (e.g., 4015)

    print("--- Elisa3 EEPROM Write Utility ---")
    print("!!! WARNING: The robot selector MUST be set to number 6 !!!")
    input("Press ENTER to proceed...")
    
    write_radio_address(PORT, NEW_ADDRESS)