#include <device/def_device_op.h>
#include <device/device.h>
#include <device/usb/usb.h>
#include <device/usb/xhci.h>
#include <irq/irq.h>
#include <libk/list.h>
#include <libk/std.h>
#include <libk/string.h>
#include <limine/requests.h>
#include <mm/malloc.h>
#include <mm/pmm.h>
#include <proc/proc.h>
#include <proc/reschedule.h>
#include <proc/suspension_q.h>
#include <status.h>
#include <sys/debug.h>
#include <sys/spin_lock.h>
#include <sys/stall.h>

/* REF:
 * https://www.intel.com/content/dam/www/public/us/en/documents/technical-specifications/extensible-host-controler-interface-usb-xhci.pdf
 */

/* clang-format off */

/* capability registers */
#define XHCI_CAPLENGTH  0x00
#define XHCI_RSVD       0x01
#define XHCI_HCIVERSION 0x02
#define XHCI_HCSPARAMS1 0x04
#define XHCI_HCSPARAMS2 0x08
#define XHCI_HCSPARAMS3 0x0C
#define XHCI_HCCPARAMS1 0x10
#define XHCI_DBOFF      0x14
#define XHCI_RTSOFF     0x18
#define XHCI_HCCPARAMS2 0x1C

/* operational registers */
#define XHCI_USBCMD     0x00
#define XHCI_USBSTS     0x04
#define XHCI_PAGESIZE   0x08
#define XHCI_DNCTRL     0x14
#define XHCI_CRCR       0x18
#define XHCI_DCBAAP     0x30
#define XHCI_CONFIG     0x38

/* port registers */
#define XHCI_PORTSC     0x00
#define XHCI_PORTPMSC   0x04
#define XHCI_PORTLI     0x08

/* runtime registers */
#define XHCI_MFINDEX    0x00
/* + IRQ sets (0x20) */
#define XHCI_IMAN       0x00
#define XHCI_IMOD       0x04
#define XHCI_ERSTSZ     0x08
#define XHCI_ERSTBA     0x10
#define XHCI_ERDP       0x18

/* event types */
#define XHCI_TRB_NORMAL         1
#define XHCI_TRB_SETUP_STAGE    2
#define XHCI_TRB_DATA_STAGE     3
#define XHCI_TRB_STATUS_STAGE   4
#define XHCI_TRB_ISOCH          5
#define XHCI_TRB_LINK           6
#define XHCI_TRB_EVENT_DATA     7
#define XHCI_TRB_NOOP           8
#define XHCI_TRB_SLOT_ENAB_CMD  9
#define XHCI_TRB_SLOT_DISB_CMD  10
#define XHCI_TRB_ADDR_DEV_CMD   11
#define XHCI_TRB_CFG_ENDP_CMD   12
#define XHCI_TRB_EVAL_CTX_CMD   13
#define XHCI_TRB_RESET_ENDP_CMD 14
#define XHCI_TRB_STOP_ENDP_CMD  15
#define XHCI_TRB_SET_DRDQP_CMD  16
#define XHCI_TRB_RESET_DEV_CMD  17
#define XHCI_TRB_FORCE_EVT_CMD  18
#define XHCI_TRB_NEGO_BANDW_CMD 19
#define XHCI_TRB_SET_LTV_CMD    20
#define XHCI_TRB_PORT_BANDW_CMD 21
#define XHCI_TRB_FORCE_HEADER   22
#define XHCI_TRB_NOOP_CMD       23
#define XHCI_TRB_GET_EXTPRP_CMD 24
#define XHCI_TRB_SET_EXTPRP_CMD 25
#define XHCI_TRB_TRANSFER_EVENT 32
#define XHCI_TRB_CMD_CMPL_EVENT 33
#define XHCI_TRB_PORT_STS_CHNG  34
#define XHCI_TRB_BANDW_RQ_EVENT 35
#define XHCI_TRB_DOORBELL_EVENT 36
#define XHCI_TRB_HOST_CTRL_EVNT 37
#define XHCI_TRB_DEV_NOTIF_EVNT 38
#define XHCI_TRB_MFINDEX_WRAP   39

/* extended capabilities */
#define XHCI_EXTCAP_USB_LEGACY_SUPPORT 1
#define XHCI_EXTCAP_SUPPORTED_PROTOCOL 2
#define XHCI_EXTCAP_EXT_POWER_MNGMNT   3
#define XHCI_EXTCAP_IO_VIRT            4
#define XHCI_EXTCAP_MSG_INTR           5
#define XHCI_EXTCAP_LOCAL_MEMORY       6
#define XHCI_EXTCAP_USB_DEBUG          7
#define XHCI_EXTCAP_EXT_MSG_INTR       8
#define XHCI_EXTCAP_VENDOR_DEFINED     9

