#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/align.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/memorybarrier.h>
#include <sys/spin_lock.h>
#include <sys/stall.h>

static uint8_t xhci_get_device_ctx_idx (struct usb_endpoint_desc* endpoint);

static void xhci_write8 (uintptr_t base, uint32_t reg, uint8_t value) {
  *(volatile uint8_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 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_create_usb_device (struct xhci* xhci, struct xhci_port* xhci_port) {
  struct xhci_usb_device* usb_device;

  list_find (struct xhci_usb_device, xhci->xhci_usb_devices, usb_device, xhci_port->port_value,
             xhci_port->port_value, usb_devices_link);

  if (usb_device != NULL)
    return;

  usb_device = malloc (sizeof (*usb_device));
  memset (usb_device, 0, sizeof (*usb_device));
  usb_device->xhci_port = xhci_port;
  usb_device->slot_id = -1;

  list_append (xhci->xhci_usb_devices, &usb_device->usb_devices_link);
}

static void xhci_delete_usb_device (struct xhci* xhci, struct xhci_port* xhci_port,
                                    struct proc* proc, struct reschedule_ctx* rctx) {
  struct xhci_usb_device* usb_device = NULL;

  list_find (struct xhci_usb_device, xhci->xhci_usb_devices, usb_device, xhci_port->port_value,
             xhci_port->port_value, usb_devices_link);

  device_delete (usb_device->device->key, proc, rctx);

  list_remove (xhci->xhci_usb_devices, &usb_device->usb_devices_link);

  if (usb_device->endpoint0_ring.phys != 0)
    pmm_free (usb_device->endpoint0_ring.phys, 1);

  pmm_free (xhci->xhci_dcbaa[usb_device->slot_id], 1);
  xhci->xhci_dcbaa[usb_device->slot_id] = 0;

  free (usb_device);
}

static void xhci_bios_handover (struct xhci* xhci) {
  uint32_t hccparams1 = xhci_read32 (xhci->xhci_mmio_base, XHCI_HCCPARAMS1);
  uint32_t ext_offset = (hccparams1 >> XHCI_HCCPARAMS1_XECP) << 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 << XHCI_USBLEGSUP_BIOS_SEMA)) {
        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 << XHCI_USBLEGSUP_BIOS_SEMA)) && (val & (1 << XHCI_USBLEGSUP_OS_SEMA)))
            break;

          stall_ms (3);
        }

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

      break;
    }

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

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

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 << XHCI_PORTSC_PP))) {
    portsc |= (1 << XHCI_PORTSC_PP);
    xhci_portsc_write (xhci, port, portsc);

    stall_ms (50);

    portsc = xhci_portsc_read (xhci, port);
  }

  portsc |= (1 << XHCI_PORTSC_CSC) | (1 << XHCI_PORTSC_PEC) | (1 << XHCI_PORTSC_PRC);

  xhci_portsc_write (xhci, port, portsc);

  if (xhci_port->type == XHCI_PORT_USB3) {
    portsc |= (1 << XHCI_PORTSC_WPR);
  } else if (xhci_port->type == XHCI_PORT_USB2) {
    portsc |= (1 << XHCI_PORTSC_PR);
  }

  xhci_portsc_write (xhci, port, portsc);

  stall_ms (250);

  portsc |= (1 << XHCI_PORTSC_PRC) |
      (1 << XHCI_PORTSC_WRC) |
      (1 << XHCI_PORTSC_CSC) |
      (1 << XHCI_PORTSC_PEC);
  portsc &= ~(1 << XHCI_PORTSC_PED);

  xhci_portsc_write (xhci, port, portsc);

  stall_ms (100);
}

