#include <stdint.h>
#include <stddef.h>
#include "pci/pci.h"
#include "pci/ata/ata.h"
#include "pci/serial/serial.h"
#include "pci/qemu/testdev/testdev.h"
#include "io/io.h"
#include "std/string.h"
#include "util/util.h"
#include "kprintf.h"

static PciDev PCI_DEV_ZERO = {0};

uint32_t pci_read32(PciDev dev, uint32_t field) {
  dev.fieldnum = (field & 0xFC) >> 2;
  dev.enable = 1;
  io_out32(PCI_CFG_ADDR, dev.bits);
  
  return io_in32(PCI_CFG_DATA);
}

uint16_t pci_read16(PciDev dev, uint32_t field) {
  dev.fieldnum = (field & 0xFC) >> 2;
  dev.enable = 1;
  io_out32(PCI_CFG_ADDR, dev.bits);

  return io_in16(PCI_CFG_DATA + (field & 0x2));
}

uint8_t pci_read8(PciDev dev, uint32_t field) {
  dev.fieldnum = (field & 0xFC) >> 2;
  dev.enable = 1;
  io_out32(PCI_CFG_ADDR, dev.bits);

  return io_in8(PCI_CFG_DATA + (field & 0x3));
}

void pci_write32(PciDev dev, uint32_t field, uint32_t v) {
  dev.fieldnum = (field & 0xFC) >> 2;
  dev.enable = 1;
  io_out32(PCI_CFG_ADDR, dev.bits);
  io_out32(PCI_CFG_DATA, v);
}

void pci_write16(PciDev dev, uint32_t field, uint16_t v) {
  dev.fieldnum = (field & 0xFC) >> 2;
  dev.enable = 1;
  io_out32(PCI_CFG_ADDR, dev.bits);
  io_out16(PCI_CFG_DATA, v);
}

void pci_write8(PciDev dev, uint32_t field, uint8_t v) {
  dev.fieldnum = (field & 0xFC) >> 2;
  dev.enable = 1;
  io_out32(PCI_CFG_ADDR, dev.bits);
  io_out8(PCI_CFG_DATA, v);
}

void pci_readbar(PciDev dev, uint32_t bar, uint32_t *addr, uint32_t *mask) {
  *addr = pci_read32(dev, bar);
  pci_write32(dev, bar, 0xffffffff);
  *mask = pci_read32(dev, bar);
  pci_write32(dev, bar, *addr);
}

void pci_getbar(PciDev dev, PciBar *bar, uint32_t barid) {
  uint32_t addrlo;
  uint32_t masklo;
  pci_readbar(dev, barid, &addrlo, &masklo);

  if (addrlo & PCI_BAR_IO) {
    bar->x.io.iobase = (uint16_t)(addrlo & ~0x3);
    bar->size = (uint16_t)(~(masklo & ~0x3) + 1);
    bar->flags = (addrlo & 0x3);
  } else if (addrlo & PCI_BAR_MEM32) {
    bar->x.mem.addr = (uint64_t)(addrlo & ~0xF);
    bar->size = ~(masklo & ~0xF) + 1;
    bar->flags = (addrlo & 0xF);
  } else if (addrlo & PCI_BAR_MEM64) {
    uint32_t addrhi;
    uint32_t maskhi;
    pci_readbar(dev, barid+4, &addrhi, &maskhi);

    bar->x.mem.addr = (uint64_t)(((uint64_t)addrhi << 32) | (addrlo & ~0xF));
    bar->size = ~(((uint64_t)maskhi << 32) | (masklo & ~0xF)) + 1;
    bar->flags = (addrlo & 0xF);
  }
}

uint32_t pci_devtype(PciDev dev) {
  uint32_t a = pci_read8(dev, PCI_CLASS) << 8;
  uint32_t b = pci_read8(dev, PCI_SUBCLASS);
  return a | b;
}

