#include "src/kernel/xhci.h"
#include "src/kernel/log.h"
#include "src/kernel/memory.h"
#include "src/kernel/panic.h"
#include "src/kernel/pci.h"
#include "src/kernel/stream.h"
#include "src/lib/layout.h"
#include "src/lib/memory.h"
#include "src/lib/util.h"

#define XHCI_PCI_CLASS 0x0C
#define XHCI_PCI_SUBCLASS 0x03
#define XHCI_PCI_INTERFACE 0x30

typedef volatile struct __attribute__((packed)) {
  uint8_t caplength;
  uint8_t reserved;
  uint16_t hciversion;
  uint32_t hcsparams1;
  uint32_t hcsparams2;
  uint32_t hcsparams3;
  uint32_t hccparams1;
  uint32_t dboff;
  uint32_t rtsoff;
  uint32_t hccparams2;
} xhci_cap_regs_t;

#define XHCI_CAP_REGS_MAX_SLOTS(c) BITS_R((c)->hcsparams1, 7, 0)
#define XHCI_CAP_REGS_MAX_INTRS(c) BITS_R((c)->hcsparams1, 18, 8)
#define XHCI_CAP_REGS_MAX_PORTS(c) BITS_R((c)->hcsparams1, 31, 24)
#define XHCI_CAP_REGS_MAX_ERSTS(c) BITS_R((c)->hcsparams2, 7, 4)
#define XHCI_CAP_REGS_MAX_SCRATCHPAD_BUFFERS(c) (BITS_R((c)->hcsparams2, 31, 27) << 5 | BITS_R((c)->hcsparams2, 4, 0))

typedef volatile struct __attribute__((packed)) {
  uint32_t usbcmd;
  uint32_t usbsts;
  uint32_t pagesize;
  uint32_t reserved[2];
  uint32_t dnctrl;
  uint64_t crcr;
  uint32_t reserved2[4];
  uint64_t dcbaap;
  uint32_t config;
} xhci_op_regs_t;

#define XHCI_OP_REGS_RS_R(o) BITS_R((o)->usbcmd, 0, 0)
#define XHCI_OP_REGS_RS_W(o, v) BITS_W((o)->usbcmd, 0, 0, v)
#define XHCI_OP_REGS_HCRST_R(o) BITS_R((o)->usbcmd, 1, 1)
#define XHCI_OP_REGS_HCRST_W(o, v) BITS_W((o)->usbcmd, 1, 1, v)
#define XHCI_OP_REGS_HCH(o) BITS_R((o)->usbsts, 0, 0)
#define XHCI_OP_REGS_HSE(o) BITS_R((o)->usbsts, 2, 2)
#define XHCI_OP_REGS_EINT_R(o) BITS_R((o)->usbsts, 3, 3)
#define XHCI_OP_REGS_EINT_W(o, v) BITS_W((o)->usbsts, 3, 3, v)
#define XHCI_OP_REGS_CNR(o) BITS_R((o)->usbsts, 11, 11)
#define XHCI_OP_REGS_HCE(o) BITS_R((o)->usbsts, 12, 12)
#define XHCI_OP_REGS_RCS(o) BITS_R((o)->crcr, 0, 0)
#define XHCI_OP_REGS_CRR(o) BITS_R((o)->crcr, 3, 3)
#define XHCI_OP_REGS_CRP(o) BITS_R((o)->crcr, 63, 6)
#define XHCI_OP_REGS_MAX_SLOTS_EN_R(o) BITS_R((o)->config, 7, 0)
#define XHCI_OP_REGS_MAX_SLOTS_EN_W(o, v) BITS_W((o)->config, 7, 0, v)

typedef volatile struct __attribute__((packed)) {
  uint32_t portsc;
  uint32_t portpmsc;
  uint32_t portli;
  uint32_t porthlpmc;
} xhci_port_regs_t;

#define XHCI_PORT_REGS_CCS(p) BITS_R((p)->portsc, 0, 0)
#define XHCI_PORT_REGS_PED(p) BITS_R((p)->portsc, 1, 1)
#define XHCI_PORT_REGS_PR_R(p) BITS_R((p)->portsc, 4, 4)
#define XHCI_PORT_REGS_PR_W(p, v) BITS_W((p)->portsc, 4, 4, v)
#define XHCI_PORT_REGS_PLL(p) BITS_R((p)->portsc, 8, 5)
#define XHCI_PORT_REGS_PP(p) BITS_R((p)->portsc, 9, 9)
#define XHCI_PORT_REGS_PS(p) BITS_R((p)->portsc, 13, 10)

typedef volatile struct __attribute__((packed)) {
  uint32_t iman;
  uint32_t imod;
  uint32_t erstsz;
  uint32_t reserved;
  uint64_t erstba;
  uint64_t erdp;
} xhci_intr_regs_t;

