#include "pvec.h"
#include "io.h"
#include "pci.h"
#include "string.h"
#include "memory.h"
#include "console.h"

struct pci_device **pci_devices;
int pci_device_count;
#define PCI_CONF_READ8(bus, dev, func, reg) (outl(PCI_CONFIG_ADDR, PCI_FORMAT(bus, dev, func, reg)), inb(PCI_CONFIG_DATA + ((reg)&3)))

void pci_config_write_char(unsigned short bus, unsigned short slot, unsigned short func, unsigned short offset, unsigned char val) {
  unsigned long address;
  unsigned long lbus = (unsigned long)bus;
  unsigned long lslot = (unsigned long)slot;
  unsigned long lfunc = (unsigned long)func;

  address = (unsigned long)((lbus << 16) | (lslot << 11) | (lfunc << 8) | (offset & 0xfc) | ((unsigned long)0x80000000));

  outportl(0xCF8, address);
  outportb(0xCFC, val);
}

unsigned short pci_config_read_char(unsigned short bus, unsigned short slot, unsigned short func, unsigned short offset) {
  unsigned long address;
  unsigned long lbus = (unsigned long)bus;
  unsigned long lslot = (unsigned long)slot;
  unsigned long lfunc = (unsigned long)func;
  unsigned char tmp = 0;

  address = (unsigned long)((lbus << 16) | (lslot << 11) | (lfunc << 8) | (offset & 0xfc) | ((unsigned long)0x80000000));

  outportl(0xCF8, address);
  tmp = (unsigned char)((inportl(0xCFC) >> ((offset & 3) * 8)) & 0xff);

  return tmp;
}

unsigned long pci_config_read_dword(unsigned short bus, unsigned short slot, unsigned short func, unsigned short offset) {
  unsigned long address;
  unsigned long lbus = (unsigned long)bus;
  unsigned long lslot = (unsigned long)slot;
  unsigned long lfunc = (unsigned long)func;
  unsigned long tmp = 0;

  address = (unsigned long)((lbus << 16) | (lslot << 11) | (lfunc << 8) | (offset & 0xfc) | ((unsigned long)0x80000000));

  outportl(0xCF8, address);
  tmp = inportl(0xCFC);

  return tmp;
}

unsigned short pci_config_read_word(unsigned short bus, unsigned short slot, unsigned short func, unsigned short offset) {
  unsigned long address;
  unsigned long lbus = (unsigned long)bus;
  unsigned long lslot = (unsigned long)slot;
  unsigned long lfunc = (unsigned long)func;
  unsigned short tmp = 0;

  address = (unsigned long)((lbus << 16) | (lslot << 11) | (lfunc << 8) | (offset & 0xfc) | ((unsigned long)0x80000000));

  outportl(0xCF8, address);
  tmp = (unsigned short)((inportl(0xCFC) >> ((offset & 2) * 8)) & 0xffff);

  return tmp;
}

void pci_config_write_word(unsigned short bus, unsigned short slot, unsigned short func, unsigned short offset, unsigned short value) {
  unsigned long address;
  unsigned long lbus = (unsigned long)bus;
  unsigned long lslot = (unsigned long)slot;
  unsigned long lfunc = (unsigned long)func;

  address = (unsigned long)((lbus << 16) | (lslot << 11) | (lfunc << 8) | (offset & 0xfc) | ((unsigned long)0x80000000));

  outportl(0xCF8, address);
  outportl(0xCFC, value);

  return;
}

void pci_config_write_dword(unsigned short bus, unsigned short slot, unsigned short func, unsigned short offset, unsigned long value) {
  unsigned long address;
  unsigned long lbus = (unsigned long)bus;
  unsigned long lslot = (unsigned long)slot;
  unsigned long lfunc = (unsigned long)func;

  address = (unsigned long)((lbus << 16) | (lslot << 11) | (lfunc << 8) | (offset & 0xfc) | ((unsigned long)0x80000000));

  outportl(0xCF8, address);
  outportl(0xCFC, value);

  return;
}

