#include <nrvn.h>
#include <vga.h>
#include <pci.h>
#include <ahci.h>
#include <memory.h>
#include <mem.h>

#define PCI_CONFIG_ADDRESS 0xcf8
#define PCI_CONFIG_DATA 0xcfc

#define PAGE_SIZE sizeof(uint64_t) * 512
extern uint64_t pml4[512];

static uint32_t pci_read_config(uint8_t bus, uint8_t slot, uint8_t func, uint8_t offset) {
	outl(PCI_CONFIG_ADDRESS,
			(bus << 16) | (slot << 11) | (func << 8) | (offset & 0xfc) | (1 << 31));
	return inl(PCI_CONFIG_DATA);
}

static uint16_t pci_check_vendor(uint8_t bus, uint8_t slot) {
	uint16_t vendor, device;
	if ((vendor = pci_read_config(bus, slot, 0, 0)) != 0xffff) {
		device = pci_read_config(bus, slot, 0, 2);
		(void) device;
	}
	return vendor;
}

static void check_dev(uint8_t bus, uint8_t dev) {
	struct pci_device device = {0};
	for (size_t i = 0; i <= 15; i ++) {
		uint32_t rply = pci_read_config(bus, dev, 0, i * 4);
		if (rply == 0xffffffff)
			return;
		((uint32_t *) &device)[i] = rply;
	}

	vga_puts("device class ");
	vga_putx(device.header.class);
	vga_puts(" subclass ");
	vga_putx(device.header.subclass);
	vga_putc('\n');

	if (device.header.class == 0x1 && device.header.subclass == 0x6) {
		memory_map(device.general.bar5, sizeof(struct hba_mem));
		probe_port((struct hba_mem *) (uintptr_t) device.general.bar5);
	}
}

static void check_bus(uint8_t bus) {
	vga_puts("checking bus ");
	vga_putn(bus);
	vga_putc('\n');
	for (uint8_t dev = 0; dev < 32; dev++)
		check_dev(bus, dev);
}

static void check_function(uint8_t bus, uint8_t dev, uint8_t func) {
	uint8_t base_class, sub_class, second_bus;
	uint32_t rply = pci_read_config(bus, dev, func, 0x8);
	base_class = rply >> 24;
	sub_class = (rply >> 16) & 0xff;

	if (base_class == 0x6 && sub_class == 0x4) {
		rply = pci_read_config(bus, dev, func, 0x18);
		second_bus = (rply >> 8) & 0xff;
		check_bus(second_bus);
	}
}

struct pci_header read_header(uint8_t bus, uint8_t dev, uint8_t func) {
	struct pci_header header = {0};
	for (size_t i = 0; i < 0x4; i++) {
		((uint32_t*) &header)[i] = pci_read_config(bus, dev, func, i * 4);
	}
	return header;
}

bool pci_init(void) {
	struct pci_header header = read_header(0, 0, 0);
	if ((header.header_type & 0x80) == 0) {
		check_bus(0);
		return true;
	}
	
	for (size_t function = 0; function < 8; function++) {
		if ((pci_read_config(0, 0, function, 0x0) & 0xffff) != 0xffff)
			break;
		check_bus(function);
	}
	return true;
}