/* clang-format on */

static void xhci_write8 (uintptr_t base, uint32_t reg, uint8_t value) {
  *(volatile uint8_t*)(base + reg) = value;
}

static void xhci_write16 (uintptr_t base, uint32_t reg, uint16_t value) {
  *(volatile uint16_t*)(base + reg) = value;
}

static void xhci_write32 (uintptr_t base, uint32_t reg, uint32_t value) {
  *(volatile uint32_t*)(base + reg) = value;
}

static uint8_t xhci_read8 (uintptr_t base, uint32_t reg) {
  return *(volatile uint8_t*)(base + reg);
}

static uint16_t xhci_read16 (uintptr_t base, uint32_t reg) {
  return *(volatile uint16_t*)(base + reg);
}

static uint32_t xhci_read32 (uintptr_t base, uint32_t reg) {
  return *(volatile uint32_t*)(base + reg);
}

static uint32_t xhci_portsc_read (struct xhci* xhci, uint8_t port) {
  return xhci_read32 (xhci->xhci_oper_base, 1024 + (16 * port));
}

static void xhci_portsc_write (struct xhci* xhci, uint8_t port, uint32_t value) {
  xhci_write32 (xhci->xhci_oper_base, 1024 + (16 * port), value);
}

static void xhci_port_reset (struct xhci* xhci, uint8_t port) {
  struct xhci_port* xhci_port;

  list_find (struct xhci_port, xhci->xhci_ports, xhci_port, port_value, port, ports_link);

  uint32_t portsc = xhci_portsc_read (xhci, port);

  /* check Port Power */
  if (!(portsc & (1 << 9))) {
    portsc |= (1 << 9);
    xhci_portsc_write (xhci, port, portsc);

    stall_ms (50);

    portsc = xhci_portsc_read (xhci, port);
  }

  /* connect status change, enable port, port reset change */
  portsc |= (1 << 17) | (1 << 18) | (1 << 21);

  xhci_portsc_write (xhci, port, portsc);

  if (xhci_port->type == XHCI_PORT_USB3) {
    /* warm port reset */
    portsc |= (1 << 31);
  } else if (xhci_port->type == XHCI_PORT_USB2) {
    /* port reset */
    portsc |= (1 << 4);
  }

  xhci_portsc_write (xhci, port, portsc);

  stall_ms (500);

  /* warm port reset change, port reset change , connect status change, port enabled change */
  portsc |= (1 << 21) | (1 << 19) | (1 << 17) | (1 << 18);
  /* port enabled */
  portsc &= ~(1 << 1);

  xhci_portsc_write (xhci, port, portsc);

  stall_ms (100);
}

static void xhci_create_pdevice (struct xhci* xhci, uint8_t port) {
  struct xhci_pdevice* pdevice;

  list_find (struct xhci_pdevice, xhci->xhci_pdevices, pdevice, port_value, port, pdevices_link);

  if (pdevice != NULL)
    return;

  pdevice = malloc (sizeof (*pdevice));
  memset (pdevice, 0, sizeof (*pdevice));
  pdevice->port_value = port;
  pdevice->slot_id = -1;

  list_append (xhci->xhci_pdevices, &pdevice->pdevices_link);
}

static void xhci_delete_pdevice (struct xhci* xhci, uint8_t port) {
  struct xhci_pdevice* pdevice;

  list_find (struct xhci_pdevice, xhci->xhci_pdevices, pdevice, port_value, port, pdevices_link);

  list_remove (xhci->xhci_pdevices, &pdevice->pdevices_link);

  if (pdevice->endpoint0_ring_phys != 0)
    pmm_free (pdevice->endpoint0_ring_phys, 1);

  free (pdevice);
}