static void xhci_reset_ports (struct xhci* xhci) {
  uint32_t hccparams1 = xhci_read32 (xhci->xhci_mmio_base, XHCI_HCCPARAMS1);
  uint32_t ext_offset = (hccparams1 >> XHCI_HCCPARAMS1_XECP) << 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_SUPPORTED_PROTOCOL) {
      uint8_t minor = (cap >> XHCI_SUPPROTO_DW0_MINOR_REV) & 0xFF;
      uint8_t major = (cap >> XHCI_SUPPROTO_DW0_MAJOR_REV) & 0xFF;

      uint32_t dword2 = xhci_read32 (cap_ptr, 8);
      uint8_t port_off = (dword2 >> XHCI_SUPPROTO_DW2_PORT_OFF) & 0xFF;
      uint8_t port_count = (dword2 >> XHCI_SUPPROTO_DW2_PORT_COUNT) & 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 << XHCI_PORTSC_CCS))) {
          DEBUG ("Device connected. resetting\n");
          xhci_port_reset (xhci, port);
          xhci_create_usb_device (xhci, xhci_port);
        }

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

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

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

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 >> XHCI_CCETRB_STS_CMPL_CODE) & 0xFF;
    uint8_t slot_id = (event->ctrl >> XHCI_CCETRB_CTRL_SLOT_ID) & 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 >> XHCI_PSCETRB_PARAM_PORT_ID) & 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 << XHCI_PORTSC_PED) | (1 << XHCI_PORTSC_PR));

    xhci_portsc_write (xhci, port, portsc);

    struct xhci_port_status_change* change = malloc (sizeof (*change));
    change->port = port;
    change->portsc = portsc;
    list_append (xhci->port_changes, &change->port_changes_link);
  } break;
  case XHCI_TRB_TRANSFER_EVENT: {
    uint8_t cmpl_code = (event->status >> XHCI_TETRB_STS_CMPL_CODE) & 0xFF;
    uint8_t slot_id = (event->ctrl >> XHCI_TETRB_CTRL_SLOT_ID) & 0xFF;
    uint8_t endpoint_id = (event->ctrl >> XHCI_TETRB_CTRL_ENDPOINT) & 0x1F;

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

    xhci->last_slot_id = slot_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.trbs[xhci->event_ring.idx];

    if ((event->ctrl & (1 << XHCI_GTRB_CYCLE_BIT)) != xhci->event_ring.cycle_bit) {
      break;
    }

    serviced = true;

    uint8_t type = (event->ctrl >> XHCI_GTRB_TRB_TYPE) & 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_ring.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_ring_put_trb (struct xhci_ring* ring, struct xhci_trb put_trb) {
  if (ring->idx == ring->size - 1) {
    struct xhci_trb* link = &ring->trbs[ring->idx];
    link->param = ring->phys;
    link->status = 0;

    write_memory_barrier ();

    link->ctrl =
        (XHCI_TRB_LINK << XHCI_LNKTRB_CTRL_TRB_TYPE) | (1 << XHCI_LNKTRB_CTRL_TC) | ring->cycle_bit;

    ring->idx = 0;
    ring->cycle_bit ^= 1;
  }

  struct xhci_trb* trb = &ring->trbs[ring->idx];
  trb->param = put_trb.param;
  trb->status = put_trb.status;

  write_memory_barrier ();

  trb->ctrl = (put_trb.ctrl & ~XHCI_LNKTRB_CTRL_TC) | ring->cycle_bit;

  ring->idx++;
}

static bool xhci_endpoint0_ctrl_in (struct xhci* xhci, struct xhci_usb_device* usb_device,
                                    uint8_t request_type, uint8_t request, uint16_t value,
                                    uint16_t index, uintptr_t data_phys, uint16_t length,
                                    uint64_t* lockflags) {
  /* clang-format off */
  uint64_t setup = ((uint64_t)length << XHCI_SSTRB_PARAM_WLENGTH)
                | ((uint64_t)index << XHCI_SSTRB_PARAM_WINDEX)
                | ((uint64_t)value << XHCI_SSTRB_PARAM_WVALUE)
                | ((uint64_t)request << XHCI_SSTRB_PARAM_BREQUEST)
                | ((uint64_t)request_type << XHCI_SSTRB_PARAM_BMREQUEST_TYPE);
  /* clang-format on */

  struct xhci_trb trb;

  /* setup stage */
  memset (&trb, 0, sizeof (trb));
  trb.param = setup;
  trb.status = 8;

  write_memory_barrier ();

  trb.ctrl = (XHCI_TRB_SETUP_STAGE << XHCI_SSTRB_CTRL_TRB_TYPE) |
      (1 << XHCI_SSTRB_CTRL_IDT) |
      (XHCI_SSTRB_TRT_IN_DATA_STAGE << XHCI_SSTRB_CTRL_TRT);
  xhci_ring_put_trb (&usb_device->endpoint0_ring, trb);

  /* data stage */
  memset (&trb, 0, sizeof (trb));
  trb.param = data_phys;
  trb.status = length;

  write_memory_barrier ();

  trb.ctrl = (XHCI_TRB_DATA_STAGE << XHCI_DSTRB_CTRL_TRB_TYPE) | (1 << XHCI_DSTRB_CTRL_DIR);
  xhci_ring_put_trb (&usb_device->endpoint0_ring, trb);

  /* status stage */
  memset (&trb, 0, sizeof (trb));
  trb.param = 0;
  trb.status = 0;

  write_memory_barrier ();

  trb.ctrl = (XHCI_TRB_STATUS_STAGE << XHCI_STSTRB_CTRL_TRB_TYPE) | (1 << XHCI_STSTRB_CTRL_IOC);
  xhci_ring_put_trb (&usb_device->endpoint0_ring, trb);

  atomic_store (&xhci->pending, true);

  int timeout = 100;

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

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

  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 bool xhci_endpoint0_ctrl_out (struct xhci* xhci, struct xhci_usb_device* usb_device,
                                     uint8_t request_type, uint8_t request, uint16_t value,
                                     uint16_t index, uintptr_t data_phys, uint16_t length,
                                     uint64_t* lockflags) {
  /* clang-format off */
  uint64_t setup = ((uint64_t)length << XHCI_SSTRB_PARAM_WLENGTH)
                | ((uint64_t)index << XHCI_SSTRB_PARAM_WINDEX)
                | ((uint64_t)value << XHCI_SSTRB_PARAM_WVALUE)
                | ((uint64_t)request << XHCI_SSTRB_PARAM_BREQUEST)
                | ((uint64_t)request_type << XHCI_SSTRB_PARAM_BMREQUEST_TYPE);
  /* clang-format on */

  struct xhci_trb trb;

  /* setup stage */
  memset (&trb, 0, sizeof (trb));
  trb.param = setup;
  trb.status = 8;

  write_memory_barrier ();

  if (length == 0) {
    trb.ctrl = (XHCI_TRB_SETUP_STAGE << XHCI_SSTRB_CTRL_TRB_TYPE) |
        (1 << XHCI_SSTRB_CTRL_IDT) |
        (XHCI_SSTRB_TRT_NO_DATA_STAGE << XHCI_SSTRB_CTRL_TRT);
    xhci_ring_put_trb (&usb_device->endpoint0_ring, trb);
  } else {
    trb.ctrl = (XHCI_TRB_SETUP_STAGE << XHCI_SSTRB_CTRL_TRB_TYPE) |
        (1 << XHCI_SSTRB_CTRL_IDT) |
        (XHCI_SSTRB_TRT_OUT_DATA_STAGE << XHCI_SSTRB_CTRL_TRT);
    xhci_ring_put_trb (&usb_device->endpoint0_ring, trb);

    /* data stage */
    memset (&trb, 0, sizeof (trb));
    trb.param = data_phys;
    trb.status = length;

    write_memory_barrier ();

    trb.ctrl = (XHCI_TRB_DATA_STAGE << XHCI_DSTRB_CTRL_TRB_TYPE);
    xhci_ring_put_trb (&usb_device->endpoint0_ring, trb);
  }

  /* status stage */
  memset (&trb, 0, sizeof (trb));
  trb.param = 0;
  trb.status = 0;

  write_memory_barrier ();

  trb.ctrl = (XHCI_TRB_STATUS_STAGE << XHCI_STSTRB_CTRL_TRB_TYPE) |
      (1 << XHCI_STSTRB_CTRL_IOC) |
      (1 << XHCI_STSTRB_CTRL_DIR);
  xhci_ring_put_trb (&usb_device->endpoint0_ring, trb);

  atomic_store (&xhci->pending, true);

  int timeout = 100;

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

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

  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->cmd_ring.idx == (xhci->cmd_ring.size - 1)) {
    struct xhci_trb* link = &xhci->cmd_ring.trbs[xhci->cmd_ring.idx];
    link->param = xhci->cmd_ring.phys;
    link->status = 0;

    write_memory_barrier ();

    link->ctrl = (XHCI_TRB_LINK << XHCI_LNKTRB_CTRL_TRB_TYPE) |
        (1 << XHCI_LNKTRB_CTRL_TC) |
        xhci->cmd_ring.cycle_bit;

    xhci->cmd_ring.idx = 0;
    xhci->cmd_ring.cycle_bit ^= 1;
  }

  struct xhci_trb* trb = &xhci->cmd_ring.trbs[xhci->cmd_ring.idx];
  trb->param = param;
  trb->status = status;

  write_memory_barrier ();

  trb->ctrl = (ctrl & ~XHCI_LNKTRB_CTRL_TC) | xhci->cmd_ring.cycle_bit;

  xhci->cmd_ring.idx++;

  atomic_store (&xhci->pending, true);

  int timeout = 100;

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

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

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

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

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

static uint8_t xhci_get_device_ctx_idx (struct usb_endpoint_desc* endpoint) {
  uint8_t endpoint_num = endpoint->endpoint_addr & 0x0F;
  bool in = (endpoint->endpoint_addr & 0x80) != 0;
  return (endpoint_num * 2) + (in ? 1 : 0);
}

static void xhci_usb_device_setup_init_endpoints (struct xhci* xhci,
                                                  struct xhci_usb_device* usb_device,
                                                  uint64_t* lockflags) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;

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

  uint8_t max_device_ctx_idx = 0;

  for (size_t ep = 0; ep < usb_device->endpoints_count; ep++) {
    struct xhci_usb_device_endpoint* endpoint = &usb_device->endpoints[ep];

    uint8_t dci = xhci_get_device_ctx_idx (&endpoint->desc);

    if (dci > max_device_ctx_idx)
      max_device_ctx_idx = dci;
  }

  void* out_ctx_virt = (void*)(xhci->xhci_dcbaa[usb_device->slot_id] + (uintptr_t)hhdm->offset);

  memcpy ((uint8_t*)input_ctx_virt + xhci->xhci_ctx_size, out_ctx_virt,
          PAGE_SIZE - xhci->xhci_ctx_size);

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

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

    ctx64->slot.dw[0] = (ctx64->slot.dw[0] & ~(0x1F << XHCI_SLCTX_CTX_ENTRIES)) |
        (max_device_ctx_idx << XHCI_SLCTX_CTX_ENTRIES);

    for (size_t ep = 0; ep < usb_device->endpoints_count; ep++) {
      struct xhci_usb_device_endpoint* endpoint = &usb_device->endpoints[ep];

      uint8_t dci = xhci_get_device_ctx_idx (&endpoint->desc);

      ctx64->ctrl.dw[1] |= (1 << dci);

      endpoint->transfer_ring.phys = pmm_alloc (1);
      endpoint->transfer_ring.trbs =
          (struct xhci_trb*)(endpoint->transfer_ring.phys + (uintptr_t)hhdm->offset);
      memset (endpoint->transfer_ring.trbs, 0, PAGE_SIZE);
      endpoint->transfer_ring.size = PAGE_SIZE / sizeof (struct xhci_trb);
      endpoint->transfer_ring.idx = 0;
      endpoint->transfer_ring.cycle_bit = 1;

      uint8_t type = endpoint->desc.attrs & 0x03;
      bool in = (endpoint->desc.endpoint_addr & 0x80) != 0;

      uint32_t xhci_type = 0;

      if (type == 2) {
        if (in)
          xhci_type = XHCI_EP_BULK_IN;
        else
          xhci_type = XHCI_EP_BULK_OUT;
      } else if (type == 3) {
        if (in)
          xhci_type = XHCI_EP_INTR_IN;
        else
          xhci_type = XHCI_EP_INTR_OUT;
      }

      ctx64->endpoints[dci - 1].dw[1] = (3 << XHCI_EPCTX_ERR_COUNT) |
          (xhci_type << XHCI_EPCTX_EP_TYPE) |
          ((uint32_t)endpoint->desc.max_packet_size << XHCI_EPCTX_MAX_PKT_SZ);
      ctx64->endpoints[dci - 1].dw[2] = (uint32_t)endpoint->transfer_ring.phys | (1 << 0);
      ctx64->endpoints[dci - 1].dw[3] = (uint32_t)(endpoint->transfer_ring.phys >> 32);
      ctx64->endpoints[dci - 1].dw[4] = (8 << XHCI_EPCTX_AVG_TRB_LEN);
    }
  } else if (xhci->xhci_ctx_size == 32) {
    struct xhci_input_ctx32* ctx32 = input_ctx_virt;

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

    ctx32->slot.dw[0] = (ctx32->slot.dw[0] & ~(0x1F << XHCI_SLCTX_CTX_ENTRIES)) |
        (max_device_ctx_idx << XHCI_SLCTX_CTX_ENTRIES);

    for (size_t ep = 0; ep < usb_device->endpoints_count; ep++) {
      struct xhci_usb_device_endpoint* endpoint = &usb_device->endpoints[ep];

      uint8_t dci = xhci_get_device_ctx_idx (&endpoint->desc);

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

      endpoint->transfer_ring.phys = pmm_alloc (1);
      endpoint->transfer_ring.trbs =
          (struct xhci_trb*)(endpoint->transfer_ring.phys + (uintptr_t)hhdm->offset);
      memset (endpoint->transfer_ring.trbs, 0, PAGE_SIZE);
      endpoint->transfer_ring.size = PAGE_SIZE / sizeof (struct xhci_trb);
      endpoint->transfer_ring.idx = 0;
      endpoint->transfer_ring.cycle_bit = 1;

      uint8_t type = endpoint->desc.attrs & 0x03;
      bool in = (endpoint->desc.endpoint_addr & 0x80) != 0;

      uint32_t xhci_type = 0;

      if (type == 2) {
        if (in)
          xhci_type = XHCI_EP_BULK_IN;
        else
          xhci_type = XHCI_EP_BULK_OUT;
      } else if (type == 3) {
        if (in)
          xhci_type = XHCI_EP_INTR_IN;
        else
          xhci_type = XHCI_EP_INTR_OUT;
      }

      ctx32->endpoints[dci - 1].dw[1] = (3 << XHCI_EPCTX_ERR_COUNT) |
          (xhci_type << XHCI_EPCTX_EP_TYPE) |
          ((uint32_t)endpoint->desc.max_packet_size << XHCI_EPCTX_MAX_PKT_SZ);
      ctx32->endpoints[dci - 1].dw[2] = (uint32_t)endpoint->transfer_ring.phys | (1 << 0);
      ctx32->endpoints[dci - 1].dw[3] = (uint32_t)(endpoint->transfer_ring.phys >> 32);
      ctx32->endpoints[dci - 1].dw[4] = (8 << XHCI_EPCTX_AVG_TRB_LEN);
    }
  }

  uint32_t ctrl = (usb_device->slot_id << 24) | (XHCI_TRB_CFG_ENDP_CMD << XHCI_GTRB_TRB_TYPE);
  xhci_send_cmd (xhci, input_ctx_phys, 0, ctrl, lockflags);

  if (xhci->last_cmpl_code != 1)
    DEBUG ("Failed to configure endpoints for this device\n");
  else
    DEBUG ("Successfully configured endpoints for this device\n");

  pmm_free (input_ctx_phys, 1);
}