void pci_read_sizes_00(struct pci_device *dev) {
  pci_config_write_dword(dev->bus, dev->slot, dev->func, 0x10, 0xffffffff);
  dev->size[0] = ~(pci_config_read_dword(dev->bus, dev->slot, dev->func, 0x10) & ~0xf) + 1;
  pci_config_write_dword(dev->bus, dev->slot, dev->func, 0x10, dev->base[0]);
  pci_config_write_dword(dev->bus, dev->slot, dev->func, 0x14, 0xffffffff);
  dev->size[1] = ~(pci_config_read_dword(dev->bus, dev->slot, dev->func, 0x14) & ~0xf) + 1;
  pci_config_write_dword(dev->bus, dev->slot, dev->func, 0x14, dev->base[1]);
  pci_config_write_dword(dev->bus, dev->slot, dev->func, 0x18, 0xffffffff);
  dev->size[2] = ~(pci_config_read_dword(dev->bus, dev->slot, dev->func, 0x18) & ~0xf) + 1;
  pci_config_write_dword(dev->bus, dev->slot, dev->func, 0x18, dev->base[2]);
  pci_config_write_dword(dev->bus, dev->slot, dev->func, 0x1C, 0xffffffff);
  dev->size[3] = ~(pci_config_read_dword(dev->bus, dev->slot, dev->func, 0x1C) & ~0xf) + 1;
  pci_config_write_dword(dev->bus, dev->slot, dev->func, 0x1C, dev->base[3]);
  pci_config_write_dword(dev->bus, dev->slot, dev->func, 0x20, 0xffffffff);
  dev->size[4] = ~(pci_config_read_dword(dev->bus, dev->slot, dev->func, 0x20) & ~0xf) + 1;
  pci_config_write_dword(dev->bus, dev->slot, dev->func, 0x20, dev->base[4]);
  pci_config_write_dword(dev->bus, dev->slot, dev->func, 0x24, 0xffffffff);
  dev->size[5] = ~(pci_config_read_dword(dev->bus, dev->slot, dev->func, 0x24) & ~0xf) + 1;
  pci_config_write_dword(dev->bus, dev->slot, dev->func, 0x24, dev->base[5]);
}

void pci_read_bases_00(struct pci_device *dev) {
  dev->base[0] = pci_config_read_dword(dev->bus, dev->slot, dev->func, 0x10);
  dev->base[1] = pci_config_read_dword(dev->bus, dev->slot, dev->func, 0x14);
  dev->base[2] = pci_config_read_dword(dev->bus, dev->slot, dev->func, 0x18);
  dev->base[3] = pci_config_read_dword(dev->bus, dev->slot, dev->func, 0x1C);
  dev->base[4] = pci_config_read_dword(dev->bus, dev->slot, dev->func, 0x20);
  dev->base[5] = pci_config_read_dword(dev->bus, dev->slot, dev->func, 0x24);
}

void init_pci(void) {
  int bus;
  int slot;
  int func;

  struct pci_device tmp_device;
  pci_device_count = 0;

  for (bus = 0; bus < 4; bus++) {
    for (slot = 0; slot < 32; slot++) {
      for (func = 0; func < 8; func++) {

        memset(&tmp_device, 0, sizeof(struct pci_device));

        if ((tmp_device.vendor = pci_config_read_word(bus, slot, 0, 0)) != 0xFFFF) {
          tmp_device.header_type = pci_config_read_word(bus, slot, func, 14) & 0x00FF;
          if (tmp_device.header_type == 0xff)
            continue;
          tmp_device.device = pci_config_read_word(bus, slot, func, 2);
          tmp_device.classtype = (pci_config_read_word(bus, slot, func, 10) >> 8) & 0x00FF;
          tmp_device.subclasstype = pci_config_read_word(bus, slot, func, 10) & 0x00FF;
          tmp_device.bus = bus;
          tmp_device.slot = slot;
          tmp_device.func = func;
          if ((tmp_device.header_type & 0x3F) == 0x00) {
            pci_read_bases_00(&tmp_device);
            pci_read_sizes_00(&tmp_device);
            tmp_device.irq = pci_config_read_word(bus, slot, func, 0x3c) & 0x00FF;
          }
          if (pci_device_count == 0) {
            pci_devices = (struct pci_device **)malloc(sizeof(struct pci_device *));
          } else {
            pci_devices = (struct pci_device **)realloc(pci_devices, sizeof(struct pci_device *) * (pci_device_count + 1));
          }

          pci_devices[pci_device_count] = (struct pci_device *)malloc(sizeof(struct pci_device));

          memcpy((char *)pci_devices[pci_device_count], (char *)&tmp_device, sizeof(struct pci_device));

          pci_device_count++;
        }
      }
    }
  }

  kprintf("Found %d PCI devices\n", pci_device_count);
}

int pci_find_device(unsigned char classt, unsigned char subclasst, struct pci_device **device, int offset) {
  int i;
  int count = 0;
  for (i = 0; i < pci_device_count; i++) {
    if (pci_devices[i]->classtype == classt && pci_devices[i]->subclasstype == subclasst) {
      if (count < offset) {
        count++;
        continue;
      }
      *device = pci_devices[i];
      return 1;
    }
  }
  return 0;
}

int pci_find_device_by_vendor(unsigned short vendort, unsigned short devicet, struct pci_device **device, int offset) {
  int i;
  int count = 0;
  for (i = 0; i < pci_device_count; i++) {
    if (pci_devices[i]->vendor == vendort && pci_devices[i]->device == devicet) {
      if (count < offset) {
        count++;
        continue;
      }
      *device = pci_devices[i];
      return 1;
    }
  }
  return 0;
}

