#include <amd64/io.h>
#include <device/pci.h>
#include <device/pci_ide.h>
#include <device/pci_info.h>
#include <libk/lengthof.h>
#include <libk/std.h>
#include <sync/spin_lock.h>
#include <sys/debug.h>

static const struct pci_driver_info pci_driver_infos[] = {
    {.class = 0x01, .subclass = 0x01, .init = &pci_ide_init},
};

static spin_lock_t pci_lock = SPIN_LOCK_INIT;

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

  spin_lock (&pci_lock);

  outl (PCI_CONFIG_ADDR, addr);
  uint32_t r = inl (PCI_CONFIG_DATA);

  spin_unlock (&pci_lock);

  return r;
}

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

  spin_lock (&pci_lock);

  outl (PCI_CONFIG_ADDR, addr);
  outl (PCI_CONFIG_DATA, value);

  spin_unlock (&pci_lock);
}

uint16_t pci_read16 (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);

  spin_lock (&pci_lock);

  outl (PCI_CONFIG_ADDR, addr);
  uint16_t r = inw (PCI_CONFIG_DATA + (offset & 2));

  spin_unlock (&pci_lock);

  return r;
}

void pci_write16 (uint8_t bus, uint8_t slot, uint8_t func, uint8_t offset, uint16_t value) {
  uint32_t addr = (uint32_t)((uint32_t)bus << 16) | ((uint32_t)slot << 11) | ((uint32_t)func << 8) |
      (offset & 0xFC) | ((uint32_t)0x80000000);

  spin_lock (&pci_lock);

  outl (PCI_CONFIG_ADDR, addr);
  outw (PCI_CONFIG_DATA + (offset & 2), value);

  spin_unlock (&pci_lock);
}

uint8_t pci_read8 (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);

  spin_lock (&pci_lock);

  outl (PCI_CONFIG_ADDR, addr);
  uint8_t r = inb (PCI_CONFIG_DATA + (offset & 3));

  spin_unlock (&pci_lock);

  return r;
}

void pci_write8 (uint8_t bus, uint8_t slot, uint8_t func, uint8_t offset, uint8_t value) {
  uint32_t addr = (uint32_t)((uint32_t)bus << 16) | ((uint32_t)slot << 11) | ((uint32_t)func << 8) |
      (offset & 0xFC) | ((uint32_t)0x80000000);

  spin_lock (&pci_lock);

  outl (PCI_CONFIG_ADDR, addr);
  outb (PCI_CONFIG_DATA + (offset & 3), value);

  spin_unlock (&pci_lock);
}

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_info pci_info = {
      .bus = bus,
      .slot = slot,
      .func = func,
      .vendor = vendor,
      .device = ((uint16_t)(reg0 >> 16)),
      .class = ((uint8_t)(reg8 >> 24)),
      .subclass = ((uint8_t)(reg8 >> 16)),
  };

  cb (pci_info);

  /* PCI 2 PCI bridge */
  if (pci_info.class == 0x06 && pci_info.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_info pci_info) {
  const char *vname, *dname, *cname;
  pci_string_identifiers (pci_info.vendor, pci_info.device, pci_info.class, pci_info.subclass,
                          &vname, &dname, &cname);

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

  for (size_t driver = 0; driver < lengthof (pci_driver_infos); driver++) {
    if (pci_driver_infos[driver].class == pci_info.class &&
        pci_driver_infos[driver].subclass == pci_info.subclass) {
      pci_driver_infos[driver].init (pci_info);
    }
  }
}

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