static void xhci_usb_device_setup_set_config (struct xhci* xhci, struct xhci_usb_device* usb_device,
                                              uint8_t config_val, uint64_t* lockflags) {
  bool ok = xhci_endpoint0_ctrl_out (xhci, usb_device, 0x00, 0x09, config_val, 0, 0, 0, lockflags);

  if (!ok) {
    DEBUG ("Failed to set the config!\n");
    return;
  }

  DEBUG ("Config set!\n");
}

static void xhci_usb_device_setup_addressing (struct xhci* xhci, struct xhci_usb_device* usb_device,
                                              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,
             usb_device->xhci_port->port_value, ports_link);

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

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

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

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

  usb_device->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[usb_device->slot_id] = out_ctx_phys;

  usb_device->endpoint0_ring.phys = pmm_alloc (1);
  usb_device->endpoint0_ring.trbs =
      (struct xhci_trb*)(usb_device->endpoint0_ring.phys + (uintptr_t)hhdm->offset);
  memset (usb_device->endpoint0_ring.trbs, 0, PAGE_SIZE);
  usb_device->endpoint0_ring.size = PAGE_SIZE / sizeof (struct xhci_trb);
  usb_device->endpoint0_ring.idx = 0;
  usb_device->endpoint0_ring.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 << XHCI_SLCTX_CTX_ENTRIES) | (speed << XHCI_SLCTX_SPEED);
    ctx64->slot.dw[1] = ((usb_device->xhci_port->port_value + 1) << XHCI_SLCTX_ROOTHUBPRNUM);

    ctx64->endpoints[0].dw[0] = 0;
    ctx64->endpoints[0].dw[1] = (3 << XHCI_EPCTX_ERR_COUNT) |
        (XHCI_EP_CTRL_BI << XHCI_EPCTX_EP_TYPE) |
        (max_packet_size << XHCI_EPCTX_MAX_PKT_SZ);
    ctx64->endpoints[0].dw[2] = (uint32_t)usb_device->endpoint0_ring.phys | (1 << 0);
    ctx64->endpoints[0].dw[3] = (uint32_t)(usb_device->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 << XHCI_SLCTX_CTX_ENTRIES) | (speed << XHCI_SLCTX_SPEED);
    ctx32->slot.dw[1] = ((usb_device->xhci_port->port_value + 1) << XHCI_SLCTX_ROOTHUBPRNUM);

    ctx32->endpoints[0].dw[1] = (3 << XHCI_EPCTX_ERR_COUNT) |
        (XHCI_EP_CTRL_BI << XHCI_EPCTX_EP_TYPE) |
        (max_packet_size << XHCI_EPCTX_MAX_PKT_SZ);
    ctx32->endpoints[0].dw[2] = (uint32_t)usb_device->endpoint0_ring.phys | (1 << 0);
    ctx32->endpoints[0].dw[3] = (uint32_t)(usb_device->endpoint0_ring.phys >> 32);
  }

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

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

  pmm_free (input_ctx_phys, 1);
}