static void xhci_event_dispatch (struct xhci* xhci, struct xhci_trb* event, uint8_t type) {
  switch (type) {
  case XHCI_TRB_CMD_CMPL_EVENT: {
    uint8_t cmpl_code = (event->status >> 24) & 0xFF;
    uint8_t slot_id = (event->ctrl >> 24) & 0xFF;

    DEBUG ("cmd completion: code=%u,slot=%u\n", cmpl_code, slot_id);

    xhci->last_slot_id = slot_id;
    xhci->last_cmpl_code = cmpl_code;
  } break;
  case XHCI_TRB_PORT_STS_CHNG: {
    uint8_t port = ((event->param >> 24) & 0xFF) - 1;

    uint32_t portsc = xhci_portsc_read (xhci, port);
    DEBUG ("Status change on port %u: 0x%08x\n", port, portsc);

    /* ack. PED + PR */
    portsc &= ~((1 << 1) | (1 << 4));

    xhci_portsc_write (xhci, port, portsc);

    if ((portsc & (1 << 17))) {
      if ((portsc & (1 << 0))) {
        DEBUG ("Device attached to port %u!\n", port);

        xhci_create_pdevice (xhci, port);
      } else {
        DEBUG ("Device detached from port %u!\n", port);

        xhci_delete_pdevice (xhci, port);
      }
    }

    if ((portsc & (1 << 21)))
      DEBUG ("Port %u reset done\n", port);
  } break;
  case XHCI_TRB_TRANSFER_EVENT: {
    uint8_t cmpl_code = (event->status >> 24) & 0xFF;
    uint8_t slot_id = (event->ctrl >> 24) & 0xFF;
    uint8_t endpoint_id = (event->ctrl >> 16) & 0x1F;

    DEBUG ("transfer completion: code=%u,slot=%u,endpoint_id=%u\n", cmpl_code, slot_id,
           endpoint_id);
    xhci->last_cmpl_code = cmpl_code;
  } break;
  default:
    DEBUG ("Unhandled event type %u at %u\n", type, xhci->event_ring_idx);
    break;
  }
}

static void xhci_poll_events (struct xhci* xhci) {
  uintptr_t ir_base = xhci->xhci_runtime_base + 0x20;

  bool serviced = false;
  for (;;) {
    struct xhci_trb* event = &xhci->event_ring[xhci->event_ring_idx];

    if ((event->ctrl & 0x1) != xhci->event_cycle_bit) {
      break;
    }

    serviced = true;

    uint8_t type = (event->ctrl >> 10) & 0x3F;

    xhci_event_dispatch (xhci, event, type);

    xhci->event_ring_idx++;

    if (xhci->event_ring_idx >= xhci->event_ring_size) {
      xhci->event_ring_idx = 0;
      xhci->event_cycle_bit ^= 1;
    }
  }

  if (serviced) {
    uintptr_t dequeue_ptr =
        xhci->event_ring_phys + (xhci->event_ring_idx * sizeof (struct xhci_trb));

    xhci_write32 (ir_base, XHCI_ERDP, (uint32_t)dequeue_ptr | (1 << 3));
    xhci_write32 (ir_base, XHCI_ERDP + 4, (uint32_t)(dequeue_ptr >> 32));

    atomic_store (&xhci->pending, false);
  }
}

static void xhci_irq (void* arg, void* regs, bool user, struct reschedule_ctx* rctx) {
  (void)user, (void)regs, (void)rctx;

  uint64_t fd;

  struct xhci* xhci = arg;

  spin_lock (&xhci->device->lock, &fd);

  uintptr_t ir_base = xhci->xhci_runtime_base + 0x20;

  /* ack */
  xhci_write32 (ir_base, XHCI_IMAN, xhci_read32 (ir_base, XHCI_IMAN) | (1 << 0));
  xhci_write32 (xhci->xhci_oper_base, XHCI_USBSTS, (1 << 3));

  xhci_poll_events (xhci);

  spin_unlock (&xhci->device->lock, fd);
}

static void xhci_endpoint0_put_trb (struct xhci_pdevice* pdevice, struct xhci_trb trb) {
  if (pdevice->endpoint0_ring_idx == pdevice->endpoint0_ring_size - 1) {
    struct xhci_trb* link = &pdevice->endpoint0_ring[pdevice->endpoint0_ring_idx];
    link->param = pdevice->endpoint0_ring_phys;
    link->status = 0;
    link->ctrl = (XHCI_TRB_LINK << 10) | (1 << 1) | pdevice->endpoint0_cycle_bit;

    pdevice->endpoint0_ring_idx = 0;
    pdevice->endpoint0_cycle_bit ^= 1;
  }

  pdevice->endpoint0_ring[pdevice->endpoint0_ring_idx] = trb;
  pdevice->endpoint0_ring[pdevice->endpoint0_ring_idx].ctrl =
      (pdevice->endpoint0_ring[pdevice->endpoint0_ring_idx].ctrl & ~1u) |
      pdevice->endpoint0_cycle_bit;

  pdevice->endpoint0_ring_idx++;
}

