#define MODULE "pci"

#include "io/PCI.h"
#include "io/IO.h"
#include "log/Log.h"
#include "thread/Spinlock.h"

#define PCI_ADDRESS 0xCF8
#define PCI_VALUE 0xCFC

#define PCI_VENDOR_FIELD 0x0
#define PCI_DEVICE_FIELD 0x2
#define PCI_SUBCLASS_FIELD 0xa
#define PCI_CLASS_FIELD 0xb
#define PCI_REVISION_ID_FIELD 0x8
#define PCI_PROG_IF_FIELD 0x9
#define PCI_HEADER_TYPE_FIELD 0xe
#define PCI_SECONDARY_BUS_NUMBER_FIELD 0x19

Spinlock pci_lock;

uint32_t PCI::raw_address(uint32_t bus, uint32_t slot, uint32_t function, int32_t offset)
{
    return 0x80000000 | (bus << 16) | (slot << 11) | (function << 8) | ((offset)&0xFC);
}

void PCI::raw_write8(uint32_t bus, uint32_t slot, uint32_t function, int32_t offset, uint8_t value)
{
    IO::outl(PCI_ADDRESS, raw_address(bus, slot, function, offset));
    IO::outl(PCI_VALUE, (uint32_t)value);
}

void PCI::raw_write16(uint32_t bus, uint32_t slot, uint32_t function, int32_t offset, uint16_t value)
{
    IO::outl(PCI_ADDRESS, raw_address(bus, slot, function, offset));
    IO::outl(PCI_VALUE, (uint32_t)value);
}

void PCI::raw_write32(uint32_t bus, uint32_t slot, uint32_t function, int32_t offset, uint32_t value)
{
    IO::outl(PCI_ADDRESS, raw_address(bus, slot, function, offset));
    IO::outl(PCI_VALUE, value);
}

uint8_t PCI::raw_read8(uint32_t bus, uint32_t slot, uint32_t function, int32_t offset)
{
    IO::outl(PCI_ADDRESS, raw_address(bus, slot, function, offset));
    return IO::inl(PCI_VALUE + (offset & 3));
}

uint16_t PCI::raw_read16(uint32_t bus, uint32_t slot, uint32_t function, int32_t offset)
{
    IO::outl(PCI_ADDRESS, raw_address(bus, slot, function, offset));
    return IO::inl(PCI_VALUE + (offset & 2));
}

uint32_t PCI::raw_read32(uint32_t bus, uint32_t slot, uint32_t function, int32_t offset)
{
    IO::outl(PCI_ADDRESS, raw_address(bus, slot, function, offset));
    return IO::inl(PCI_VALUE);
}

PCI::DeviceID PCI::get_device_id(uint32_t bus, uint32_t slot, uint32_t function)
{
    uint16_t vendor = PCI::raw_read16(bus, slot, function, PCI_VENDOR_FIELD);
    uint16_t device = PCI::raw_read16(bus, slot, function, PCI_DEVICE_FIELD);
    return {vendor, device};
}

PCI::DeviceType PCI::get_device_type(uint32_t bus, uint32_t slot, uint32_t function)
{
    uint8_t dev_subclass = PCI::raw_read8(bus, slot, function, PCI_SUBCLASS_FIELD);
    uint8_t dev_class = PCI::raw_read8(bus, slot, function, PCI_CLASS_FIELD);
    uint8_t prog_if = PCI::raw_read8(bus, slot, function, PCI_PROG_IF_FIELD);
    uint8_t revision = PCI::raw_read8(bus, slot, function, PCI_REVISION_ID_FIELD);
    return {dev_class, dev_subclass, prog_if, revision};
}

static void pci_scan_bus(uint8_t bus, void (*callback)(PCI::DeviceID&, PCI::DeviceType&))
{
    for (uint8_t slot = 0; slot < 32; slot++)
    {
        uint8_t num_functions = 1;
        for (uint8_t function = 0; function < num_functions; function++)
        {
            PCI::DeviceID device_id = PCI::get_device_id(bus, slot, function);
            if (device_id.vendor == 0xFFFF || device_id.device == 0xFFFF) continue;
            PCI::DeviceType device_type = PCI::get_device_type(bus, slot, function);
            uint8_t header = PCI::raw_read8(bus, slot, function, PCI_HEADER_TYPE_FIELD);
            if (header & 0x80) // multi function device
            {
                num_functions = 8;
            }
            if ((header & 0x7F) == 1)
            {
                uint8_t sub_bus = PCI::raw_read8(bus, slot, function, PCI_SECONDARY_BUS_NUMBER_FIELD);
                pci_scan_bus(sub_bus, callback);
            }
            pci_lock.release();
            callback(device_id, device_type);
            pci_lock.acquire();
        }
    }
}

void PCI::scan(void (*callback)(PCI::DeviceID&, PCI::DeviceType&))
{
    pci_lock.acquire();
    pci_scan_bus(0, callback);
    pci_lock.release();
}