static bool xhci_usb_device_setup_get_info (struct xhci* xhci, struct xhci_usb_device* usb_device,
                                            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, usb_device, 0x80, 6, (1 << 8), 0, response_buf_phys, 18,
                                    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);

  usb_device->device_desc = *usb_desc;

  pmm_free (response_buf_phys, 1);
  return true;
}

static bool xhci_usb_device_setup_get_config (struct xhci* xhci, struct xhci_usb_device* usb_device,
                                              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, usb_device, 0x80, 6, (2 << 8), 0, response_buf_phys, 9,
                                    lockflags);

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

  struct usb_config_desc* usb_config = response_buf;

  DEBUG ("USB device config: total_len=%u num_ifs=%u, cfgval=%u cfg=%u attrs=%02x maxpow=%u\n",
         usb_config->total_length, usb_config->num_ifs, usb_config->config_value,
         usb_config->config, usb_config->attrs, usb_config->max_power);

  usb_device->config_desc = *usb_config;

  pmm_free (response_buf_phys, 1);
  return true;
}

static bool xhci_usb_device_setup_get_config_full (struct xhci* xhci,
                                                   struct xhci_usb_device* usb_device,
                                                   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, usb_device, 0x80, 6, (2 << 8), 0, response_buf_phys,
                                    usb_device->config_desc.total_length, lockflags);

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

  uint8_t* ptr = (uint8_t*)response_buf;
  uint8_t* end = ptr + usb_device->config_desc.total_length;

  while (ptr < end) {
    struct usb_desc_hdr* hdr = (struct usb_desc_hdr*)ptr;

    if (hdr->length == 0)
      break;

    switch (hdr->desc_type) {
    case USB_DESC_CONFIG: {
      struct xhci_usb_device_config* config = &usb_device->configs[usb_device->configs_count++];
      memset (config, 0, sizeof (*config));
      config->desc = *(struct usb_config_desc*)ptr;

      DEBUG (
          "Found USB device config: total_len=%u num_ifs=%u, cfgval=%u cfg=%u attrs=%02x maxpow=%u\n",
          config->desc.total_length, config->desc.num_ifs, config->desc.config_value,
          config->desc.config, config->desc.attrs, config->desc.max_power);
    } break;
    case USB_DESC_IF: {
      struct xhci_usb_device_if* if_ = &usb_device->ifs[usb_device->ifs_count++];
      memset (if_, 0, sizeof (*if_));
      if_->desc = *(struct usb_if_desc*)ptr;

      DEBUG (
          "Found USB device interface: ifnum=%u altsetting=%u eps=%u ifclass=%u ifsubclass=%u ifproto=%u if=%u\n",
          if_->desc.if_num, if_->desc.alt_setting, if_->desc.num_endpoints, if_->desc.if_class,
          if_->desc.if_subclass, if_->desc.if_proto, if_->desc.if1);
    } break;
    case USB_DESC_ENDPOINT: {
      struct xhci_usb_device_endpoint* endpoint =
          &usb_device->endpoints[usb_device->endpoints_count++];
      memset (endpoint, 0, sizeof (*endpoint));
      endpoint->desc = *(struct usb_endpoint_desc*)ptr;

      DEBUG ("Found USB device endpoint: addr=%u attrs=%u maxpkt=%u intvl=%u\n",
             endpoint->desc.endpoint_addr, endpoint->desc.attrs, endpoint->desc.max_packet_size,
             endpoint->desc.interval);
    } break;
    default:
      DEBUG ("Found unknown USB descriptor of type %u\n", hdr->desc_type);
      break;
    }

    ptr += hdr->length;
  }

  pmm_free (response_buf_phys, 1);
  return true;
}