static bool xhci_endpoint0_ctrl_in (struct xhci* xhci, struct xhci_pdevice* pdevice,
                                    uint8_t request_type, uint8_t request, uint16_t value,
                                    uint16_t index, uintptr_t data_phys, uint16_t length,
                                    uint64_t* lockflags) {
  uint64_t setup = ((uint64_t)length << 48) | ((uint64_t)index << 32) | ((uint64_t)value << 16) |
      ((uint64_t)request << 8) | ((uint64_t)request_type);

  struct xhci_trb trb;

  xhci->last_cmpl_code = 0;
  atomic_store (&xhci->pending, true);

  /* setup stage */
  memset (&trb, 0, sizeof (trb));
  trb.param = setup;
  trb.status = 8;
  trb.ctrl = (XHCI_TRB_SETUP_STAGE << 10) | (1 << 6) | (1 << 4);
  xhci_endpoint0_put_trb (pdevice, trb);

  /* data stage */
  memset (&trb, 0, sizeof (trb));
  trb.param = data_phys;
  trb.status = length;
  trb.ctrl = (XHCI_TRB_DATA_STAGE << 10) | (1 << 16) | (1 << 4);
  xhci_endpoint0_put_trb (pdevice, trb);

  /* status stage */
  memset (&trb, 0, sizeof (trb));
  trb.param = 0;
  trb.status = 0;
  trb.ctrl = (XHCI_TRB_STATUS_STAGE << 10) | (1 << 5);
  xhci_endpoint0_put_trb (pdevice, trb);

  xhci_write32 (xhci->xhci_doorbell_base, pdevice->slot_id * 4, 1);

  int timeout = 100;

  spin_unlock (&xhci->device->lock, *lockflags);

  while (atomic_load (&xhci->pending) && --timeout > 0)
    stall_ms (10);

  spin_lock (&xhci->device->lock, lockflags);

  if (timeout == 0) {
    DEBUG ("timed out\n");
    return false;
  }

  return xhci->last_cmpl_code == 1;
}

static void xhci_send_cmd (struct xhci* xhci, uint64_t param, uint32_t status, uint32_t ctrl,
                           uint64_t* lockflags) {
  if (xhci->irqs_support) {
    if (xhci->cmd_ring_idx == (xhci->cmd_ring_size - 1)) {
      struct xhci_trb* link = &xhci->cmd_ring[xhci->cmd_ring_idx];
      link->param = xhci->cmd_ring_phys;
      link->status = 0;
      link->ctrl = (XHCI_TRB_LINK << 10) | (1 << 1) | xhci->cmd_cycle_bit;

      xhci->cmd_ring_idx = 0;
      xhci->cmd_cycle_bit ^= 1;
    }

    struct xhci_trb* trb = &xhci->cmd_ring[xhci->cmd_ring_idx];
    trb->param = param;
    trb->status = status;
    trb->ctrl = (ctrl & ~0x1) | xhci->cmd_cycle_bit;

    xhci->cmd_ring_idx++;

    atomic_store (&xhci->pending, true);

    xhci_write32 (xhci->xhci_doorbell_base, 0, 0);

    int timeout = 100;

    spin_unlock (&xhci->device->lock, *lockflags);

    while (atomic_load (&xhci->pending) && --timeout > 0)
      stall_ms (10);

    spin_lock (&xhci->device->lock, lockflags);

    if (timeout == 0)
      DEBUG ("timed out\n");
  } else {
    while (atomic_load (&xhci->pending)) {
      xhci_poll_events (xhci);
      spin_lock_relax ();
    }
  }
}

static void xhci_bios_handover (struct xhci* xhci) {
  uint32_t hccparams1 = xhci_read32 (xhci->xhci_mmio_base, XHCI_HCCPARAMS1);
  uint32_t ext_offset = (hccparams1 >> 16) << 2;

  if (ext_offset == 0)
    return;

  while (ext_offset) {
    uintptr_t cap_ptr = xhci->xhci_mmio_base + ext_offset;
    uint32_t cap = xhci_read32 (cap_ptr, 0);
    uint8_t cap_id = cap & 0xFF;

    if (cap_id == XHCI_EXTCAP_USB_LEGACY_SUPPORT) {
      /* Make or break on real hardware. We need to take over ownership from the BIOS. */
      if (cap & (1 << 16)) {
        DEBUG ("BIOS owns XHCI, requesting handover!\n");

        xhci_write8 (cap_ptr, 3, 1);

        /* Wait for BIOS Semaphore owned bit to be cleared */
        int timeout = 1000;
        while (--timeout > 0) {
          uint32_t val = xhci_read32 (cap_ptr, 0);
          if (!(val & (1 << 16)) && (val & (1 << 24)))
            break;

          stall_ms (100);
        }

        DEBUG ("XHCI Handover OK\n");
      }

      break;
    }

    uint8_t next = (cap >> 8) & 0xFF;
    if (!next)
      break;

    ext_offset += (next << 2);
  }
}

