#include "arch/PCI.h"
#include "Log.h"

namespace PCI
{
    Device::ID read_id(const Device::Address& address)
    {
        u16 vendor = read16(address, Field::VendorID);
        u16 device = read16(address, Field::DeviceID);
        return { vendor, device };
    }

    Device::Type read_type(const Device::Address& address)
    {
        u8 klass = read8(address, Field::Class);
        u8 subclass = read8(address, Field::Subclass);
        u8 prog_if = read8(address, Field::ProgIF);
        return { klass, subclass, prog_if };
    }

    constexpr u8 PCI_BRIDGE_CLASS = 0x06;
    constexpr u8 PCI_TO_PCI_BRIDGE_SUBCLASS = 0x04;

    static bool is_pci_to_pci_bridge(const Device::Address& address)
    {
        return read8(address, Field::Class) == PCI_BRIDGE_CLASS &&
               read8(address, Field::Subclass) == PCI_TO_PCI_BRIDGE_SUBCLASS;
    }

    void scan_bus(u32);

    void scan_function(const Device::Address& address)
    {
        Device::ID id = read_id(address);
        Device::Type type = read_type(address);

        kinfoln("Found PCI device %.4x:%.4x (Class %.2x, Subclass %.2x, ProgIF %.2x) at PCI address %u:%u:%u",
                id.vendor, id.device, type.klass, type.subclass, type.prog_if, address.bus, address.slot,
                address.function);

        if (is_pci_to_pci_bridge(address))
        {
            kdbgln("PCI-to-PCI bridge detected, performing another bus scan");
            scan_bus(read8(address, Field::SecondaryBus));
        }
    }

    void scan_slot(u32 bus, u32 slot)
    {
        Device::Address address { bus, slot, 0 };

        // No device on this slot!
        if (read16(address, Field::VendorID) == PCI::INVALID_ID) return;

        scan_function(address);

        u8 header_type = read8(address, Field::HeaderType);

        // Multiple-function PCI device
        if ((header_type & 0x80) == 0x80)
        {
            for (u32 function = 1; function < 8; function++)
            {
                address.function = function;
                if (read16(address, Field::VendorID) == PCI::INVALID_ID) continue;
                scan_function(address);
            }
        }
    }

    void scan_bus(u32 bus)
    {
        for (u32 slot = 0; slot < 32; slot++) { scan_slot(bus, slot); }
    }

    void scan()
    {
        u8 header_type = read8({ 0, 0, 0 }, Field::HeaderType);

        // Single-function PCI bus
        if ((header_type & 0x80) == 0)
        {
            kdbgln("PCI bus is single-function");
            scan_bus(0);
        }
        else
        {
            kdbgln("PCI bus is multiple-function");
            for (u32 function = 0; function < 8; function++)
            {
                if (read16({ 0, 0, function }, Field::VendorID) != PCI::INVALID_ID)
                {
                    kdbgln("PCI bus has function %u", function);
                    scan_bus(function);
                }
            }
        }
    }
}