static void xhci_poll_setup_init_ifs (struct xhci* xhci, struct xhci_usb_device* usb_device,
                                      uint64_t* lockflags) {
  struct reschedule_ctx rctx;
  memset (&rctx, 0, sizeof (rctx));

  for (size_t if0 = 0; if0 < usb_device->ifs_count; if0++) {
    struct xhci_usb_device_if* if_ = &usb_device->ifs[if0];

    for (size_t i = 0; i < USB_DRIVER_MAX_MATCHES; i++) {
      struct usb_driver_info* info = &usb_driver_infos[i];

      if (if_->desc.if_class == info->if_class &&
          if_->desc.if_subclass == info->if_subclass &&
          if_->desc.if_proto == info->if_proto) {
        struct device* device = info->init (xhci, usb_device, thiscpu->kproc, &rctx, lockflags);

        if (device == NULL)
          DEBUG ("USB driver failed to initialize. Skipping device!\n");

        usb_device->device = device;
      }
    }
  }
}

static void xhci_poll_setup_devices (struct xhci* xhci, uint64_t* lockflags) {
  struct list_node_link *usb_device_link, *tmp_usb_device_link;

  list_foreach (xhci->xhci_usb_devices, usb_device_link, tmp_usb_device_link) {
    struct xhci_usb_device* usb_device =
        list_entry (usb_device_link, struct xhci_usb_device, usb_devices_link);

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

    xhci_usb_device_setup_addressing (xhci, usb_device, lockflags);
    xhci_usb_device_setup_get_info (xhci, usb_device, lockflags);
    xhci_usb_device_setup_get_config (xhci, usb_device, lockflags);
    xhci_usb_device_setup_get_config_full (xhci, usb_device, lockflags);
    xhci_usb_device_setup_set_config (xhci, usb_device, usb_device->config_desc.config_value,
                                      lockflags);
    xhci_usb_device_setup_init_endpoints (xhci, usb_device, lockflags);
    xhci_poll_setup_init_ifs (xhci, usb_device, lockflags);
  }
}

