#!/usr/bin/env python3
"""
Verify KB2040 composite USB exposes both HID and Vendor/WebUSB interfaces.

Usage:
  python3 tests/test_kb2040_dual.py

Dependencies:
  pip install pyusb

Notes:
- May require sudo unless udev permissions are configured.
- Defaults to VID:PID 239a:8104 per user report.
"""

import argparse
import sys
import time

try:
    import usb.core
    import usb.util
except ImportError:
    print("ERROR: Missing dependency. Install with: pip install pyusb")
    sys.exit(2)


def find_device(vid: int, pid: int):
    return usb.core.find(idVendor=vid, idProduct=pid)


def claim_vendor_interface(dev, vendor_itf):
    itf_num = vendor_itf.bInterfaceNumber
    try:
        if dev.is_kernel_driver_active(itf_num):
            dev.detach_kernel_driver(itf_num)
    except (NotImplementedError, usb.core.USBError):
        pass
    usb.util.claim_interface(dev, itf_num)


def release_vendor_interface(dev, vendor_itf):
    itf_num = vendor_itf.bInterfaceNumber
    try:
        usb.util.release_interface(dev, itf_num)
    except usb.core.USBError:
        pass


def find_endpoints(vendor_itf):
    ep_out = None
    ep_in = None
    for ep in vendor_itf:
        direction = usb.util.endpoint_direction(ep.bEndpointAddress)
        if direction == usb.util.ENDPOINT_OUT:
            ep_out = ep
        elif direction == usb.util.ENDPOINT_IN:
            ep_in = ep
    return ep_out, ep_in


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--vid", default="239a", help="USB Vendor ID in hex (default: 239a)")
    parser.add_argument("--pid", default="8104", help="USB Product ID in hex (default: 8104)")
    parser.add_argument("--payload", default="hello", help="Payload for vendor echo test")
    parser.add_argument("--timeout-ms", type=int, default=2000, help="Read timeout in ms")
    args = parser.parse_args()

    vid = int(args.vid, 16)
    pid = int(args.pid, 16)

    passed = 0
    failed = 0

    dev = find_device(vid, pid)
    if dev is None:
        print(f"[FAIL] Device {vid:04x}:{pid:04x} not found")
        return 1

    print(f"[PASS] Device {dev.idVendor:04x}:{dev.idProduct:04x} found")
    passed += 1

    try:
        dev.set_configuration()
    except usb.core.USBError:
        # Usually already configured; continue.
        pass

    cfg = dev.get_active_configuration()
    hid_itf = None
    vendor_itf = None

    for itf in cfg:
        if itf.bInterfaceClass == 0x03:  # HID
            hid_itf = itf
        elif itf.bInterfaceClass == 0xFF:  # Vendor-specific
            vendor_itf = itf

    if hid_itf is not None:
        print(f"[PASS] HID interface found (if={hid_itf.bInterfaceNumber})")
        passed += 1
    else:
        print("[FAIL] HID interface not found")
        failed += 1

    if vendor_itf is not None:
        print(f"[PASS] Vendor/WebUSB interface found (if={vendor_itf.bInterfaceNumber})")
        passed += 1
    else:
        print("[FAIL] Vendor/WebUSB interface not found")
        failed += 1

    if vendor_itf is not None:
        try:
            claim_vendor_interface(dev, vendor_itf)
            ep_out, ep_in = find_endpoints(vendor_itf)
            if ep_out is None or ep_in is None:
                print("[FAIL] Could not find bulk IN/OUT endpoints on vendor interface")
                failed += 1
            else:
                payload = args.payload.encode("utf-8")
                ep_out.write(payload)
                time.sleep(0.05)
                data = ep_in.read(ep_in.wMaxPacketSize, timeout=args.timeout_ms)
                got = bytes(data)
                expected = b"KB2040:" + payload
                if got == expected:
                    print(f"[PASS] Vendor echo OK: sent={payload!r} got={got!r}")
                    passed += 1
                else:
                    print(f"[FAIL] Vendor echo mismatch: expected={expected!r} got={got!r}")
                    failed += 1
        except usb.core.USBTimeoutError:
            print("[FAIL] Vendor read timed out")
            failed += 1
        except usb.core.USBError as e:
            print(f"[FAIL] USB error during vendor test: {e}")
            failed += 1
        finally:
            release_vendor_interface(dev, vendor_itf)

    print(f"\nSummary: {passed} passed, {failed} failed")
    if failed == 0:
        print("Dual interface support is working (HID + Vendor/WebUSB).")
        return 0
    return 1


if __name__ == "__main__":
    sys.exit(main())