static void xhci_reset_ports (struct xhci* xhci) {
  uint32_t hccparams1 = xhci_read32 (xhci->xhci_mmio_base, XHCI_HCCPARAMS1);
  uint32_t ext_offset = (hccparams1 >> 16) << 2;

  if (ext_offset == 0)
    return;

  while (ext_offset) {
    uintptr_t cap_ptr = xhci->xhci_mmio_base + ext_offset;
    uint32_t cap = xhci_read32 (cap_ptr, 0);
    uint8_t cap_id = cap & 0xFF;
    uint8_t minor = (cap >> 16) & 0xFF;
    uint8_t major = (cap >> 24) & 0xFF;

    if (cap_id == XHCI_EXTCAP_SUPPORTED_PROTOCOL) {
      uint32_t dword2 = xhci_read32 (cap_ptr, 8);
      uint8_t port_off = dword2 & 0xFF;
      uint8_t port_count = (dword2 >> 8) & 0xFF;

      uint8_t first_port = port_off - 1;
      uint8_t last_port = first_port + port_count - 1;

      for (uint8_t port = first_port; port <= last_port; port++) {
        struct xhci_port* xhci_port = malloc (sizeof (*xhci_port));
        memset (xhci_port, 0, sizeof (*xhci_port));

        if (major == 3)
          xhci_port->type = XHCI_PORT_USB3;
        else
          xhci_port->type = XHCI_PORT_USB2;

        xhci_port->port_value = port;

        list_append (xhci->xhci_ports, &xhci_port->ports_link);

        uint32_t portsc = xhci_portsc_read (xhci, port);

        if ((portsc & (1 << 0))) {
          DEBUG ("Device connected. resetting\n");
          xhci_port_reset (xhci, port);
          xhci_create_pdevice (xhci, port);
        }

        DEBUG ("PORT %u: USB %u.%u\n", port, major, minor);
      }
    }

    uint8_t next = (cap >> 8) & 0xFF;
    if (!next)
      break;

    ext_offset += (next << 2);
  }
}