static xhci_cap_regs_t *cap_regs = (xhci_cap_regs_t *)KERNEL_VIRTUAL_XHCI;
static xhci_op_regs_t *op_regs;
static xhci_port_regs_t *port_regs;
static xhci_intr_regs_t *intr_regs;

typedef volatile uint32_t xhci_db_regs;

#define XHCI_DB_REGS_TASK_ID(d) BITS_R((d), 31, 16)
#define XHCI_DB_REGS_TARGET(d) BITS_R((d), 7, 0)

static xhci_db_regs *db_regs;

#define MAX_SLOTS 32

static void *scratchpads[32];

static void *dcbaa[MAX_SLOTS + 1] __attribute__((aligned(PAGE_SIZE)));

typedef struct __attribute__((packed)) {
  uint64_t base;
  uint32_t size;
  uint32_t reserved;
} xchi_erst_entry_t;

typedef struct __attribute__((packed)) {
  uint64_t parameter;
  uint32_t status;
  uint32_t control;
} xhci_trb_t;

#define XHCI_TRB_CYCLE_R(t) BITS_R((t)->control, 0, 0)
#define XHCI_TRB_CYCLE_W(t, v) BITS_W((t)->control, 0, 0, v)
#define XHCI_TRB_TOGGLE_CYCLE_R(t) BITS_R((t)->control, 1, 1)
#define XHCI_TRB_TOGGLE_CYCLE_W(t, v) BITS_W((t)->control, 1, 1, v)
#define XHCI_TRB_IOC_R(t) BITS_R((t)->control, 5, 5)
#define XHCI_TRB_IOC_W(t, v) BITS_W((t)->control, 5, 5, v)
#define XHCI_TRB_TRB_TYPE_R(t) BITS_R((t)->control, 15, 10)
#define XHCI_TRB_TRB_TYPE_W(t, v) BITS_W((t)->control, 15, 10, v)
#define XHCI_TRB_SLOT_TYPE_R(t) BITS_R((t)->control, 20, 16)
#define XHCI_TRB_SLOT_TYPE_W(t, v) BITS_W((t)->control, 20, 16, v)
#define XHCI_TRB_ENDPOINT_ID_R(t) BITS_R((t)->control, 20, 16)
#define XHCI_TRB_ENDPOINT_ID_W(t, v) BITS_W((t)->control, 20, 16, v)
#define XHCI_TRB_COMPLETION_CODE(t) BITS_R((t)->status, 31, 24)
#define XHCI_TRB_SLOT_ID_R(t) BITS_R((t)->control, 31, 24)
#define XHCI_TRB_SLOT_ID_W(t, v) BITS_W((t)->control, 31, 24, v)

#define XHCI_TRB_TYPE_TRANSFER_NORMAL 1
#define XHCI_TRB_TYPE_TRANSFER_SETUP_STAGE 2
#define XHCI_TRB_TYPE_TRANSFER_DATA_STAGE 3
#define XHCI_TRB_TYPE_TRANSFER_STATUS_STAGE 4
#define XHCI_TRB_TYPE_LINK 6
#define XHCI_TRB_TYPE_COMMAND_ENABLE_SLOT 9
#define XHCI_TRB_TYPE_COMMAND_DISABLE_SLOT 10
#define XHCI_TRB_TYPE_COMMAND_ADDRESS_DEVICE 11
#define XHCI_TRB_TYPE_COMMAND_CONFIGURE_ENDPOINT 12
#define XHCI_TRB_TYPE_COMMAND_STOP_ENDPOINT 15
#define XHCI_TRB_TYPE_COMMAND_NOOP 23
#define XHCI_TRB_TYPE_EVENT_TRANSFER_COMPLETION 32
#define XHCI_TRB_TYPE_EVENT_COMMAND_COMPLETION 33
#define XHCI_TRB_TYPE_EVENT_PORT_STATUS_CHANGE 34

static xchi_erst_entry_t erst[1] __attribute__((aligned(PAGE_SIZE)));

#define RING_LENGTH 32

#define RING(name)                                                                                                                         \
  static uint8_t name##_r_cycle = 1;                                                                                                       \
  static uint8_t name##_r_i = 0;                                                                                                           \
  static xhci_trb_t name##_r[RING_LENGTH + 1] __attribute__((aligned(PAGE_SIZE)));