static void xhci_poll_port_changes (struct xhci* xhci, struct proc* proc,
                                    struct reschedule_ctx* rctx) {
  struct list_node_link *port_change_link, *tmp_port_change_link;
  struct xhci_port* xhci_port;

  list_foreach (xhci->port_changes, port_change_link, tmp_port_change_link) {
    struct xhci_port_status_change* change =
        list_entry (port_change_link, struct xhci_port_status_change, port_changes_link);

    list_remove (xhci->port_changes, &change->port_changes_link);

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

    if ((change->portsc & (1 << XHCI_PORTSC_CSC))) {
      if ((change->portsc & (1 << XHCI_PORTSC_CCS))) {
        DEBUG ("Device attached to port %u!\n", change->port);

        xhci_port_reset (xhci, change->port);
        xhci_create_usb_device (xhci, xhci_port);
      } else {
        DEBUG ("Device detached from port %u!\n", change->port);

        xhci_delete_usb_device (xhci, xhci_port, proc, rctx);
      }
    }

    free (change);
  }
}

int xhci_bulk_transfer (struct xhci* xhci, struct xhci_usb_device* usb_device,
                        uint8_t endpoint_addr, uintptr_t buffer_phys, size_t buffer_size,
                        uint64_t* lockflags) {
  struct xhci_usb_device_endpoint* endpoint = NULL;

  for (size_t ep = 0; ep < XHCI_USB_DEVICE_ENDPOINTS_MAX; ep++) {
    endpoint = &usb_device->endpoints[ep];

    if (endpoint->desc.endpoint_addr == endpoint_addr)
      break;
  }

  if (endpoint == NULL)
    return -ST_NOT_FOUND;

  /* not bulk */
  if ((endpoint->desc.attrs & 0x03) != 0x02)
    return -ST_NOT_FOUND;

  if (usb_device->slot_id < 0)
    return -ST_NOT_FOUND;

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

  uint32_t rem = buffer_size;
  uint32_t current = buffer_phys;

  while (rem > 0) {
    uint32_t chunk = rem > 0x1FFFFU ? 0x1FFFFU : rem;

    struct xhci_trb trb;
    memset (&trb, 0, sizeof (trb));
    trb.param = current;
    trb.status = chunk;

    write_memory_barrier ();

    trb.ctrl = (XHCI_TRB_NORMAL << XHCI_GTRB_TRB_TYPE);

    if (rem != chunk) {
      trb.ctrl |= (1 << XHCI_DSTRB_CTRL_CH);
    } else {
      trb.ctrl |= (1 << XHCI_DSTRB_CTRL_IOC) | (1 << XHCI_DSTRB_CTRL_ISP);
    }

    xhci_ring_put_trb (&endpoint->transfer_ring, trb);

    current += chunk;
    rem -= chunk;
  }

  uint32_t dci = xhci_get_device_ctx_idx (&endpoint->desc);

  int timeout = 100;

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

  xhci_write32 (xhci->xhci_doorbell_base, usb_device->slot_id * 4, dci);

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

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

  if (timeout == 0) {
    DEBUG ("bulk transfer timed out\n");
    return -ST_TRY_AGAIN;
  }

  return xhci->last_cmpl_code == 1 ? ST_OK : -ST_USB_CTRL_ERROR;
}

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

  xhci_poll_port_changes (xhci, proc, rctx);
  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 >> XHCI_HCSPARAMS2_MAX_SCRTCH_HI) & 0x1F) << 5) |
      ((hcsparams2 >> XHCI_HCSPARAMS2_MAX_SCRTCH_LO) & 0x1F);

  uint32_t hccparams1 = xhci_read32 (xhci->xhci_mmio_base, XHCI_HCCPARAMS1);
  xhci->xhci_ctx_size = (hccparams1 & (1 << XHCI_HCCPARAMS1_CSZ)) ? 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 (100);

  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.trbs = (struct xhci_trb*)(xhci->cmd_ring.phys + (uintptr_t)hhdm->offset);
  memset (xhci->cmd_ring.trbs, 0, PAGE_SIZE);
  xhci->cmd_ring.size = PAGE_SIZE / sizeof (struct xhci_trb);
  xhci->cmd_ring.idx = 0;
  xhci->cmd_ring.cycle_bit = 1;

  uint64_t crcr = xhci->cmd_ring.phys | xhci->cmd_ring.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.trbs = (struct xhci_trb*)(xhci->event_ring.phys + (uintptr_t)hhdm->offset);
  memset (xhci->event_ring.trbs, 0, PAGE_SIZE);
  xhci->event_ring.size = PAGE_SIZE / sizeof (struct xhci_trb);
  xhci->event_ring.idx = 0;
  xhci->event_ring.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));

  /* 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 (250);

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

    xhci_delete_usb_device (xhci, port, proc, rctx);
    free (port);
  }

  irq_detach (xhci->irq);

  free (xhci);
}