static void xhci_pdevice_setup_addressing (struct xhci* xhci, struct xhci_pdevice* pdevice,
                                           uint64_t* lockflags) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;

  struct xhci_port* xhci_port;
  list_find (struct xhci_port, xhci->xhci_ports, xhci_port, port_value, pdevice->port_value,
             ports_link);

  uint32_t portsc = xhci_portsc_read (xhci, pdevice->port_value);

  uint32_t speed = (portsc >> 10) & 0x0F;

  xhci->last_cmpl_code = 0;
  xhci_send_cmd (xhci, 0, 0, XHCI_TRB_SLOT_ENAB_CMD << 10, lockflags);

  if (xhci->last_cmpl_code != 1) {
    DEBUG ("Enable slot failed\n");
    return;
  }

  pdevice->slot_id = xhci->last_slot_id;

  uintptr_t out_ctx_phys = pmm_alloc (1);
  void* out_ctx_virt = (void*)(out_ctx_phys + (uintptr_t)hhdm->offset);
  memset (out_ctx_virt, 0, PAGE_SIZE);
  xhci->xhci_dcbaa[pdevice->slot_id] = out_ctx_phys;

  pdevice->endpoint0_ring_phys = pmm_alloc (1);
  pdevice->endpoint0_ring =
      (struct xhci_trb*)(pdevice->endpoint0_ring_phys + (uintptr_t)hhdm->offset);
  memset (pdevice->endpoint0_ring, 0, PAGE_SIZE);
  pdevice->endpoint0_ring_size = PAGE_SIZE / sizeof (struct xhci_trb);
  pdevice->endpoint0_ring_idx = 0;
  pdevice->endpoint0_cycle_bit = 1;

  uintptr_t input_ctx_phys = pmm_alloc (1);
  void* input_ctx_virt = (void*)(input_ctx_phys + (uintptr_t)hhdm->offset);
  memset (input_ctx_virt, 0, PAGE_SIZE);

  if (xhci->xhci_ctx_size == 64) {
    struct xhci_input_ctx64* ctx64 = input_ctx_virt;

    size_t max_packet_size = (xhci_port->type == XHCI_PORT_USB3) ? 512 : 64;

    /* Add slot and endpoint 0 */
    ctx64->ctrl.dw[1] = (1 << 0) | (1 << 1);

    ctx64->slot.dw[0] = (1 << 27) | (speed << 20);
    ctx64->slot.dw[1] = ((pdevice->port_value + 1) << 16);

    ctx64->endpoints[0].dw[0] = 0;
    ctx64->endpoints[0].dw[1] = 3 | (4 << 3) | (max_packet_size << 16);
    ctx64->endpoints[0].dw[2] = (uint32_t)pdevice->endpoint0_ring_phys | (1 << 0);
    ctx64->endpoints[0].dw[3] = (uint32_t)(pdevice->endpoint0_ring_phys >> 32);
  } else {
    struct xhci_input_ctx32* ctx32 = input_ctx_virt;

    size_t max_packet_size = (xhci_port->type == XHCI_PORT_USB3) ? 512 : 64;

    ctx32->ctrl.dw[1] = (1 << 0) | (1 << 1);

    ctx32->slot.dw[0] = (1 << 27) | (speed << 20);
    ctx32->slot.dw[1] = ((pdevice->port_value + 1) << 16);

    ctx32->endpoints[0].dw[1] = 3 | (4 << 3) | (max_packet_size << 16);
    ctx32->endpoints[0].dw[2] = (uint32_t)pdevice->endpoint0_ring_phys | (1 << 0);
    ctx32->endpoints[0].dw[3] = (uint32_t)(pdevice->endpoint0_ring_phys >> 32);
  }

  xhci->last_cmpl_code = 0;
  uint32_t ctrl = (pdevice->slot_id << 24) | (XHCI_TRB_ADDR_DEV_CMD << 10);
  xhci_send_cmd (xhci, input_ctx_phys, 0, ctrl, lockflags);

  stall_ms (500);

  if (xhci->last_cmpl_code != 1) {
    DEBUG ("Failed to address device. port = %u, slot = %u\n", pdevice->port_value,
           pdevice->slot_id);
  } else {
    DEBUG ("Device on port %u addressed on slot %u!\n", pdevice->port_value, pdevice->slot_id);
  }

  pmm_free (input_ctx_phys, 1);
}

static bool xhci_pdevice_setup_get_info (struct xhci* xhci, struct xhci_pdevice* pdevice,
                                         uint64_t* lockflags) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;

  uintptr_t response_buf_phys = pmm_alloc (1);

  void* response_buf = (void*)(response_buf_phys + (uintptr_t)hhdm->offset);
  memset (response_buf, 0, PAGE_SIZE);

  bool ok = xhci_endpoint0_ctrl_in (xhci, pdevice, 0x80, 6, (1 << 8), 0, response_buf_phys,
                                    sizeof (struct usb_device_desc), lockflags);

  if (!ok) {
    pmm_free (response_buf_phys, 1);
    return false;
  }

  struct usb_device_desc* usb_desc = response_buf;

  DEBUG ("USB device info: vndr=%04x prdct=%04x class=%02x subclass=%02x mps=%u cfgs=%u\n",
         usb_desc->vendor_id, usb_desc->product_id, usb_desc->dev_class, usb_desc->dev_subclass,
         usb_desc->max_packet_size, usb_desc->num_configs);

  pmm_free (response_buf_phys, 1);
  return true;
}

static void xhci_poll_setup_devices (struct xhci* xhci, uint64_t* lockflags) {
  struct list_node_link *pdevice_link, *tmp_pdevice_link;

  list_foreach (xhci->xhci_pdevices, pdevice_link, tmp_pdevice_link) {
    struct xhci_pdevice* pdevice = list_entry (pdevice_link, struct xhci_pdevice, pdevices_link);

    /* slot assigned */
    if (pdevice->slot_id != -1)
      continue;

    xhci_pdevice_setup_addressing (xhci, pdevice, lockflags);

    xhci_pdevice_setup_get_info (xhci, pdevice, lockflags);
  }
}

