#include <string.h>
#include <assert.h>
#include <algorithm>

#include <inc/x86_64.hpp>
#include <inc/pci.hpp>
#include <inc/pcireg.h>
#include <inc/logger.hpp>
#include <inc/memory.hpp>
#include <inc/utils.hpp>

#define ARRAY_SIZE(a) (sizeof(a) / sizeof(a[0]))

using x86_64::inl;
using x86_64::outl;
using x86_64::inb;
using x86_64::outb;

namespace PCI {
	// ==== PCI device config storage ====
	
	#define MAX_N_PCI_DEVICES 128
	
	static Func stored_pci_devices[MAX_N_PCI_DEVICES];
	static int n_stored_pci_devices = 0;
	
	static int pci_store(Func *pcif) {
		if (n_stored_pci_devices < MAX_N_PCI_DEVICES) {
			pci_func_enable(pcif);
			memcpy(&stored_pci_devices[n_stored_pci_devices++], pcif, sizeof(Func));
			LINFO("PCI enabled: %04x.%04x",
				PCI_VENDOR(pcif->dev_id), PCI_PRODUCT(pcif->dev_id));
			return 0;
		} else {
			return -1;
		}
	}
	
	// ========
	
	// Flag to do "lspci" at bootup
	static int pci_show_devs = 1;
	static int pci_show_addrs = 0;
	
	// PCI "configuration mechanism one"
	static uint32_t pci_conf1_addr_ioport = 0x0cf8;
	static uint32_t pci_conf1_data_ioport = 0x0cfc;
	
	// Forward declarations
	static int pci_bridge_attach(Func *pcif);
	
	// PCI driver table
	struct Driver {
		uint32_t key1, key2;
		int (*attachfn) (Func *pcif);
	};
	
	// pci_attach_class matches the class and subclass of a PCI device
	Driver pci_attach_class[] = {
		{ PCI_CLASS_BRIDGE, PCI_SUBCLASS_BRIDGE_PCI, &pci_bridge_attach },
		{ 0, 0, 0 },
	};
	
	// pci_attach_vendor matches the vendor ID and device ID of a PCI device. key1
	// and key2 should be the vendor ID and device ID respectively
	Driver pci_attach_vendor[] = {
		// Only store e1000 network cards
		{ 0x8086, 0x15a3, &pci_store },  // e1000e
		{ 0x8086, 0x10d3, &pci_store },  // e1000e
		{ 0x8086, 0x153a, &pci_store },  // I217-LM (e1000e)
		{ 0x8086, 0x100e, &pci_store },  // e1000
		{ 0x8086, 0x15bc, &pci_store },  // I219-V
		{ 0x8086, 0x15b8, &pci_store },  // I219-V H310CM-ITX/ac
		{ 0x8086, 0x0d55, &pci_store },  // B460M TUF Gaming
		{ 0x8086, 0x15fa, &pci_store },  // H510M-HDV/M.2
		{ 0x8086, 0x1539, &pci_store },  // I211AT
		{ 0x8086, 0x15f3, &pci_store },  // I225-V
		{ 0x8086, 0x46d1, &pci_store },  // VGA (Intel N100)
		{ 0x10ec, 0x8168, &pci_store },  // RTL8111/8168
		{ 0x8086, 0x1a1c, &pci_store },  // I219-LM (13th gen)
		{ 0, 0, 0 },
	};
	
	static void pci_conf1_set_addr(uint32_t bus, uint32_t dev,
		uint32_t func, uint32_t offset) {
		assert(bus < 256);
		assert(dev < 32);
		assert(func < 8);
		assert(offset < 256);
		assert((offset & 0x3) == 0);
		
		uint32_t v = (1 << 31) |         // config-space
			(bus << 16) | (dev << 11) | (func << 8) | (offset);
		outl(pci_conf1_addr_ioport, v);
	}
	
	static uint32_t pci_conf_read(Func *f, uint32_t off) {
		pci_conf1_set_addr(f->bus->busno, f->dev, f->func, off);
		return inl(pci_conf1_data_ioport);
	}
	
	static void pci_conf_write(Func *f, uint32_t off, uint32_t v) {
		pci_conf1_set_addr(f->bus->busno, f->dev, f->func, off);
		outl(pci_conf1_data_ioport, v);
	}
	
	static uint8_t pci_conf_read_byte(Func *f, uint32_t off) {
		pci_conf1_set_addr(f->bus->busno, f->dev, f->func, off & ~3);
		return inb(pci_conf1_data_ioport + (off & 3));
	}
	
	static void pci_conf_write_byte(Func *f, uint32_t off, uint8_t v) {
		pci_conf1_set_addr(f->bus->busno, f->dev, f->func, off & ~3);
		outb(pci_conf1_data_ioport + (off & 3), v);
	}
	
