#include <amd64/io.h>
#include <device/pci.h>
#include <libk/std.h>
#include <sys/debug.h>

static uint32_t pci_read32 (uint8_t bus, uint8_t slot, uint8_t func, uint8_t offset) {
  uint32_t addr = (uint32_t)((uint32_t)bus << 16) | ((uint32_t)slot << 11) | ((uint32_t)func << 8) |
      (offset & 0xFc) | ((uint32_t)0x80000000);
  outl (PCI_CONFIG_ADDR, addr);
  return inl (PCI_CONFIG_DATA);
}

static void pci_write32 (uint8_t bus, uint8_t slot, uint8_t func, uint8_t offset, uint32_t value) {
  uint32_t addr = (uint32_t)((uint32_t)bus << 16) | ((uint32_t)slot << 11) | ((uint32_t)func << 8) |
      (offset & 0xFc) | ((uint32_t)0x80000000);
  outl (PCI_CONFIG_ADDR, addr);
  outl (PCI_CONFIG_DATA, value);
}

static void pci_check_bus (uint8_t bus, pci_cb_func_t cb);

static void pci_check_func (uint8_t bus, uint8_t slot, uint8_t func, pci_cb_func_t cb) {
  uint32_t reg0 = pci_read32 (bus, slot, func, PCI_VENDOR_ID);
  uint16_t vendor = (uint16_t)(reg0 & 0xFFFF);

  if (vendor == 0xFFFF)
    return;

  uint32_t reg8 = pci_read32 (bus, slot, func, PCI_REVISION_ID);

  struct pci_dev pci_dev = {
      .bus = bus,
      .slot = slot,
      .func = func,
      .vendor = vendor,
      .device = ((uint16_t)(reg0 >> 16)),
      .classcode = ((uint8_t)(reg8 >> 24)),
      .subclass = ((uint8_t)(reg8 >> 16)),
  };

  cb (pci_dev);

  /* PCI 2 PCI bridge */
  if (pci_dev.classcode == 0x06 && pci_dev.subclass == 0x04) {
    uint32_t reg18 = pci_read32 (bus, slot, func, 0x18);
    uint8_t secondary = (uint8_t)(reg18 >> 8);
    pci_check_bus (secondary, cb);
  }
}

static void pci_check_device (uint8_t bus, uint8_t slot, pci_cb_func_t cb) {
  uint32_t reg0 = pci_read32 (bus, slot, 0, PCI_VENDOR_ID);

  if ((uint16_t)(reg0 & 0xFFFF) == 0xFFFF)
    return;

  pci_check_func (bus, slot, 0, cb);

  /* multifunc device */
  uint32_t reg0xc = pci_read32 (bus, slot, 0, PCI_CACHELINE);
  if ((reg0xc >> 16) & 0x80) {
    for (uint8_t func = 1; func < 8; func++)
      pci_check_func (bus, slot, func, cb);
  }
}

static void pci_check_bus (uint8_t bus, pci_cb_func_t cb) {
  for (uint8_t slot = 0; slot < 32; slot++)
    pci_check_device (bus, slot, cb);
}

static void pci_enumerate (pci_cb_func_t cb) {
  uint32_t reg0xc = pci_read32 (0, 0, 0, PCI_CACHELINE);
  bool is_multictrl = (reg0xc >> 16) & 0x80;

  if (!is_multictrl)
    pci_check_bus (0, cb);
  else {
    for (uint8_t func = 0; func < 8; func++) {
      if ((pci_read32 (0, 0, func, PCI_VENDOR_ID) & 0xFFFF) != 0xFFFF)
        pci_check_bus (func, cb);
    }
  }
}

static void pci_string_identifiers (uint16_t vendor_id, uint16_t device_id, uint8_t class_id,
                                    uint8_t subclass_id, const char** vname, const char** dname,
                                    const char** cname) {
  *vname = "Unknown vendor";
  *dname = "Unknown device";
  *cname = "Unknown class";

  for (size_t i = 0; pci_vendors[i].name; i++) {
    if (pci_vendors[i].id == vendor_id) {
      *vname = pci_vendors[i].name;
      break;
    }
  }

  for (size_t i = 0; pci_device_names[i].name; i++) {
    if (pci_device_names[i].vendor_id == vendor_id && pci_device_names[i].device_id == device_id) {
      *dname = pci_device_names[i].name;
      break;
    }
  }

  for (size_t i = 0; pci_classes[i].name; i++) {
    if (pci_classes[i].class == class_id && pci_classes[i].subclass == subclass_id) {
      *cname = pci_classes[i].name;
      break;
    }
  }
}

static void pci_discovery_cb (struct pci_dev pci_dev) {
  const char *vname, *dname, *cname;
  pci_string_identifiers (pci_dev.vendor, pci_dev.device, pci_dev.classcode, pci_dev.subclass,
                          &vname, &dname, &cname);

  DEBUG ("PCI DEVICE: %04x:%04x at %02d:%02d:%02d / %s; %s; %s\n", pci_dev.vendor, pci_dev.device,
         pci_dev.bus, pci_dev.slot, pci_dev.func, vname, dname, cname);
}

void pci_init (void) { pci_enumerate (&pci_discovery_cb); }