DEFINE_DEVICE_OP (xhci_poll_driver) {
  struct xhci* xhci = device->udata;

  xhci_poll_setup_devices (xhci, lockflags);

  return ST_OK;
}

DEFINE_DEVICE_INIT (xhci_init) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;

  struct xhci* xhci = malloc (sizeof (*xhci));

  if (xhci == NULL)
    return false;

  struct xhci_init* init = arg;

  memset (xhci, 0, sizeof (*xhci));
  xhci->device = device;
  xhci->xhci_mmio_base = init->xhci_mmio_base;
  xhci->irqs_support = init->irqs_support;
  xhci->irq = init->irq;

  device->udata = xhci;

  uint32_t usbcmd, config, cap_length;

  cap_length = xhci_read8 (xhci->xhci_mmio_base, XHCI_CAPLENGTH);

  xhci->xhci_oper_base = xhci->xhci_mmio_base + cap_length;

  uint32_t rtsoff = xhci_read32 (xhci->xhci_mmio_base, XHCI_RTSOFF);
  xhci->xhci_runtime_base = xhci->xhci_mmio_base + rtsoff;

  uint32_t dboff = xhci_read32 (xhci->xhci_mmio_base, XHCI_DBOFF);
  xhci->xhci_doorbell_base = xhci->xhci_mmio_base + dboff;

  uint32_t hcsparams2 = xhci_read32 (xhci->xhci_mmio_base, XHCI_HCSPARAMS2);
  xhci->max_scratchpad = (((hcsparams2 >> 21) & 0x1F) << 5) | ((hcsparams2 >> 27) & 0x1F);

  uint32_t hccparams1 = xhci_read32 (xhci->xhci_mmio_base, XHCI_HCCPARAMS1);
  xhci->xhci_ctx_size = (hccparams1 & (1 << 2)) ? 64 : 32;

  DEBUG ("starting init sequence\n");

  /* stop running / clear Run/Stop bit */
  usbcmd = xhci_read32 (xhci->xhci_oper_base, XHCI_USBCMD);
  usbcmd &= ~(1 << 0);
  xhci_write32 (xhci->xhci_oper_base, XHCI_USBCMD, usbcmd);

  stall_ms (1000);

  xhci_bios_handover (xhci);

  /* reset controller */
  usbcmd = xhci_read32 (xhci->xhci_oper_base, XHCI_USBCMD);
  usbcmd |= (1 << 1);
  xhci_write32 (xhci->xhci_oper_base, XHCI_USBCMD, usbcmd);

  while (xhci_read32 (xhci->xhci_oper_base, XHCI_USBSTS) & (1 << 11))
    spin_lock_relax ();

  DEBUG ("controller reset\n");

  xhci->max_slots = xhci_read32 (xhci->xhci_mmio_base, XHCI_HCSPARAMS1) & 0xFF;
  DEBUG ("max_slots=%u\n", xhci->max_slots);

  /* enable device notifications */
  xhci_write32 (xhci->xhci_oper_base, XHCI_DNCTRL, 0xFFFF);

  /* enable slots */
  config = xhci_read32 (xhci->xhci_oper_base, XHCI_CONFIG);
  xhci_write32 (xhci->xhci_oper_base, XHCI_CONFIG, (config & ~0xFF) | (xhci->max_slots & 0xFF));

  /* Prepare DCBAA */
  xhci->xhci_dcbaa_phys = pmm_alloc (1);
  xhci->xhci_dcbaa = (uintptr_t*)(xhci->xhci_dcbaa_phys + (uintptr_t)hhdm->offset);
  memset (xhci->xhci_dcbaa, 0, PAGE_SIZE);

  if (xhci->max_scratchpad > 0) {
    xhci->scratchpads_phys = pmm_alloc (1);
    xhci->scratchpads = (uintptr_t*)(xhci->scratchpads_phys + (uintptr_t)hhdm->offset);
    memset (xhci->scratchpads, 0, PAGE_SIZE);

    for (size_t sp = 0; sp < xhci->max_scratchpad; sp++) {
      xhci->scratchpads[sp] = pmm_alloc (1);
    }

    xhci->xhci_dcbaa[0] = xhci->scratchpads_phys;
  }

  xhci_write32 (xhci->xhci_oper_base, XHCI_DCBAAP, (uint32_t)xhci->xhci_dcbaa_phys);
  xhci_write32 (xhci->xhci_oper_base, XHCI_DCBAAP + 4, (uint32_t)(xhci->xhci_dcbaa_phys >> 32));

  xhci->cmd_ring_phys = pmm_alloc (1);
  xhci->cmd_ring = (struct xhci_trb*)(xhci->cmd_ring_phys + (uintptr_t)hhdm->offset);
  memset (xhci->cmd_ring, 0, PAGE_SIZE);
  xhci->cmd_ring_size = PAGE_SIZE / sizeof (struct xhci_trb);
  xhci->cmd_ring_idx = 0;
  xhci->cmd_cycle_bit = 1;

  uint64_t crcr = xhci->cmd_ring_phys | xhci->cmd_cycle_bit;
  xhci_write32 (xhci->xhci_oper_base, XHCI_CRCR, (uint32_t)crcr);
  xhci_write32 (xhci->xhci_oper_base, XHCI_CRCR + 4, (uint32_t)(crcr >> 32));

  xhci->event_ring_phys = pmm_alloc (1);
  xhci->event_ring = (struct xhci_trb*)(xhci->event_ring_phys + (uintptr_t)hhdm->offset);
  memset (xhci->event_ring, 0, PAGE_SIZE);
  xhci->event_ring_size = PAGE_SIZE / sizeof (struct xhci_trb);
  xhci->event_ring_idx = 0;
  xhci->event_cycle_bit = 1;

  xhci->erst_phys = pmm_alloc (1);
  xhci->erst = (struct xhci_erst_entry*)(xhci->erst_phys + (uintptr_t)hhdm->offset);
  memset (xhci->erst, 0, PAGE_SIZE);
  xhci->erst[0].ptr = xhci->event_ring_phys;
  xhci->erst[0].size = xhci->event_ring_size;
  xhci->erst[0]._rsvd = 0;

  uintptr_t ir_base = xhci->xhci_runtime_base + 0x20;
  xhci_write32 (ir_base, XHCI_ERSTSZ, 1);
  xhci_write32 (ir_base, XHCI_ERSTBA, (uint32_t)xhci->erst_phys);
  xhci_write32 (ir_base, XHCI_ERSTBA + 4, (uint32_t)(xhci->erst_phys >> 32));

  xhci_write32 (ir_base, XHCI_ERDP, (uint32_t)xhci->event_ring_phys | (1 << 3));
  xhci_write32 (ir_base, XHCI_ERDP + 4, (uint32_t)(xhci->event_ring_phys >> 32));

  if (xhci->irqs_support) {
    /* enable interrupter */
    irq_attach (&xhci_irq, xhci, xhci->irq);
    xhci_write32 (ir_base, XHCI_IMAN, xhci_read32 (ir_base, XHCI_IMAN) | (1 << 1));
  }

  usbcmd = xhci_read32 (xhci->xhci_oper_base, XHCI_USBCMD);
  xhci_write32 (xhci->xhci_oper_base, XHCI_USBCMD, usbcmd | (1 << 0) | (1 << 2));

  stall_ms (500);

  xhci_reset_ports (xhci);

  return true;
}

