#include <stdint.h>
#include <stddef.h>
#include "pci/pci.h"
#include "pci/ata/ata.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_SIZE_MAP[100] = {0};

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

  uint32_t size = PCI_SIZE_MAP[field];
  uint8_t u8; uint16_t u16; uint32_t u32;
  switch (size) {
    case 1:
      u8 = io_in8(PCI_CFG_DATA + (field & 0x3));
      return (uint32_t)u8;
    case 2:
      u16 = io_in16(PCI_CFG_DATA + (field & 0x2));
      return (uint32_t)u16;
    case 4:
      u32 = io_in32(PCI_CFG_DATA);
      return u32;
    default:
      return 0xFFFF;
  }
}

void pci_write(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);
}

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

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

uint32_t pci_isend(PciDev dev) {
  return !(pci_read(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_read(dev, PCI_VENDORID);
    uint32_t devid = pci_read(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_read(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_read(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\n");
  }

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

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

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

  return PCI_DEV_ZERO;
}

void pci_init_size_map(void) {
  PCI_SIZE_MAP[PCI_VENDORID]    = 2;
  PCI_SIZE_MAP[PCI_DEVICEID]    = 2;
  PCI_SIZE_MAP[PCI_CMD]         = 2;
  PCI_SIZE_MAP[PCI_STATUS]      = 2;
  PCI_SIZE_MAP[PCI_SUBCLASS]    = 1;
  PCI_SIZE_MAP[PCI_CLASS]       = 1;
  PCI_SIZE_MAP[PCI_CACHELINESZ] = 1;
  PCI_SIZE_MAP[PCI_LTNCY_TIMER] = 1;
  PCI_SIZE_MAP[PCI_HDRTYPE]     = 1;
  PCI_SIZE_MAP[PCI_BIST]        = 1;
  PCI_SIZE_MAP[PCI_BAR0]        = 4;
  PCI_SIZE_MAP[PCI_BAR1]        = 4;
  PCI_SIZE_MAP[PCI_BAR2]        = 4;
  PCI_SIZE_MAP[PCI_BAR3]        = 4;
  PCI_SIZE_MAP[PCI_BAR4]        = 4;
  PCI_SIZE_MAP[PCI_BAR5]        = 4;
  PCI_SIZE_MAP[PCI_INTRLINE]    = 1;
  PCI_SIZE_MAP[PCI_SCNDRY_BUS]  = 1;
}


PciInitFn PCI_INIT_ARRAY[PCI_INIT_ARRAY_MAX] = {
  &pci_ata_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_size_map();
  pci_init_devs();
}