#include <socks/kext.h>
#include <socks/device.h>
#include <socks/printk.h>
#include <socks/pci.h>
#include <socks/libc/stdio.h>
#include <arch/ports.h>

static struct driver *pci_driver = NULL;
static struct bus_device *pci_bus = NULL;

static void pci_scan_bus(pci_func_t f, int type, int bus, void *arg);

void pci_write_field(uint32_t device, int field, int size, uint32_t value)
{
	outportl(PCI_ADDRESS_PORT, pci_get_addr(device, field));
	outportl(PCI_VALUE_PORT, value);
}

uint32_t pci_read_field(uint32_t device, int field, int size)
{
	outportl(PCI_ADDRESS_PORT, pci_get_addr(device, field));

	if (size == 4) {
		uint32_t t = inportl(PCI_VALUE_PORT);
		return t;
	} else if (size == 2) {
		uint16_t t = inportw(PCI_VALUE_PORT + (field & 2));
		return t;
	} else if (size == 1) {
		uint8_t t = inportb(PCI_VALUE_PORT + (field & 3));
		return t;
	}

	return 0xFFFF;
}

uint16_t pci_find_type(uint32_t dev)
{
	return (pci_read_field(dev, PCI_CLASS, 1) << 8) | pci_read_field(dev, PCI_SUBCLASS, 1);
}

static void pci_scan_hit(pci_func_t f, uint32_t dev, void *arg)
{
	int dev_vend = (int)pci_read_field(dev, PCI_VENDOR_ID, 2);
	int dev_dvid = (int)pci_read_field(dev, PCI_DEVICE_ID, 2);

	f(dev, dev_vend, dev_dvid, arg);
}

static void pci_scan_func(pci_func_t f, int type, int bus, int slot, int func, void *arg)
{
	uint32_t dev = pci_box_device(bus, slot, func);
	if (type == -1 || type == pci_find_type(dev)) {
		pci_scan_hit(f, dev, arg);
	}

	if (pci_find_type(dev) == PCI_TYPE_BRIDGE) {
		pci_scan_bus(f, type, pci_read_field(dev, PCI_SECONDARY_BUS, 1), arg);
	}
}

static void pci_scan_slot(pci_func_t f, int type, int bus, int slot, void *arg)
{
	uint32_t dev = pci_box_device(bus, slot, 0);
	if (pci_read_field(dev, PCI_VENDOR_ID, 2) == PCI_NONE) {
		return;
	}

	pci_scan_func(f, type, bus, slot, 0, arg);
	if (!pci_read_field(dev, PCI_HEADER_TYPE, 1)) {
		return;
	}

	for (int func = 1; func < 8; func++) {
		uint32_t dev = pci_box_device(bus, slot, func);
		if (pci_read_field(dev, PCI_VENDOR_ID, 2) != PCI_NONE) {
			pci_scan_func(f, type, bus, slot, func, arg);
		}
	}
}

static void pci_scan_bus(pci_func_t f, int type, int bus, void *arg)
{
	for (int slot = 0; slot < 32; ++slot) {
		pci_scan_slot(f, type, bus, slot, arg);
	}
}

void pci_scan(pci_func_t f, int type, void *arg)
{
	if ((pci_read_field(0, PCI_HEADER_TYPE, 1) & 0x80) == 0) {
		pci_scan_bus(f, type, 0, arg);
		return;
	}

	int hit = 0;
	for (int func = 0; func < 8; ++func) {
		uint32_t dev = pci_box_device(0, 0, func);
		if (pci_read_field(dev, PCI_VENDOR_ID, 2) == PCI_NONE) {
			break;
		}

		hit = 1;
		pci_scan_bus(f, type, func, arg);
	}

	if (hit) {
		return;
	}

	for (int bus = 0; bus < 256; ++bus) {
		for (int slot = 0; slot < 32; ++slot) {
			pci_scan_slot(f, type, bus, slot, arg);
		}
	}
}

static void init_pci_device(uint32_t device, uint16_t vendid, uint16_t devid, void *arg)
{
	struct device *dev = generic_device_create();
	snprintf(dev->dev_name, sizeof dev->dev_name, "%02x:%02x.%u",
	         pci_get_bus(device),
	         pci_get_slot(device),
	         pci_get_func(device));

	device_register(dev, pci_driver, bus_device_base(pci_bus));
}

static kern_status_t pci_bus_scan(struct device *bus)
{
	printk("pci: scanning for devices...");
	pci_scan(init_pci_device, -1, NULL);
	return KERN_OK;
}

static struct bus_device_ops pci_bus_ops = {
	.scan = pci_bus_scan,
};

static kern_status_t online(struct kext *self)
{
	pci_driver = driver_create(self, "pci");
	if (!pci_driver) {
		return KERN_NO_MEMORY;
	}

	driver_register(pci_driver);

	pci_bus = bus_device_create();
	struct device *pci_base = bus_device_base(pci_bus);
	snprintf(pci_base->dev_name, sizeof pci_base->dev_name, "pci");
	pci_bus->b_ops = &pci_bus_ops;

	device_register(pci_base, pci_driver, root_device());

	return KERN_OK;
}

DEFINE_KEXT("net.doorstuck.socks.pci",
	    online, NULL,
	    KEXT_NO_DEPENDENCIES);