DEFINE_DEVICE_FINI (xhci_fini) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;

  struct xhci* xhci = device->udata;

  if (xhci->max_scratchpad > 0) {
    uintptr_t scratchpads_phys = xhci->xhci_dcbaa[0];

    uintptr_t* scratchpads = (uintptr_t*)(scratchpads_phys + (uintptr_t)hhdm->offset);

    for (size_t i = 0; i < xhci->max_scratchpad; i++) {
      if (scratchpads[i] != 0)
        pmm_free (scratchpads[i], 1);
    }

    pmm_free (scratchpads_phys, 1);
  }

  for (size_t i = 0; i < PAGE_SIZE / sizeof (xhci->xhci_dcbaa[0]); i++) {
    uintptr_t out_ctx = xhci->xhci_dcbaa[i];

    if (out_ctx != 0)
      pmm_free (out_ctx, 1);
  }

  pmm_free (xhci->xhci_dcbaa_phys, 1);
  pmm_free (xhci->cmd_ring_phys, 1);
  pmm_free (xhci->event_ring_phys, 1);
  pmm_free (xhci->erst_phys, 1);

  struct list_node_link *port_link, *tmp_port_link;
  list_foreach (xhci->xhci_ports, port_link, tmp_port_link) {
    struct xhci_port* port = list_entry (port_link, struct xhci_port, ports_link);

    list_remove (xhci->xhci_ports, &port->ports_link);
    free (port);
  }

  irq_detach (xhci->irq);

  free (xhci);
}