import os, sys, time
from pathlib import Path
import serial.tools.list_ports

try:
    import usb.core
    import usb.backend.libusb1
    import libusb_package
except ImportError:
    usb = None
    libusb_package = None

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

def get_usb_backend():
    """Helper to load the libusb backend for pyusb."""
    if usb is None or libusb_package is None:
        return None

    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.
    """
    sysfs_count = count_usb_devices_sysfs(vid, pid)
    if sysfs_count is not None:
        return sysfs_count

    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 count_usb_devices_sysfs(vid, pid):
    """
    Counts USB devices from Linux sysfs. This avoids missing the short boot ROM
    window when libusb cannot enumerate the device due to udev permissions.
    """
    if not sys.platform.startswith("linux"):
        return None

    usb_devices_path = Path("/sys/bus/usb/devices")
    if not usb_devices_path.exists():
        return None

    expected_vid = f"{vid:04x}"
    expected_pid = f"{pid:04x}"
    count = 0

    for device_path in usb_devices_path.iterdir():
        id_vendor_path = device_path / "idVendor"
        id_product_path = device_path / "idProduct"
        if not id_vendor_path.exists() or not id_product_path.exists():
            continue

        try:
            current_vid = id_vendor_path.read_text(encoding="ascii").strip().lower()
            current_pid = id_product_path.read_text(encoding="ascii").strip().lower()
        except OSError:
            continue

        if current_vid == expected_vid and current_pid == expected_pid:
            count += 1

    return count

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.")