	static int pci_attach_match(uint32_t key1, uint32_t key2,
		Driver *list, Func *pcif) {
		for (uint32_t i = 0; list[i].attachfn; i++) {
			if (list[i].key1 == key1 && list[i].key2 == key2) {
				int r = list[i].attachfn(pcif);
				if (r > 0) return r;
				if (r < 0) {
					LERROR("pci_attach_match: attaching %x.%x (%p): e",
						key1, key2, list[i].attachfn, r);
				}
			}
		}
		
		return 0;
	}
	
	static int pci_attach(Func *f) {
		return
			pci_attach_match(
				PCI_CLASS(f->dev_class),
				PCI_SUBCLASS(f->dev_class),
				&pci_attach_class[0],
				f
			) || pci_attach_match(
				PCI_VENDOR(f->dev_id),
				PCI_PRODUCT(f->dev_id),
				&pci_attach_vendor[0],
				f
			);
	}
	
	static const char *pci_class[] = {
		[0x0] = "Unknown",
		[0x1] = "Storage controller",
		[0x2] = "Network controller",
		[0x3] = "Display controller",
		[0x4] = "Multimedia device",
		[0x5] = "Memory controller",
		[0x6] = "Bridge device",
	};
	
	static void pci_print_func(Func *f) {
		const char *the_class = pci_class[0];
		if (PCI_CLASS(f->dev_class) < ARRAY_SIZE(pci_class)) {
			the_class = pci_class[PCI_CLASS(f->dev_class)];
		}
		
		LDEBUG("PCI: %02x:%02x.%d: %04x:%04x: class: %x.%x (%s) irq: %d",
			f->bus->busno, f->dev, f->func,
			PCI_VENDOR(f->dev_id), PCI_PRODUCT(f->dev_id),
			PCI_CLASS(f->dev_class), PCI_SUBCLASS(f->dev_class), the_class,
			f->irq_line);
	}
	
	static int pci_scan_bus(Bus *bus) {
		int totaldev = 0;
		Func df;
		memset(&df, 0, sizeof(df));
		df.bus = bus;
		
		for (df.dev = 0; df.dev < 32; df.dev++) {
			uint32_t bhlc = pci_conf_read(&df, PCI_BHLC_REG);
			if (PCI_HDRTYPE_TYPE(bhlc) > 1) {	    // Unsupported or no device
				continue;
			}
			
			totaldev++;
			
			Func f = df;
			for (f.func = 0; f.func < (PCI_HDRTYPE_MULTIFN(bhlc) ? 8 : 1); f.func++) {
				Func af = f;
				
				af.dev_id = pci_conf_read(&f, PCI_ID_REG);
				if (PCI_VENDOR(af.dev_id) == 0xffff) {
					continue;
				}
				
				uint32_t intr = pci_conf_read(&af, PCI_INTERRUPT_REG);
				af.irq_line = PCI_INTERRUPT_LINE(intr);
				
				af.dev_class = pci_conf_read(&af, PCI_CLASS_REG);
				if (pci_show_devs) {
					pci_print_func(&af);
				}
				pci_attach(&af);
			}
		}
		
		return totaldev;
	}
	
	static int pci_bridge_attach(Func *pcif) {
		uint32_t ioreg  = pci_conf_read(pcif, PCI_BRIDGE_STATIO_REG);
		uint32_t busreg = pci_conf_read(pcif, PCI_BRIDGE_BUS_REG);
		
		if (PCI_BRIDGE_IO_32BITS(ioreg)) {
			LWARN("PCI: %02x:%02x.%d: 32-bit bridge IO not supported.",
				pcif->bus->busno, pcif->dev, pcif->func);
			return 0;
		}
		
		Bus nbus;
		memset(&nbus, 0, sizeof(nbus));
		nbus.parent_bridge = pcif;
		nbus.busno = (busreg >> PCI_BRIDGE_BUS_SECONDARY_SHIFT) & 0xff;
		
		if (pci_show_devs) {
			LDEBUG("PCI: %02x:%02x.%d: bridge to PCI bus %d--%d\n",
				pcif->bus->busno, pcif->dev, pcif->func,
				nbus.busno,
				(busreg >> PCI_BRIDGE_BUS_SUBORDINATE_SHIFT) & 0xff);
		}
		
		pci_scan_bus(&nbus);
		return 1;
	}
	
	// External PCI subsystem interface
	
