import os, sys, time
import usb.core
import usb.backend.libusb1
import libusb_package
import serial.tools.list_ports

SERIAL_VID = 0x1209
SERIAL_PID = 0xABD1
BOOT_VID = 0x29F1
BOOT_PID = 0x0230

def get_usb_backend():
    """Helper to load the libusb backend for pyusb."""
    return usb.backend.libusb1.get_backend(
        find_library=libusb_package.find_library
    )

def find_serial_port(vid, pid):
    """
    Returns the port name (e.g., COM3) if exactly one device matches.
    Returns None if 0 matches.
    Raises Exception if > 1 matches.
    """
    matches = []
    ports = serial.tools.list_ports.comports()

    for port in ports:
        if port.vid == vid and port.pid == pid:
            matches.append(port)

    count = len(matches)

    if count == 0:
        return None
    elif count > 1:
        msg = (f"CRITICAL: Found {count} Serial devices with VID:0x{vid:04X} PID:0x{pid:04X}. "
               f"Expected exactly 1. Ports: {[m.device for m in matches]}")
        raise Exception(msg)
    else:
        # Exactly one
        return matches[0].device

def count_bootloader_devices(vid, pid):
    """
    Returns the number of USB devices matching the Bootloader VID/PID.
    """
    backend = get_usb_backend()
    if not backend:
        raise Exception("libusb backend not found. Cannot scan USB devices.")

    # list(devices) is required because usb.core.find returns a generator
    devices = list(usb.core.find(find_all=True, backend=backend, idVendor=vid, idProduct=pid))
    return len(devices)

def enter_bootloader_mode(port):
    """
    Opens the serial port and performs the signaling to jump to bootloader.
    """
    print(f"Attempting to switch device on {port} to Bootloader mode...")
    ser = None
    try:
        ser = serial.Serial(
            port=port,
            baudrate=300,
            bytesize=serial.FIVEBITS,
            parity=serial.PARITY_MARK,
            stopbits=serial.STOPBITS_TWO,
            timeout=1,
            xonxoff=False,
            rtscts=False,
            dsrdtr=False,
        )

        # Specific signaling sequence
        ser.send_break()
        ser.dtr = True
        time.sleep(0.05)
        ser.rts = True
        time.sleep(0.05)

        print("  Signal sent. Closing port.")

    except serial.SerialException as e:
        print(f"  Warning: Serial communication error (device might have reset already): {e}")
    finally:
        if ser and ser.is_open:
            try:
                ser.close()
            except:
                pass

def wait_for_bootloader_appearance(vid, pid, timeout):
    """
    Polls for the bootloader device to appear.
    Raises exception if multiple appear.
    """
    print("Waiting for device to enumerate in Bootloader mode...")
    start_time = time.time()

    while time.time() - start_time < timeout:
        count = count_bootloader_devices(vid, pid)

        if count == 1:
            print("  Success! Device detected in Bootloader mode.")
            return True
        elif count > 1:
            raise Exception(f"Error: Multiple ({count}) Bootloader devices detected during wait!")

        time.sleep(0.25)

    return False

def ensure_device_in_bootloader(timeout):
    print("--- Starting Device Check ---")

    try:
        serial_port = find_serial_port(SERIAL_VID, SERIAL_PID)
    except Exception as e:
        print(f"Aborting: {e}")
        sys.exit(1)

    if serial_port:
        print(f"Found Serial Device at {serial_port}.")

        enter_bootloader_mode(serial_port)

        if wait_for_bootloader_appearance(BOOT_VID, BOOT_PID, timeout):
            print("Operation Complete: Device is ready in Bootloader mode.")
        else:
            print("Error: Timed out waiting for device to appear in Bootloader mode.")
            sys.exit(1)

    else:
        print("No Serial Device found. Checking for existing Bootloader devices...")

        boot_count = count_bootloader_devices(BOOT_VID, BOOT_PID)

        if boot_count == 0:
            print("Error: No devices found in Serial(CanMV) OR Bootloader mode.")
            sys.exit(1)
        elif boot_count > 1:
            print(f"Error: Found {boot_count} devices in Bootloader mode. Expected exactly 1.")
            sys.exit(1)
        else:
            # boot_count == 1
            print("Success: Found exactly one device already in Bootloader mode.")
