import mm.mm as mm
import string
import drivers.serial.uart.serial as serial
import drivers.usb.uhci as uhci
import drivers.usb.pci as pci
import platform.pch.timer as timer
import viperlib


define USB_DESC_DEVICE = 1
define USB_DESC_CONFIG = 2
define USB_DESC_STRING = 3
define USB_DESC_INTERFACE = 4
define USB_DESC_ENDPOINT = 5

define USB_REQ_GET_STATUS = 0
define USB_REQ_CLEAR_FEATURE = 1
define USB_REQ_SET_FEATURE = 3
define USB_REQ_SET_ADDRESS = 5
define USB_REQ_GET_DESCRIPTOR = 6
define USB_REQ_SET_DESCRIPTOR = 7
define USB_REQ_GET_CONFIGURATION = 8
define USB_REQ_SET_CONFIGURATION = 9

define USB_REQ_SET_INTERFACE = 11

define HID_CLASS = 0x03
define HID_SUBCLASS_BOOT = 0x01
define HID_PROTO_KEYBOARD = 0x01
define HID_PROTO_MOUSE = 0x02

define HID_REQ_SET_PROTOCOL = 0x0B
define HID_REQ_SET_IDLE = 0x0A
define HID_REQ_GET_REPORT = 0x01

define USB_MAX_DEVICES = 8

class usb_dev_desc:
    bLength:         uint8
    bDescriptorType: uint8
    bcdUSB:          uint16
    bDeviceClass:    uint8
    bDeviceSubClass: uint8
    bDeviceProtocol: uint8
    bMaxPacketSize0: uint8
    idVendor:        uint16
    idProduct:       uint16
    bcdDevice:       uint16
    iManufacturer:   uint8
    iProduct:        uint8
    iSerialNumber:   uint8
    bNumConfigurations: uint8

class usb_config_desc:
    bLength:         uint8
    bDescriptorType: uint8
    wTotalLength:    uint16
    bNumInterfaces:  uint8
    bConfigurationValue: uint8
    iConfiguration:  uint8
    bmAttributes:    uint8
    bMaxPower:       uint8

class usb_iface_desc:
    bLength:            uint8
    bDescriptorType:    uint8
    bInterfaceNumber:   uint8
    bAlternateSetting:  uint8
    bNumEndpoints:      uint8
    bInterfaceClass:    uint8
    bInterfaceSubClass: uint8
    bInterfaceProtocol: uint8
    iInterface:         uint8

class usb_ep_desc:
    bLength:          uint8
    bDescriptorType:  uint8
    bEndpointAddress: uint8
    bmAttributes:     uint8
    wMaxPacketSize:   uint16
    bInterval:        uint8

class usb_setup_pkt:
    bmRequestType: uint8
    bRequest:      uint8
    wValue:        uint16
    wIndex:        uint16
    wLength:       uint16

class usb_device:
    addr:           uint8
    ls:             int
    max_pkt0:       uint8
    iface_class:    uint8
    iface_subclass: uint8
    iface_protocol: uint8
    iface_num:      uint8
    ep_in:          uint8
    ep_in_max:      uint16
    ep_in_interval: uint8
    ep_in_toggle:   uint8
    configured:     int

_devices: [USB_MAX_DEVICES]usb_device
_device_count: int = 0

def _setup_packet(bmRequestType: uint8, bRequest: uint8, wValue: uint16, wIndex: uint16, wLength: uint16, buf: usb_setup_pkt*):
    buf.bmRequestType = bmRequestType
    buf.bRequest = bRequest
    buf.wValue = wValue
    buf.wIndex = wIndex
    buf.wLength = wLength

def _send_control(dev_addr: uint8, ls: int, max_pkt: uint8, setup: usb_setup_pkt*, data_buf: rawptr, data_len: uint16, data_in: int) -> int:
    return uhci.control_transfer(dev_addr, 0, ls, max_pkt, setup, data_buf, data_len, data_in)

def get_descriptor(dev_addr: uint8, ls: int, max_pkt: uint8, desc_type: uint8, desc_idx: uint8, buf: rawptr, length: uint16) -> int:
    setup: usb_setup_pkt
    _setup_packet(0x80, USB_REQ_GET_DESCRIPTOR, (uint16(desc_type) << 8) | desc_idx, 0, length, &setup)
    return _send_control(dev_addr, ls, max_pkt, &setup, buf, length, 1)

def set_address(dev_addr: uint8, ls: int, max_pkt: uint8, new_addr: uint8) -> int:
    setup: usb_setup_pkt
    _setup_packet(0x00, USB_REQ_SET_ADDRESS, new_addr, 0, 0, &setup)
    return _send_control(dev_addr, ls, max_pkt, &setup, None, 0, 0)

