#include <amd64/apic.h>
#include <amd64/intr_defs.h>
#include <device/def_device_op.h>
#include <device/device.h>
#include <device/pci/pci.h>
#include <device/pci/pci_info.h>
#include <device/pci/pci_xhci.h>
#include <device/usb/xhci.h>
#include <devices.h>
#include <libk/align.h>
#include <libk/lengthof.h>
#include <libk/printf.h>
#include <libk/std.h>
#include <limine/requests.h>
#include <proc/proc.h>
#include <proc/reschedule.h>
#include <sys/debug.h>
#include <sys/mm.h>

static bool xhci_init_done = false;

bool pci_xhci_init(struct proc* proc, struct reschedule_ctx* rctx, struct pci_info pci_info) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;

  uint8_t progif = pci_read8(pci_info.bus, pci_info.slot, pci_info.func, PCI_PROG_IF);

  /* not an XHCI controller */
  if (progif != 0x30) {
    return true;
  }

  if (xhci_init_done) {
    DEBUG("Cannot initialize more XHCI controllers\n");
    return false;
  }

  uint16_t pci_cmd = pci_read16(pci_info.bus, pci_info.slot, pci_info.func, PCI_COMMAND);

  uint16_t new_cmd = pci_cmd;
  new_cmd |= (1 << PCI_CMD_MEMSPACE);
  new_cmd |= (1 << PCI_CMD_BUSMASTER);
  new_cmd &= ~(1 << PCI_CMD_INTRDISABLE);

  if (pci_cmd != new_cmd) {
    pci_write16(pci_info.bus, pci_info.slot, pci_info.func, PCI_COMMAND, new_cmd);
  }

  uintptr_t xhci_phys = 0;
  size_t map_pages = 0;

  uint32_t bar0 = pci_read32(pci_info.bus, pci_info.slot, pci_info.func, PCI_BAR0);
  map_pages = pci_get_bar_size(pci_info.bus, pci_info.slot, pci_info.func, PCI_BAR0);
  map_pages = div_align_up(map_pages, PAGE_SIZE);

  if ((bar0 & 0x6) == 0x4) {
    uint32_t bar1 = pci_read32(pci_info.bus, pci_info.slot, pci_info.func, PCI_BAR1);
    xhci_phys = ((uint64_t)bar1 << 32) | (bar0 & ~0xF);
    DEBUG("XHCI phys base addr=%p (64 bit)\n", xhci_phys);
  } else {
    xhci_phys = (bar0 & ~0xF);
    DEBUG("XHCI phys base addr=%p (32 bit)\n", xhci_phys);
  }

  if (xhci_phys == 0) {
    DEBUG("WARNING xhci_phys is NULL!\n");
    return false;
  }

  uintptr_t xhci_base = xhci_phys + (uintptr_t)hhdm->offset;

  DEBUG("BAR size = %zu pages\n", map_pages);

  for (size_t page = 0; page < map_pages; page++) {
    mm_map_kernel_page(xhci_phys + page * PAGE_SIZE, xhci_base + page * PAGE_SIZE,
                       MM_PG_RW | MM_PG_PRESENT | MM_PG_NOCACHE);
  }

  bool irqs_support = false;

  if (!pci_msi_init(pci_info.bus, pci_info.slot, pci_info.func, INTR_XHCI, thiscpu->lapic_id)) {
    uint8_t intr_line = pci_read8(pci_info.bus, pci_info.slot, pci_info.func, PCI_INTERRUPT_LINE);
    uint8_t intr_pin = pci_read8(pci_info.bus, pci_info.slot, pci_info.func, PCI_INTERRUPT_PIN);

    if (intr_pin != 0) {
      irqs_support = true;
      ioapic_route_irq(INTR_XHCI, intr_line, 0, thiscpu->lapic_id);
    }
  } else {
    irqs_support = true;
  }

  DEBUG("IRQ support=%d\n", irqs_support);

  if (!irqs_support) {
    return false;
  }

  struct xhci_init init = {
      .xhci_mmio_base = xhci_base,
      .irqs_support = irqs_support,
      .irq = INTR_XHCI,
  };

  device_op_func_t ops[] = {
      [XUSBCTRL_POLL_DRIVER] = &xhci_poll_driver,
  };

  device_create(DEVICE_TYPE_USB_CTRL, "xhci", ops, lengthof(ops), &xhci_init, &xhci_fini, &init,
                proc, rctx);

  xhci_init_done = true;

  return true;
}