import os, sys, time, struct
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
_bootloader_serial_handle = None

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 find_cdc_control_interface(dev):
    """
    Finds the CDC ACM control interface used for USB serial control transfers.
    """
    try:
        cfg = dev.get_active_configuration()
        for intf in cfg:
            if intf.bInterfaceClass == 0x02 and intf.bInterfaceSubClass == 0x02:
                return intf.bInterfaceNumber
    except usb.core.USBError:
        pass

    try:
        for cfg in dev:
            for intf in cfg:
                if intf.bInterfaceClass == 0x02 and intf.bInterfaceSubClass == 0x02:
                    return intf.bInterfaceNumber
    except usb.core.USBError:
        pass

    return 0

def detach_kernel_drivers(dev):
    """
    Detaches Linux CDC ACM kernel drivers so control transfers can be sent.
    """
    try:
        for cfg in dev:
            for intf in cfg:
                ifnum = intf.bInterfaceNumber
                try:
                    if dev.is_kernel_driver_active(ifnum):
                        dev.detach_kernel_driver(ifnum)
                except (usb.core.USBError, NotImplementedError):
                    pass
    except usb.core.USBError:
        pass

def enter_bootloader_mode_linux_usb(port):
    """
    Switches the Linux CDC ACM device into bootloader mode using the SDK's
    USB control-transfer sequence.
    """
    backend = get_usb_backend()
    if not backend:
        print("  Warning: pyusb/libusb backend not available.")
        return False

    print(f"Attempting to switch device on {port} to Bootloader mode via USB control transfer...")

    try:
        dev = usb.core.find(
            idVendor=SERIAL_VID,
            idProduct=SERIAL_PID,
            backend=backend
        )
        if dev is None:
            print(f"  Warning: USB device {SERIAL_VID:04X}:{SERIAL_PID:04X} not found via pyusb.")
            return False

        detach_kernel_drivers(dev)
        control_interface = find_cdc_control_interface(dev)

        # CDC SET_LINE_CODING: 300 baud, 2 stop bits, mark parity, 5 data bits.
        line_coding = struct.pack("<I", 300) + bytes([2, 3, 5])
        dev.ctrl_transfer(
            bmRequestType=0x21,
            bRequest=0x20,
            wValue=0,
            wIndex=control_interface,
            data_or_wLength=line_coding,
            timeout=1000
        )

        # CDC SEND_BREAK.
        dev.ctrl_transfer(
            bmRequestType=0x21,
            bRequest=0x23,
            wValue=0xFFFF,
            wIndex=control_interface,
            data_or_wLength=None,
            timeout=1000
        )

        # CDC SET_CONTROL_LINE_STATE: DTR=1 and RTS=1.
        dev.ctrl_transfer(
            bmRequestType=0x21,
            bRequest=0x22,
            wValue=0x0003,
            wIndex=control_interface,
            data_or_wLength=None,
            timeout=1000
        )

        print("  Signal sent via USB control transfer.")
        return True
    except usb.core.USBError as e:
        print(f"  Warning: USB control transfer failed: {e}")
        return False

def enter_bootloader_mode_serial(port, linux_sequence=False):
    """
    Opens the serial port and performs the signaling to jump to bootloader.
    """
    global _bootloader_serial_handle

    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,
        )

        if linux_sequence:
            time.sleep(0.2)
            ser.dtr = False
            ser.rts = False
            time.sleep(0.1)

            ser.send_break()
            time.sleep(0.2)

            ser.dtr = True
            time.sleep(0.1)

            ser.rts = True
            time.sleep(0.2)

            _bootloader_serial_handle = ser
            ser = None
            print("  Signal sent. Keeping port open while device reboots.")
            return

        # Existing Windows-compatible 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 close_bootloader_serial_handle():
    global _bootloader_serial_handle

    if _bootloader_serial_handle and _bootloader_serial_handle.is_open:
        try:
            _bootloader_serial_handle.close()
        except:
            pass
    _bootloader_serial_handle = None

def enter_bootloader_mode(port):
    """
    Switches the connected board into bootloader mode.
    """
    if sys.platform.startswith("linux"):
        if enter_bootloader_mode_linux_usb(port):
            return
        print("  Falling back to serial bootloader signaling.")
        enter_bootloader_mode_serial(port, linux_sequence=True)
        return

    enter_bootloader_mode_serial(port)

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

        try:
            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)
        finally:
            close_bootloader_serial_handle()

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