def set_configuration(dev_addr: uint8, ls: int, max_pkt: uint8, config_val: uint8) -> int:
    setup: usb_setup_pkt
    _setup_packet(0x00, USB_REQ_SET_CONFIGURATION, config_val, 0, 0, &setup)
    return _send_control(dev_addr, ls, max_pkt, &setup, None, 0, 0)

def hid_set_protocol(dev_addr: uint8, ls: int, max_pkt: uint8, iface: uint8, protocol: uint8) -> int:
    setup: usb_setup_pkt
    _setup_packet(0x21, HID_REQ_SET_PROTOCOL, protocol, uint16(iface), 0, &setup)
    return _send_control(dev_addr, ls, max_pkt, &setup, None, 0, 0)

def hid_set_idle(dev_addr: uint8, ls: int, max_pkt: uint8, iface: uint8, duration: uint8) -> int:
    setup: usb_setup_pkt
    _setup_packet(0x21, HID_REQ_SET_IDLE, uint16(duration) << 8, uint16(iface), 0, &setup)
    return _send_control(dev_addr, ls, max_pkt, &setup, None, 0, 0)

def _parse_config(dev: usb_device*, config_buf: rawptr, total_len: uint16) -> int:
    offset: int = 0
    found_hid: int = 0
    # 每次解析都从干净状态开始：调用方可能在"描述符没读全"时重试，残留的 iface/ep 字段
    # 会让上一次的半截结果串到这一次（例如 iface_class=3 但 ep_in=0 的不一致组合）。
    dev.iface_class = 0
    dev.iface_subclass = 0
    dev.iface_protocol = 0
    dev.iface_num = 0
    dev.ep_in = 0
    dev.ep_in_max = 0
    dev.ep_in_interval = 0
    while offset < int(total_len):
        p: uint8* = config_buf + uint64(offset)
        desc_len: uint8 = *p
        p2: uint8* = config_buf + uint64(offset) + 1
        desc_type: uint8 = *p2
        if desc_len == 0: break
        # 边界：描述符整体须落在 config_buf（total_len 字节）内，且不得短于
        # 所访问的结构体（interface=9 / endpoint=7），否则越界读。
        if offset + int(desc_len) > int(total_len): break
        if desc_type == USB_DESC_INTERFACE and int(desc_len) >= 9:
            iface: usb_iface_desc* = usb_iface_desc*(config_buf + uint64(offset))
            if iface.bInterfaceClass == HID_CLASS:
                found_hid = 1
                dev.iface_class = HID_CLASS
                dev.iface_subclass = iface.bInterfaceSubClass
                dev.iface_protocol = iface.bInterfaceProtocol
                dev.iface_num = iface.bInterfaceNumber
        if desc_type == USB_DESC_ENDPOINT and found_hid and int(desc_len) >= 7:
            ep: usb_ep_desc* = usb_ep_desc*(config_buf + uint64(offset))
            if (ep.bEndpointAddress & 0x80) != 0:
                dev.ep_in = ep.bEndpointAddress & 0x0F
                dev.ep_in_max = ep.wMaxPacketSize
                dev.ep_in_interval = ep.bInterval
                found_hid = 0
        offset = offset + int(desc_len)
    # 返回实际走完的偏移：调用方据此判断描述符链是否完整（== total_len 才算读全）
    return offset