#define RING_WRAP(name)                                                                                                                    \
  xhci_trb_t *name##_wrapper = &name##_r[RING_LENGTH];                                                                                     \
  name##_wrapper->parameter = (uint64_t)VIRT_TO_PHYS(name##_r);                                                                            \
  name##_wrapper->status = 0;                                                                                                              \
  XHCI_TRB_TRB_TYPE_W(name##_wrapper, XHCI_TRB_TYPE_LINK);                                                                                 \
  XHCI_TRB_TOGGLE_CYCLE_W(name##_wrapper, 1);

#define RING_ADVANCE(name)                                                                                                                 \
  name##_r_i = (name##_r_i + 1) % RING_LENGTH;                                                                                             \
  XHCI_TRB_CYCLE_W(&name##_r[RING_LENGTH], name##_r_cycle);                                                                                \
  name##_r_cycle ^= (name##_r_i == 0);

RING(cmd);
RING(evt);

typedef struct __attribute__((packed)) {
  uint32_t drop_flags;
  uint32_t add_flags;
  uint32_t reserved[6];
} xhci_input_control_ctx_t;

typedef struct __attribute__((packed)) {
  uint32_t dw0;
  uint32_t dw1;
  uint32_t dw2;
  uint32_t dw3;
  uint32_t reserved[4];
} xhci_slot_ctx_t;

#define XHCI_SLOT_CTX_SPEED_R(s) BITS_R((s)->dw0, 23, 20)
#define XHCI_SLOT_CTX_SPEED_W(s, v) BITS_W((s)->dw0, 23, 20, v)
#define XHCI_SLOT_CTX_CTX_ENTRIES_R(s) BITS_R((s)->dw0, 31, 27)
#define XHCI_SLOT_CTX_CTX_ENTRIES_W(s, v) BITS_W((s)->dw0, 31, 27, v)
#define XHCI_SLOT_CTX_ROOT_HUB_PORT_R(s) BITS_R((s)->dw1, 23, 16)
#define XHCI_SLOT_CTX_ROOT_HUB_PORT_W(s, v) BITS_W((s)->dw1, 23, 16, v)

typedef struct __attribute__((packed)) {
  uint32_t dw0;
  uint32_t dw1;
  uint64_t dequeue;
  uint32_t dw4;
  uint32_t reserved[3];
} xhci_ep_ctx_t;

#define XHCI_EP_CTX_INTERVAL_R(e) BITS_R((e)->dw0, 23, 16)
#define XHCI_EP_CTX_INTERVAL_W(e, v) BITS_W((e)->dw0, 23, 16, v)
#define XHCI_EP_CTX_EP_TYPE_R(e) BITS_R((e)->dw1, 5, 3)
#define XHCI_EP_CTX_EP_TYPE_W(e, v) BITS_W((e)->dw1, 5, 3, v)
#define XHCI_EP_CTX_MAX_PACKET_SIZE_R(e) BITS_R((e)->dw1, 31, 16)
#define XHCI_EP_CTX_MAX_PACKET_SIZE_W(e, v) BITS_W((e)->dw1, 31, 16, v)

#define XHCI_EP_CTX_MAX_PACKET_SIZE 64

typedef struct __attribute__((packed)) {
  xhci_input_control_ctx_t control;
  xhci_slot_ctx_t slot;
  xhci_ep_ctx_t ep[31];
} xhci_input_ctx_t;

static xhci_input_ctx_t input_ctx __attribute__((aligned(64)));

typedef struct __attribute__((packed)) {
  xhci_slot_ctx_t slot;
  xhci_ep_ctx_t endpoints[31];
} xchi_device_ctx_t;

static xchi_device_ctx_t device_ctx __attribute__((aligned(64)));

static uint8_t tsf_r_attached_slot = 0;
RING(tsf);

static uint8_t ep_r_attached_endpoint = 0;
RING(ep);

static uint16_t ep_transfer_size = 0;
static uint8_t ep_buffer[RING_LENGTH][XHCI_EP_CTX_MAX_PACKET_SIZE] __attribute__((aligned(64)));
static uint8_t ep_processed_events = RING_LENGTH;

static const xhci_trb_t *command_exec_sync(const xhci_trb_t *cmd) {
  LOG_LN_STEP("Queueing command TRB...");
  LOG_VAL_TRACE(cmd->parameter, "%lx");
  LOG_VAL_TRACE(cmd->status, "%x");
  LOG_VAL_TRACE(cmd->control, "%x");

  xhci_trb_t *cre = &cmd_r[cmd_r_i];
  memory_copy(cmd, sizeof(xhci_trb_t), cre);
  XHCI_TRB_CYCLE_W(cre, cmd_r_cycle);
  RING_ADVANCE(cmd);

  LOG_LN_STEP("Ringing command doorbell...");
  LOG_VAL_TRACE(op_regs->usbsts, "%x");
  db_regs[0] = 0;
  LOG_VAL_TRACE(op_regs->usbsts, "%x");

  LOG_LN_STEP("Waiting for completion...");
  while (1) {
    xhci_trb_t *cmp = &evt_r[evt_r_i];
    while (XHCI_TRB_CYCLE_R(cmp) != evt_r_cycle)
      ;
    RING_ADVANCE(evt);

    if (XHCI_TRB_TRB_TYPE_R(cmp) != XHCI_TRB_TYPE_EVENT_COMMAND_COMPLETION || (void *)cmp->parameter != VIRT_TO_PHYS(cre)) {
      LOG_LN_STEP("Received irrelevant event, skipping...");

      LOG_VAL_TRACE(cmp->parameter, "%lx");
      LOG_VAL_TRACE(cmp->status, "%x");
      LOG_VAL_TRACE(cmp->control, "%x");
    } else {
      LOG_LN_STEP("Received command completion event...");

      LOG_VAL_TRACE(cmp->parameter, "%lx");
      LOG_VAL_TRACE(cmp->status, "%x");
      LOG_VAL_TRACE(cmp->control, "%x");

      ASSERT(XHCI_TRB_COMPLETION_CODE(cmp) == 1, "command_exec_sync: Command failed.");

      LOG_LN_STEP("Done.");

      return cmp;
    }
  }
}

static const xhci_trb_t *control_transfer_exec_sync(uint64_t setup, void *data, uint16_t length) {
  ASSERT(tsf_r_attached_slot, "control_transfer_exec_sync: Slot not attached.");

  LOG_LN_STEP("Queueing transfer TRBs...");

  xhci_trb_t *tre = &tsf_r[tsf_r_i];

  tre->parameter = setup;
  tre->status = sizeof(setup);
  XHCI_TRB_TRB_TYPE_W(tre, XHCI_TRB_TYPE_TRANSFER_SETUP_STAGE);
  BITS_W(tre->control, 17, 16, 3);
  BITS_W(tre->control, 6, 6, 1);
  XHCI_TRB_CYCLE_W(tre, tsf_r_cycle);

  LOG_VAL_TRACE(tre->parameter, "%lx");
  LOG_VAL_TRACE(tre->status, "%x");
  LOG_VAL_TRACE(tre->control, "%x");

  RING_ADVANCE(tsf);

  if (data != NUL) {
    tre = &tsf_r[tsf_r_i];

    tre->parameter = (uint64_t)VIRT_TO_PHYS(data);
    tre->status = length;
    XHCI_TRB_TRB_TYPE_W(tre, XHCI_TRB_TYPE_TRANSFER_DATA_STAGE);
    BITS_W(tre->control, 16, 16, 1);
    XHCI_TRB_CYCLE_W(tre, tsf_r_cycle);

    LOG_VAL_TRACE(tre->parameter, "%lx");
    LOG_VAL_TRACE(tre->status, "%x");
    LOG_VAL_TRACE(tre->control, "%x");

    RING_ADVANCE(tsf);
  }

  tre = &tsf_r[tsf_r_i];

  XHCI_TRB_TRB_TYPE_W(tre, XHCI_TRB_TYPE_TRANSFER_STATUS_STAGE);
  BITS_W(tre->control, 16, 16, 0);
  XHCI_TRB_IOC_W(tre, 1);
  XHCI_TRB_CYCLE_W(tre, tsf_r_cycle);

  RING_ADVANCE(tsf);

  LOG_LN_STEP("Ringing slot doorbell...");
  db_regs[tsf_r_attached_slot] = 1;

  LOG_LN_STEP("Waiting for completion...");
  while (1) {
    xhci_trb_t *cmp = &evt_r[evt_r_i];
    while (XHCI_TRB_CYCLE_R(cmp) != evt_r_cycle)
      ;

    RING_ADVANCE(evt);

    if (XHCI_TRB_TRB_TYPE_R(cmp) != XHCI_TRB_TYPE_EVENT_TRANSFER_COMPLETION || (void *)cmp->parameter != VIRT_TO_PHYS(tre)) {
      LOG_LN_STEP("Received irrelevant event, skipping...");

      LOG_VAL_TRACE(cmp->parameter, "%lx");
      LOG_VAL_TRACE(cmp->status, "%x");
      LOG_VAL_TRACE(cmp->control, "%x");
    } else {
      LOG_LN_STEP("Received transfer completion event...");

      LOG_VAL_TRACE(cmp->parameter, "%lx");
      LOG_VAL_TRACE(cmp->status, "%x");
      LOG_VAL_TRACE(cmp->control, "%x");

      ASSERT(XHCI_TRB_COMPLETION_CODE(cmp) == 1, "control_transfer_exec_sync: Transfer failed.");

      LOG_LN_STEP("Done.");

      return cmp;
    }
  }
}

void xhci_init() {
  LOG_LN_INFO("Searching for suitable controller...");
  pci_bdf_t bdf = pci_find_by_class(XHCI_PCI_CLASS, XHCI_PCI_SUBCLASS, XHCI_PCI_INTERFACE);
  ASSERT(bdf != (pci_bdf_t)-1, "xhci_init: No suitable controllers found.");

  LOG_LN_INFO("Mapping controller to virtual memory...");
  pci_map(bdf, cap_regs, 4);
  op_regs = (xhci_op_regs_t *)((uint8_t *)cap_regs + cap_regs->caplength);
  port_regs = (xhci_port_regs_t *)((uint8_t *)cap_regs + cap_regs->caplength + 0x400);
  intr_regs = (xhci_intr_regs_t *)((uint8_t *)cap_regs + cap_regs->rtsoff + 0x20);
  db_regs = (uint32_t *)((uint8_t *)cap_regs + cap_regs->dboff);

  LOG_VAL_DATA(cap_regs->caplength, "%hhx");
  LOG_VAL_DATA(cap_regs->hciversion, "%hx");
  LOG_VAL_DATA(cap_regs->hcsparams1, "%x");
  LOG_VAL_DATA(cap_regs->hcsparams2, "%x");
  LOG_VAL_DATA(cap_regs->hcsparams3, "%x");
  LOG_VAL_DATA(cap_regs->hccparams1, "%x");
  LOG_VAL_DATA(cap_regs->dboff, "%x");
  LOG_VAL_DATA(cap_regs->rtsoff, "%x");
  LOG_VAL_DATA(cap_regs->hccparams2, "%x");

  LOG_VAL_DATA(op_regs->usbcmd, "%x");
  LOG_VAL_DATA(op_regs->usbsts, "%x");
  LOG_VAL_DATA(op_regs->pagesize, "%x");
  LOG_VAL_DATA(op_regs->dnctrl, "%x");
  LOG_VAL_DATA(op_regs->crcr, "%lx");
  LOG_VAL_DATA(op_regs->dcbaap, "%lx");
  LOG_VAL_DATA(op_regs->config, "%x");

  LOG_LN_INFO("Stopping controller...");
  XHCI_OP_REGS_RS_W(op_regs, 0);
  LOG_VAL_DATA(op_regs->usbcmd, "%x");
  LOG_VAL_DATA(op_regs->usbsts, "%x");
  while (!XHCI_OP_REGS_HCH(op_regs))
    ;
  LOG_VAL_DATA(op_regs->usbcmd, "%x");
  LOG_VAL_DATA(op_regs->usbsts, "%x");

  LOG_LN_INFO("Resetting controller...");
  XHCI_OP_REGS_EINT_W(op_regs, 0);
  XHCI_OP_REGS_HCRST_W(op_regs, 1);
  LOG_VAL_DATA(op_regs->usbcmd, "%x");
  LOG_VAL_DATA(op_regs->usbsts, "%x");
  while (XHCI_OP_REGS_HCRST_R(op_regs))
    ;
  LOG_VAL_DATA(op_regs->usbcmd, "%x");
  LOG_VAL_DATA(op_regs->usbsts, "%x");
  while (XHCI_OP_REGS_CNR(op_regs))
    ;
  LOG_VAL_DATA(op_regs->usbcmd, "%x");
  LOG_VAL_DATA(op_regs->usbsts, "%x");

  LOG_LN_INFO("Setting up command ring...");
  RING_WRAP(cmd);
  XHCI_TRB_CYCLE_W(&cmd_r[RING_LENGTH], cmd_r_cycle);

  LOG_LN_INFO("Setting up scratchpad buffers...");
  ASSERT(XHCI_CAP_REGS_MAX_SCRATCHPAD_BUFFERS(cap_regs) <= sizeof(scratchpads), "Too many scratchpad buffers expected.");
  for (uint32_t i = 0; i < XHCI_CAP_REGS_MAX_SCRATCHPAD_BUFFERS(cap_regs); i++) {
    scratchpads[i] = memory_page_allocate();
  }
  dcbaa[0] = VIRT_TO_PHYS(scratchpads);

  LOG_LN_INFO("Configuring controller...");
  XHCI_OP_REGS_MAX_SLOTS_EN_W(op_regs, MAX_SLOTS);
  op_regs->dcbaap = (uint64_t)VIRT_TO_PHYS(dcbaa);
  op_regs->crcr = (uint64_t)VIRT_TO_PHYS(cmd_r) | cmd_r_cycle;
  erst[0].size = RING_LENGTH;
  erst[0].base = (uint64_t)VIRT_TO_PHYS(evt_r);
  intr_regs->erstsz = sizeof(erst) / sizeof(xchi_erst_entry_t);
  intr_regs->erstba = (uint64_t)VIRT_TO_PHYS(erst);
  intr_regs->erdp = erst[0].base;

  LOG_VAL_DATA(op_regs->config, "%x");
  LOG_VAL_DATA(op_regs->dcbaap, "%lx");
  LOG_VAL_DATA(op_regs->crcr, "%lx");
  LOG_VAL_DATA(intr_regs->erstba, "%lx");
  LOG_VAL_DATA(intr_regs->erstsz, "%x");
  LOG_VAL_DATA(intr_regs->erdp, "%lx");
  LOG_VAL_DATA(erst[0].size, "%hx");
  LOG_VAL_DATA(erst[0].base, "%lx");

  LOG_LN_INFO("Starting controller...");
  LOG_VAL_DATA(op_regs->usbcmd, "%x");
  LOG_VAL_DATA(op_regs->usbsts, "%x");
  XHCI_OP_REGS_RS_W(op_regs, 1);
  LOG_VAL_DATA(op_regs->usbcmd, "%x");
  LOG_VAL_DATA(op_regs->usbsts, "%x");
  while (XHCI_OP_REGS_HCH(op_regs))
    ;
  LOG_VAL_DATA(op_regs->usbcmd, "%x");
  LOG_VAL_DATA(op_regs->usbsts, "%x");
  sleep(1000); // Wait for ports to settle

  PRINT_LN(kernel_log, "Enumerating ports...");
  xhci_enumerate(kernel_log);
  PRINT_LN(kernel_log, "Done.");
}

void xhci_enumerate(stream_t *out) {
  for (uint16_t i = 0; i < XHCI_CAP_REGS_MAX_PORTS(cap_regs); i++) {
    PRINT_LN(out, "port=%d CCS=%d PED=%d PP=%d PS=%d portsc=%x", i, XHCI_PORT_REGS_CCS(&port_regs[i]), XHCI_PORT_REGS_PED(&port_regs[i]),
             XHCI_PORT_REGS_PP(&port_regs[i]), XHCI_PORT_REGS_PS(&port_regs[i]), port_regs[i].portsc);
  }
}

uint8_t xhci_port_count() {
  return XHCI_CAP_REGS_MAX_PORTS(cap_regs);
}

uint8_t xhci_port_connected(xhci_port_t port) {
  ASSERT(port < XHCI_CAP_REGS_MAX_PORTS(cap_regs), "xhci_port_connected: Port doesn't exist.");

  return XHCI_PORT_REGS_CCS(&port_regs[port]);
}

xhci_slot_t xhci_attach(xhci_port_t port) {
  LOG_LN_STEP("Attaching port %d ...", port);

  ASSERT(port < XHCI_CAP_REGS_MAX_PORTS(cap_regs), "xhci_attach: Port doesn't exist.");
  ASSERT(!tsf_r_attached_slot, "xhci_attach: Can only attach one port at a time.");
  ASSERT(XHCI_PORT_REGS_CCS(&port_regs[port]), "xhci_attach: Port is not connected.");

  LOG_LN_STEP("Setting up transfer ring...");
  memory_set(0, sizeof(tsf_r), tsf_r);
  tsf_r_i = 0;
  tsf_r_cycle = 1;
  RING_WRAP(tsf);

  LOG_LN_STEP("Resetting device...");
  LOG_VAL_TRACE(port_regs[port].portsc, "%x");
  port_regs[port].portsc = (port_regs[port].portsc & ~(uint32_t)0x00FE0000) | 0x00FE0000;
  LOG_VAL_TRACE(port_regs[port].portsc, "%x");
  port_regs[port].portsc = (port_regs[port].portsc & ~(uint32_t)0x00FE0000) | (1 << 4);
  LOG_VAL_TRACE(port_regs[port].portsc, "%x");
  while (XHCI_PORT_REGS_PR_R(&port_regs[port]))
    ;
  LOG_VAL_TRACE(port_regs[port].portsc, "%x");
  while (!XHCI_PORT_REGS_PED(&port_regs[port]) && XHCI_PORT_REGS_PLL(&port_regs[port]) != 0)
    ;
  LOG_VAL_TRACE(port_regs[port].portsc, "%x");

  LOG_LN_STEP("Enabling a slot...");
  xhci_trb_t cmd_es;
  memory_set(0, sizeof(xhci_trb_t), &cmd_es);
  XHCI_TRB_TRB_TYPE_W(&cmd_es, XHCI_TRB_TYPE_COMMAND_ENABLE_SLOT);
  xhci_slot_t slot = XHCI_TRB_SLOT_ID_R(command_exec_sync(&cmd_es));

  LOG_LN_STEP("Addressing device...");
  memory_set(0, sizeof(input_ctx), &input_ctx);
  input_ctx.control.add_flags = 0b11;
  XHCI_SLOT_CTX_CTX_ENTRIES_W(&input_ctx.slot, 1);
  XHCI_SLOT_CTX_ROOT_HUB_PORT_W(&input_ctx.slot, port + 1);
  XHCI_SLOT_CTX_SPEED_W(&input_ctx.slot, XHCI_PORT_REGS_PS((&port_regs[port])));
  XHCI_EP_CTX_EP_TYPE_W(&input_ctx.ep[0], 4);
  XHCI_EP_CTX_MAX_PACKET_SIZE_W(&input_ctx.ep[0], 64);
  input_ctx.ep[0].dequeue = (uint64_t)VIRT_TO_PHYS(tsf_r) | tsf_r_cycle;

  dcbaa[slot] = VIRT_TO_PHYS(&device_ctx);

  xhci_trb_t cmd_ad;
  memory_set(0, sizeof(xhci_trb_t), &cmd_ad);
  cmd_ad.parameter = (uint64_t)VIRT_TO_PHYS(&input_ctx);
  XHCI_TRB_TRB_TYPE_W(&cmd_ad, XHCI_TRB_TYPE_COMMAND_ADDRESS_DEVICE);
  XHCI_TRB_SLOT_ID_W(&cmd_ad, slot);
  command_exec_sync(&cmd_ad);

  tsf_r_attached_slot = slot;

  LOG_LN_STEP("Done.");

  return slot;
}

uint64_t xhci_control_transfer(xhci_slot_t slot, uint64_t setup, void *data, uint16_t length) {
  ASSERT(slot < MAX_SLOTS, "xhci_control_transfer: Slot does not exist.");
  ASSERT(tsf_r_attached_slot == slot, "xhci_control_transfer: Slot is not attached.");
  ASSERT(!ep_r_attached_endpoint, "xhci_control_transfer: Can not control transfer while in polling mode.");

  return length - (control_transfer_exec_sync(setup, data, length)->status & 0xFFFFFF);
}

xhci_endpoint_t xhci_open_endpoint(xhci_slot_t slot, uint8_t address, uint8_t type, uint16_t max_packet_size, uint8_t interval) {
  LOG_LN_INFO("Opening endpoint slot %d address %hhx...", slot, address);

  ASSERT(slot < MAX_SLOTS, "xhci_control_transfer: Slot does not exist.");
  ASSERT(tsf_r_attached_slot == slot, "xhci_open_endpoint: Slot is not attached.");
  ASSERT(!ep_r_attached_endpoint, "xhci_open_endpoint: Can only open one endpoint at a time.");
  ASSERT(max_packet_size <= XHCI_EP_CTX_MAX_PACKET_SIZE, "xhci_open_endpoint: Packet size too big.")

  uint8_t ep_num = address & 0x0F;
  uint8_t ep_dir = (address >> 7) & 1;
  uint8_t dci = ep_num * 2 + ep_dir;

  LOG_LN_INFO("Setting up endpoint ring...");
  RING_WRAP(ep);
  XHCI_TRB_CYCLE_W(&ep_r[RING_LENGTH], ep_r_cycle);

  LOG_LN_INFO("Configuring endpoint...");
  memory_set(0, sizeof(xhci_input_ctx_t), &input_ctx);
  input_ctx.slot = device_ctx.slot;
  input_ctx.control.add_flags = (1 << 0) | (1 << dci);
  XHCI_SLOT_CTX_CTX_ENTRIES_W(&input_ctx.slot, dci);

  XHCI_EP_CTX_EP_TYPE_W(&input_ctx.ep[dci - 1], type);
  XHCI_EP_CTX_MAX_PACKET_SIZE_W(&input_ctx.ep[dci - 1], max_packet_size);
  XHCI_EP_CTX_INTERVAL_W(&input_ctx.ep[dci - 1], interval);
  input_ctx.ep[dci - 1].dequeue = (uint64_t)VIRT_TO_PHYS(ep_r) | ep_r_cycle;

  xhci_trb_t cmd_ce;
  memory_set(0, sizeof(xhci_trb_t), &cmd_ce);
  XHCI_TRB_TRB_TYPE_W(&cmd_ce, XHCI_TRB_TYPE_COMMAND_CONFIGURE_ENDPOINT);
  cmd_ce.parameter = (uint64_t)VIRT_TO_PHYS(&input_ctx);
  XHCI_TRB_SLOT_ID_W(&cmd_ce, slot);
  command_exec_sync(&cmd_ce);

  ep_r_attached_endpoint = dci;
  ep_transfer_size = max_packet_size;

  return dci;
}

uint16_t xhci_read_endpoint(xhci_slot_t slot, xhci_endpoint_t endpoint, uint16_t max, void *to) {
  LOG_LN_STEP("Reading from endpoint %hhx...", endpoint);

  ASSERT(slot < MAX_SLOTS, "xhci_control_transfer: Slot does not exist.");
  ASSERT(tsf_r_attached_slot == slot, "xhci_read_endpoint: Slot is not attached.");
  ASSERT(ep_r_attached_endpoint == endpoint, "xhci_read_endpoint: Endpoint not open.");

  if (ep_processed_events == RING_LENGTH) {
    LOG_LN_STEP("Event batch finished, posting new transfer requests...");
    ep_processed_events = 0;

    ASSERT(ep_r_i == 0, "xhci_read_endpoint: Endpoint ring cycle broken.");
    for (uint64_t i = 0; i < RING_LENGTH; i++) {
      xhci_trb_t *ere = &ep_r[i];

      ere->parameter = (uint64_t)VIRT_TO_PHYS(&ep_buffer[i]);
      ere->status = ep_transfer_size;
      XHCI_TRB_TRB_TYPE_W(ere, XHCI_TRB_TYPE_TRANSFER_NORMAL);
      XHCI_TRB_IOC_W(ere, 1);
      XHCI_TRB_CYCLE_W(ere, ep_r_cycle);
      RING_ADVANCE(ep);
    }

    db_regs[tsf_r_attached_slot] = ep_r_attached_endpoint;

    return 0;
  }

  uint16_t transferred = 0;
  while (XHCI_TRB_CYCLE_R(&evt_r[evt_r_i]) == evt_r_cycle && !transferred) {
    xhci_trb_t *cmp = &evt_r[evt_r_i];

    if (XHCI_TRB_TRB_TYPE_R(cmp) != XHCI_TRB_TYPE_EVENT_TRANSFER_COMPLETION) {
      LOG_LN_STEP("Received irrelevant event, skipping...");
      LOG_VAL_TRACE(cmp->parameter, "%lx");
      LOG_VAL_TRACE(cmp->status, "%x");
      LOG_VAL_TRACE(cmp->control, "%x");

    } else if (XHCI_TRB_COMPLETION_CODE(cmp) != 1) {
      LOG_LN_STEP("Received unsuccessful transfer event, skipping...");
      ep_processed_events++;

      LOG_VAL_TRACE(cmp->parameter, "%lx");
      LOG_VAL_TRACE(cmp->status, "%x");
      LOG_VAL_TRACE(cmp->control, "%x");
    } else {
      LOG_LN_STEP("Received transfer event, returning...");
      ep_processed_events++;

      LOG_VAL_TRACE(cmp->parameter, "%lx");
      LOG_VAL_TRACE(cmp->status, "%x");
      LOG_VAL_TRACE(cmp->control, "%x");

      xhci_trb_t *cmd = PHYS_TO_VIRT(cmp->parameter);
      uint16_t size = max < cmd->status ? max : (uint16_t)cmd->status;
      memory_copy(PHYS_TO_VIRT(cmd->parameter), size, to);

      LOG_VAL_TRACE(size, "%d");
      LOG_VAL_TRACE(*(uint64_t *)to, "%lx");

      transferred = size;
    }

    RING_ADVANCE(evt);
    intr_regs->erdp = (uint64_t)VIRT_TO_PHYS(&evt_r[evt_r_i]);
  }

  return transferred;
}

void xhci_close_endpoint(xhci_slot_t slot, xhci_endpoint_t endpoint) {
  LOG_LN_INFO("Closing endpoind %hhx...", endpoint);

  ASSERT(slot < MAX_SLOTS, "xhci_control_transfer: Slot does not exist.");
  ASSERT(tsf_r_attached_slot == slot, "xhci_close_endpoint: Slot is not attached.");
  ASSERT(ep_r_attached_endpoint == endpoint, "xhci_close_endpoint: Endpoint not open.");

  ep_r_attached_endpoint = NUL;

  xhci_trb_t cmd_se;
  memory_set(0, sizeof(xhci_trb_t), &cmd_se);
  XHCI_TRB_TRB_TYPE_W(&cmd_se, XHCI_TRB_TYPE_COMMAND_STOP_ENDPOINT);
  XHCI_TRB_SLOT_ID_W(&cmd_se, slot);
  XHCI_TRB_ENDPOINT_ID_W(&cmd_se, endpoint);
  command_exec_sync(&cmd_se);
}

void xhci_detach(xhci_slot_t slot) {
  LOG_LN_STEP("Detaching slot %d...", slot);

  ASSERT(slot < MAX_SLOTS, "xhci_control_transfer: Slot does not exist.");
  ASSERT(tsf_r_attached_slot == slot, "xhci_detach: Slot not attached.");
  ASSERT(!ep_r_attached_endpoint, "xhci_detach: Can not detach slot while in polling mode.");

  xhci_trb_t cmd_ds;
  memory_set(0, sizeof(xhci_trb_t), &cmd_ds);
  XHCI_TRB_TRB_TYPE_W(&cmd_ds, XHCI_TRB_TYPE_COMMAND_DISABLE_SLOT);
  XHCI_TRB_SLOT_ID_W(&cmd_ds, slot);
  command_exec_sync(&cmd_ds);

  dcbaa[slot] = 0;
  tsf_r_attached_slot = 0;

  LOG_LN_STEP("Done.");
}