#include <device/def_device_op.h>
#include <device/device.h>
#include <device/xhci.h>
#include <libk/list.h>
#include <libk/std.h>
#include <libk/string.h>
#include <mm/malloc.h>
#include <proc/proc.h>
#include <proc/reschedule.h>
#include <proc/suspension_q.h>
#include <sys/debug.h>
#include <sys/spin_lock.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

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

DEFINE_DEVICE_INIT (xhci_init) {
  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;

  device->udata = xhci;

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

  DEBUG ("starting init sequence\n");

  /* assert CNR is 0 */
  while (xhci_read32 (xhci->xhci_oper_base, XHCI_USBSTS) & (1 << 11))
    spin_lock_relax ();

  /* STOP */
  uint32_t usbcmd = xhci_read32 (xhci->xhci_oper_base, XHCI_USBCMD);
  xhci_write32 (xhci->xhci_oper_base, XHCI_USBCMD, usbcmd & ~0x01);

  /* wait for HCH bit */
  int timeout = 100000;
  while (!(xhci_read32 (xhci->xhci_oper_base, XHCI_USBSTS) & (1 << 12))) {
    if (--timeout == 0)
      break;

    spin_lock_relax ();
  }

  /* RESET */
  xhci_write32 (xhci->xhci_oper_base, XHCI_USBCMD, (1 << 1));

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

  /* Stall while controller not ready */
  while (xhci_read32 (xhci->xhci_oper_base, XHCI_USBSTS) & (1 << 11))
    spin_lock_relax ();

  DEBUG ("XHCI init done\n");

  return true;
}

DEFINE_DEVICE_FINI (xhci_fini) {
  struct xhci* xhci = device->udata;

  free (xhci);
}