def enumerate_port(port: int) -> int:
    global _device_count
    if _device_count >= USB_MAX_DEVICES: return -1

    ls: int = uhci.port_is_low_speed(port)

    reset_ok: int = -1
    for reset_retry in range(3):
        if uhci.port_reset(port) == 0:
            reset_ok = 0
            break
        timer.timer_msleep(50)
    if reset_ok != 0:
        return -2

    timer.timer_msleep(10)

    dev_buf: rawptr = mm.malloc(usb_dev_desc.__sizeof__())
    if dev_buf is None: return -3
    string.memset(dev_buf, 0, usb_dev_desc.__sizeof__())

    result: int = get_descriptor(0, ls, 8, USB_DESC_DEVICE, 0, dev_buf, 8)
    retry: int
    for retry in range(3):
        if result == 0: break
        timer.timer_msleep(50)
        uhci.port_reset(port)
        timer.timer_msleep(10)
        result = get_descriptor(0, ls, 8, USB_DESC_DEVICE, 0, dev_buf, 8)
    if result != 0:
        mm.free(dev_buf)
        return -4

    desc: usb_dev_desc* = dev_buf
    max_pkt: uint8 = desc.bMaxPacketSize0
    if max_pkt < 8: max_pkt = 8

    new_addr: uint8 = uint8(_device_count + 1)
    result = set_address(0, ls, max_pkt, new_addr)
    for retry in range(3):
        if result == 0: break
        timer.timer_msleep(20)
        result = set_address(0, ls, max_pkt, new_addr)
    if result != 0:
        mm.free(dev_buf)
        return -5

    timer.timer_msleep(10)

    dev: usb_device* = &_devices[_device_count]
    dev.addr = new_addr
    dev.ls = ls
    dev.max_pkt0 = max_pkt
    dev.configured = 0
    dev.iface_class = 0
    dev.iface_subclass = 0
    dev.iface_protocol = 0
    dev.ep_in = 0
    dev.ep_in_max = 0
    dev.ep_in_interval = 0
    dev.ep_in_toggle = 0

    config_buf: rawptr = mm.malloc(256)
    if config_buf is None:
        mm.free(dev_buf)
        return -6
    string.memset(config_buf, 0, 256)

    result = get_descriptor(new_addr, ls, max_pkt, USB_DESC_CONFIG, 0, config_buf, 9)
    for retry in range(3):
        if result == 0: break
        timer.timer_msleep(50)
        result = get_descriptor(new_addr, ls, max_pkt, USB_DESC_CONFIG, 0, config_buf, 9)
    if result != 0:
        mm.free(config_buf)
        mm.free(dev_buf)
        return -7

    cfg: usb_config_desc* = config_buf
    total_len: uint16 = cfg.wTotalLength
    if total_len > 256: total_len = 256

    # 读完整配置描述符：不能只看控制传输是否"报成功"——短读时错误位可能全 0，缓冲区里只
    # 落下前一段（现场：USB 鼠标的 25 字节描述符只到前 9 字节），描述符链会在中途遇到
    # bLength==0 提前结束 ⇒ 接口描述符解析不到 ⇒ iface_class=0 ⇒ 驱动找不到设备。
    # 因此用"描述符链是否恰好走完 total_len"来判定，不满足就重读；重读仍不完整直接报错
    # （fail-fast，不带着半截描述符继续）。
    result = -1
    consumed: int = 0
    for retry in range(3):
        result = get_descriptor(new_addr, ls, max_pkt, USB_DESC_CONFIG, 0, config_buf, total_len)
        if result == 0:
            consumed = _parse_config(dev, config_buf, total_len)
            if consumed == int(total_len): break
        timer.timer_msleep(50)
    if result != 0 or consumed != int(total_len):
        serial.puts("[usb] config descriptor incomplete\n")
        mm.free(config_buf)
        mm.free(dev_buf)
        return -8

    result = set_configuration(new_addr, ls, max_pkt, cfg.bConfigurationValue)
    if result != 0:
        mm.free(config_buf)
        mm.free(dev_buf)
        return -9

    dev.configured = 1

    if dev.iface_class == HID_CLASS:
        hid_set_protocol(new_addr, ls, max_pkt, dev.iface_num, 0)
        hid_set_idle(new_addr, ls, max_pkt, dev.iface_num, 0)

    _device_count = _device_count + 1
    mm.free(config_buf)
    mm.free(dev_buf)
    return int(new_addr)

def init() -> int:
    global _device_count
    _device_count = 0

    io_base: uint16 = pci.find_uhci()
    if io_base == 0:
        serial.puts("[usb] UHCI controller not found\n")
        return -1

    buf: [64]char
    viperlib.snprintf(&buf, 64, "[usb] UHCI at I/O 0x%04X\n", uint32(io_base))
    serial.puts(buf)

    if uhci.init(io_base) != 0:
        serial.puts("[usb] UHCI init failed\n")
        return -2

    serial.puts("[usb] UHCI initialized\n")

    timer.timer_msleep(50)

    for port in range(2):
        if uhci.port_is_connected(port):
            viperlib.snprintf(&buf, 64, "[usb] port %d connected\n", uint32(port))
            serial.puts(buf)
            dev_id: int = enumerate_port(port)
            if dev_id > 0:
                dev: usb_device* = &_devices[dev_id - 1]
                viperlib.snprintf(&buf, 64, "[usb] device addr=%d class=%02X proto=%02X ep=%d\n",
                    uint32(dev.addr), uint32(dev.iface_class),
                    uint32(dev.iface_protocol), uint32(dev.ep_in))
                serial.puts(buf)
            else:
                viperlib.snprintf(&buf, 64, "[usb] enumerate port %d failed: %d\n", uint32(port), int32(dev_id))
                serial.puts(buf)

    return _device_count

def find_hid_device(protocol: uint8) -> usb_device*:
    i: int
    for i in range(_device_count):
        dev: usb_device* = &_devices[i]
        if dev.iface_class == HID_CLASS:
            if dev.iface_protocol == protocol:
                if dev.configured != 0:
                    return dev
    return None

def read_interrupt(dev: usb_device*, buf: rawptr) -> int:
    return 0