int pci_find_device_by_vendor_virtio(unsigned short subsystemt, struct pci_device **device, int offset) {
  int i;
  int count = 0;
  for (i = 0; i < pci_device_count; i++) {
    if (pci_devices[i]->vendor == 0x1AF4 && pci_devices[i]->device >= 0x1000 && pci_devices[i]->device <= 0x103F) {
      unsigned short subsystem = pci_config_read_word(pci_devices[i]->bus, pci_devices[i]->slot, pci_devices[i]->func, 0x2e);

      if (subsystem == subsystemt) {
        if (count < offset) {
          count++;
          continue;
        }
        *device = pci_devices[i];
        return 1;
      }
    }
  }
  return 0;
}

void pci_set_master(struct pci_device *dev, int enable) {
  unsigned short command;
  unsigned short oldcommand = pci_config_read_word(dev->bus, dev->slot, dev->func, 0x04);

  if (enable) {
    command = oldcommand | PCI_COMMAND_MASTER_ENABLE;
  } else {
    command = oldcommand & ~PCI_COMMAND_MASTER_ENABLE;
  }

  if (command != oldcommand) {
    pci_config_write_word(dev->bus, dev->slot, dev->func, PCI_COMMAND_STATUS_REG, command);
  }
}

void pci_set_io_enable(struct pci_device *dev, int enable) {
  unsigned short command;
  unsigned short oldcommand = pci_config_read_word(dev->bus, dev->slot, dev->func, 0x04);

  if (enable) {
    command = oldcommand | PCI_COMMAND_IO_ENABLE;
  } else {
    command = oldcommand & ~PCI_COMMAND_IO_ENABLE;
  }

  if (command != oldcommand) {
    pci_config_write_word(dev->bus, dev->slot, dev->func, PCI_COMMAND_STATUS_REG, command);
  }
}

void pci_set_mem_enable(struct pci_device *dev, int enable) {
  unsigned short command;
  unsigned short oldcommand = pci_config_read_word(dev->bus, dev->slot, dev->func, 0x04);

  if (enable) {
    command = oldcommand | PCI_COMMAND_MEM_ENABLE;
  } else {
    command = oldcommand & ~PCI_COMMAND_MEM_ENABLE;
  }

  if (command != oldcommand) {
    pci_config_write_word(dev->bus, dev->slot, dev->func, PCI_COMMAND_STATUS_REG, command);
  }
}

int pci_fill_capabilities(struct pci_device *dev) {
  unsigned char next;

  pci_set_mem_enable(dev, 1);
  pci_set_io_enable(dev, 1);
  pci_set_master(dev, 1);

  unsigned short status = pci_config_read_word(dev->bus, dev->slot, dev->func, PCI_COMMAND_STATUS_REG + 2);

  // Check if the device has a capabalities list
  if (!(status & PCI_COMMAND_CAPABALITES_LIST)) {
    kprintf("No dev capabilities\n");
    return -1;
  }

  unsigned long cap_register = pci_config_read_dword(dev->bus, dev->slot, dev->func, PCI_CAP_REG);
  unsigned char cap_pointer = PCI_CAP_POINTER(cap_register) & PCI_CAP_MASK;

  while (cap_pointer) {
    next = pci_config_read_char(dev->bus, dev->slot, dev->func, cap_pointer + PCI_CAP_NEXT) & PCI_CAP_MASK;
    unsigned char type = pci_config_read_char(dev->bus, dev->slot, dev->func, cap_pointer + PCI_CAP_CFG_TYPE);
    unsigned char bar = pci_config_read_char(dev->bus, dev->slot, dev->func, cap_pointer + PCI_CAP_BAR);
    unsigned long offset = pci_config_read_dword(dev->bus, dev->slot, dev->func, cap_pointer + PCI_CAP_OFF);

    // Location of the given capability in the PCI config space.
    dev->cap[type] = cap_pointer;
    dev->cap_bar[type] = bar;
    dev->cap_off[type] = offset;

    // cprintf("cap type: %d pointer: %p\n", type, cap_pointer);

    cap_pointer = next;
  }

  return 0;
}

struct pci_report_t {
  unsigned short vendor;
  unsigned short device;
  unsigned short bus;
  unsigned short slot;
  unsigned short func;
  unsigned char irq;
};

int pci_dev_info(unsigned char *buffer, int len, int last_pci_dev) {
  struct pci_report_t *pciinfo = (struct pci_report_t *)buffer;
  int ret = 0;

  for (int i = last_pci_dev; ret < len / sizeof(struct pci_report_t) && i < pci_device_count; i++) {
    if (i > last_pci_dev) {
      pciinfo[ret].vendor = pci_devices[i]->vendor;
      pciinfo[ret].device = pci_devices[i]->device;
      pciinfo[ret].bus = pci_devices[i]->bus;
      pciinfo[ret].slot = pci_devices[i]->slot;
      pciinfo[ret].func = pci_devices[i]->func;
      pciinfo[ret].irq = pci_devices[i]->irq;
      ret++;
    }
  }

  return ret;
}