	void pci_func_enable(Func *f) {
		pci_conf_write(f, PCI_COMMAND_STATUS_REG,
			PCI_COMMAND_IO_ENABLE
			| PCI_COMMAND_MEM_ENABLE
			| PCI_COMMAND_MASTER_ENABLE);
		
		#define PCI_LATENCY_TIMER 0x0d
		uint8_t lat = pci_conf_read_byte(f, PCI_LATENCY_TIMER);
		if (lat < 32) {
			LWARN("PCI: %02x:%02x.%d: latency timer too low (%d), "
				"increasing to 32",
				f->bus->busno, f->dev, f->func, lat);
			pci_conf_write_byte(f, PCI_LATENCY_TIMER, 32);
		}
		
		uint32_t bar_width;
		uint32_t bar;
		for (bar = PCI_MAPREG_START; bar < PCI_MAPREG_END; bar += bar_width) {
			uint32_t oldv = pci_conf_read(f, bar);
			
			bar_width = 4;
			pci_conf_write(f, bar, 0xffffffff);
			uint32_t rv = pci_conf_read(f, bar);
			
			if (rv == 0) {
				continue;
			}
			
			int regnum = PCI_MAPREG_NUM(bar);
			uint32_t base, size;
			if (PCI_MAPREG_TYPE(rv) == PCI_MAPREG_TYPE_MEM) {
				if (PCI_MAPREG_MEM_TYPE(rv) == PCI_MAPREG_MEM_TYPE_64BIT) {
					bar_width = 8;
				}
				
				f->reg_is_io[regnum] = false;
				
				size = PCI_MAPREG_MEM_SIZE(rv);
				base = PCI_MAPREG_MEM_ADDR(oldv);
				if (pci_show_addrs) {
					LDEBUG("  mem region %d: %d bytes at 0x%x",
						regnum, size, base);
				}
			} else {
				f->reg_is_io[regnum] = true;
				
				size = PCI_MAPREG_IO_SIZE(rv);
				base = PCI_MAPREG_IO_ADDR(oldv);
				if (pci_show_addrs) {
					LDEBUG("  io region %d: %d bytes at 0x%x",
						regnum, size, base);
				}
			}
			
			pci_conf_write(f, bar, oldv);
			f->reg_base[regnum] = base;
			f->reg_size[regnum] = size;
			
			if (size && !base) {
				LWARN("PCI device %02x:%02x.%d (%04x:%04x) "
					"may be misconfigured: "
					"region %d: base 0x%x, size %d",
					f->bus->busno, f->dev, f->func,
					PCI_VENDOR(f->dev_id), PCI_PRODUCT(f->dev_id),
					regnum, base, size);
			}
		}
		
		LDEBUG("PCI function %02x:%02x.%d (%04x:%04x) enabled",
			f->bus->busno, f->dev, f->func,
			PCI_VENDOR(f->dev_id), PCI_PRODUCT(f->dev_id));
	}
	
	int init() {
		LDEBUG_ENTER_RET();
		
		static Bus root_bus;
		memset(&root_bus, 0, sizeof(root_bus));
		
		int ret = pci_scan_bus(&root_bus);
		return ret;
	}
	
	// ==== For user mode PCI drivers ====
	
	uint64_t map_device(uint32_t key1, uint32_t key2, uint64_t base, uint64_t maxlen, uint32_t bar) {
		bool found = false;
		uint64_t reg_base, reg_size;
		for (int i = 0; i < n_stored_pci_devices; i++) {
			uint32_t dev_id = stored_pci_devices[i].dev_id;
			if (PCI_VENDOR(dev_id) == key1 && PCI_PRODUCT(dev_id) == key2) {
				if (stored_pci_devices[i].reg_is_io[bar]) {
					continue;
				}
				
				reg_base = stored_pci_devices[i].reg_base[bar];
				reg_size = stored_pci_devices[i].reg_size[bar];
				found = true;
				break;
			}
		}
		
		if (found) {
			LDEBUG("pci: map_device %04x:%04x at 0x%llx len 0x%llx (bar %d) -> 0x%llx",
				key1, key2, base, maxlen, bar, reg_base);
			maxlen = std::min(reg_size, maxlen);
			uint64_t end = base + Utils::round_up(maxlen, Memory::PAGE_SIZE);
			Memory::map_region_cache_disabled(base, end, reg_base);
			
			// reload page table
			x86_64::lcr3(x86_64::rcr3());
			
			return maxlen;
		} else {
			return -1ull;
		}
	}
	
	uint64_t get_device_reg_base(uint32_t key1, uint32_t key2, uint32_t bar) {
		for (int i = 0; i < n_stored_pci_devices; i++) {
			uint32_t dev_id = stored_pci_devices[i].dev_id;
			if (PCI_VENDOR(dev_id) == key1 && PCI_PRODUCT(dev_id) == key2) {
				if (!stored_pci_devices[i].reg_is_io[bar]) {
					continue;
				}
				
				uint64_t reg_base = stored_pci_devices[i].reg_base[bar];
				return reg_base;
			}
		}
		
		return -1ull;
	}
	
}