uint32_t pci_scndrybus(PciDev dev) {
  return pci_read8(dev, PCI_SCNDRY_BUS);
}

uint32_t pci_isend(PciDev dev) {
  return !(pci_read8(dev, PCI_HDRTYPE));
}

PciDev pci_scanfn(uint16_t vendorid, uint16_t deviceid, uint32_t bus,
                  uint32_t device, uint32_t fn, int devtype) {
  PciDev dev; memset(&dev, 0, sizeof(dev));
  dev.busnum = bus;
  dev.devnum = device;
  dev.fnnum = fn;

  // bridge
  if (pci_devtype(dev) == 0x0604) {
    pci_scanbus(vendorid, deviceid, pci_scndrybus(dev), devtype);
  }
  if (devtype == -1 || (uint32_t)devtype == pci_devtype(dev)) {
    uint32_t venid = pci_read16(dev, PCI_VENDORID);
    uint32_t devid = pci_read16(dev, PCI_DEVICEID);

    if (devid == deviceid && venid == vendorid) {
      return dev;
    }
  }
  return PCI_DEV_ZERO;
}

PciDev pci_scanbus(uint16_t vendorid, uint16_t deviceid,
                   uint32_t bus, int devtype) {
  for (uint32_t device = 0; device < PCI_DEV_PER_BUS; device++) {
    PciDev d = pci_scandev(vendorid, deviceid, bus, device, devtype);
    if (d.bits) {
      return d;
    }
  }
  return PCI_DEV_ZERO;
}

PciDev pci_scandev(uint16_t vendorid, uint16_t deviceid,
                  uint32_t bus, uint32_t device, int devtype) {
  PciDev dev; memset(&dev, 0, sizeof(dev));
  dev.busnum = bus;
  dev.devnum = device;

  if (pci_read16(dev, PCI_VENDORID) == 0xFFFF) {
    return PCI_DEV_ZERO;
  }

  PciDev d = pci_scanfn(vendorid, deviceid, bus, device, 0, devtype);
  if (d.bits) {
    return d;
  }

  if (pci_isend(dev)) {
    return PCI_DEV_ZERO;
  }

  for (uint32_t fn = 1; fn < PCI_FN_PER_DEV; fn++) {
    if (pci_read16(dev, PCI_VENDORID) != 0xFFFF) {
      d = pci_scanfn(vendorid, deviceid, bus, device, fn, devtype);
      if (d.bits) {
        return d;
      }
    }
  }
  return PCI_DEV_ZERO;
}

PciDev pci_getdev(uint16_t vendorid, uint16_t deviceid, int devtype) {
  PciDev d = pci_scanbus(vendorid, deviceid, 0, devtype);
  if (d.bits) {
    return d;
  }

  if (pci_isend(PCI_DEV_ZERO)) {
    ERR("pci", "pci_getdev() failed for 0x%04x/0x%04x,%d\n", vendorid, deviceid, devtype);
    return PCI_DEV_ZERO;
  }

  for (uint32_t fn = 1; fn < PCI_FN_PER_DEV; fn++) {
    PciDev d2; memset(&d2, 0, sizeof(d2));
    d2.fnnum = fn;

    if (pci_read16(d2, PCI_VENDORID) == 0xFFFF) {
      break;
    }

    d = pci_scanbus(vendorid, deviceid, fn, devtype);
    if (d.bits) {
      return d;
    }
  }

  return PCI_DEV_ZERO;
}

PciInitFn PCI_INIT_ARRAY[PCI_INIT_ARRAY_MAX] = {
  &pci_ata_init,
  &pci_serial_init,
  &pci_qemu_testdev_init,
};

void pci_init_devs(void) {
  for (size_t i = 0; i < LEN(PCI_INIT_ARRAY); i++) {
    if (PCI_INIT_ARRAY[i] != NULL)
      PCI_INIT_ARRAY[i]();
  }
}

void pci_init(void) {
  pci_init_devs();
}