/*
    * Pci.cpp
    * PCI Express enumeration and configuration space access
    * Copyright (c) 2025 Daniel Hammer
*/

#include "Pci.hpp"
#include <Arch/Interrupts.hpp>
#include <Terminal/Terminal.hpp>
#include <CppLib/Stream.hpp>
#include <Memory/HHDM.hpp>
#include <Memory/Paging.hpp>
#include <Io/IoPort.hpp>
#include <CppLib/Spinlock.hpp>
#include <Timekeeping/Timer.hpp>

using namespace Kt;

namespace Pci {

    // Legacy PCI config I/O ports
    static constexpr uint16_t ConfigAddressPort = 0xCF8;
    static constexpr uint16_t ConfigDataPort    = 0xCFC;

    // CF8/CFC is one machine-wide address/data latch, not a per-CPU interface.
    // Keep the address write and matching data access indivisible: otherwise a
    // second CPU can replace CF8 between them and make us read or write an
    // unrelated device's register.  Besides bogus probing, a crossed runtime
    // command/MSI write can disable interrupts or bus mastering underneath an
    // active GPU, USB controller, NIC, or HDA controller.
    static kcp::Spinlock g_legacyConfigLock;

    // PCI config space register offsets
    static constexpr uint16_t RegVendorId   = 0x00;
    static constexpr uint16_t RegDeviceId   = 0x02;
    static constexpr uint16_t RegCommand    = 0x04;
    static constexpr uint16_t RegStatus     = 0x06;
    static constexpr uint16_t RegRevisionId = 0x08;
    static constexpr uint16_t RegProgIf     = 0x09;
    static constexpr uint16_t RegSubClass   = 0x0A;
    static constexpr uint16_t RegClassCode  = 0x0B;
    static constexpr uint16_t RegHeaderType = 0x0E;

    static kcp::vector<PciDevice> g_devices{};

    // State: ECAM base (0 if not available, use legacy)
    static uint64_t g_ecamBase = 0;
    static uint8_t g_ecamStartBus = 0;
    static uint8_t g_ecamEndBus = 0;
    static bool g_useEcam = false;

    // -------------------------------------------------------------------------
    // MCFG table discovery (same pattern as MADT)
    // -------------------------------------------------------------------------

    static bool SignatureMatch(const char* sig, const char* target, int len) {
        for (int i = 0; i < len; i++) {
            if (sig[i] != target[i]) return false;
        }
        return true;
    }

    static Hal::ACPI::CommonSDTHeader* FindMcfgInXsdt(Hal::ACPI::CommonSDTHeader* xsdt) {
        uint32_t entryCount = (xsdt->Length - sizeof(Hal::ACPI::CommonSDTHeader)) / 8;
        uint64_t* entries = (uint64_t*)((uint64_t)xsdt + sizeof(Hal::ACPI::CommonSDTHeader));

        for (uint32_t i = 0; i < entryCount; i++) {
            auto* header = (Hal::ACPI::CommonSDTHeader*)Memory::HHDM(entries[i]);
            if (SignatureMatch(header->Signature, "MCFG", 4)) {
                return header;
            }
        }

        return nullptr;
    }

    static bool ParseMcfg(Hal::ACPI::CommonSDTHeader* xsdt) {
        auto* mcfgHeader = FindMcfgInXsdt(xsdt);
        if (mcfgHeader == nullptr) {
            KernelLogStream(WARNING, "PCI") << "MCFG table not found, falling back to legacy config access";
            return false;
        }

        if (!Hal::ACPI::TestChecksum(mcfgHeader)) {
            KernelLogStream(ERROR, "PCI") << "MCFG checksum failed";
            return false;
        }

        KernelLogStream(OK, "PCI") << "Found MCFG table";

        auto* mcfg = (McfgHeader*)mcfgHeader;
        uint32_t entriesSize = mcfg->SDTHeader.Length - sizeof(McfgHeader);
        uint32_t entryCount = entriesSize / sizeof(McfgEntry);

        if (entryCount == 0) {
            KernelLogStream(WARNING, "PCI") << "MCFG contains no entries";
            return false;
        }

        auto* entries = (McfgEntry*)((uint64_t)mcfg + sizeof(McfgHeader));

        // Use the first MCFG entry (segment group 0)
        g_ecamBase = entries[0].BaseAddress;
        g_ecamStartBus = entries[0].StartBus;
        g_ecamEndBus = entries[0].EndBus;

        KernelLogStream(INFO, "PCI") << "ECAM base: " << base::hex << g_ecamBase
            << " buses " << base::dec << (uint64_t)g_ecamStartBus
            << "-" << (uint64_t)g_ecamEndBus;

        // Map the ECAM MMIO region
        // Size = (endBus - startBus + 1) * 32 devices * 8 functions * 4096 bytes
        uint64_t busCount = (uint64_t)(g_ecamEndBus - g_ecamStartBus + 1);
        uint64_t ecamSize = busCount * 32 * 8 * 4096;

        if (Memory::VMM::g_paging) {
            for (uint64_t offset = 0; offset < ecamSize; offset += 0x1000) {
                uint64_t phys = g_ecamBase + offset;
                Memory::VMM::g_paging->MapMMIO(phys, Memory::HHDM(phys));
            }
            KernelLogStream(DEBUG, "PCI") << "Mapped ECAM region: " << base::hex << ecamSize << " bytes";
        }

        g_useEcam = true;
        return true;
    }

    // -------------------------------------------------------------------------
    // ECAM (memory-mapped) configuration space access
    // -------------------------------------------------------------------------

    static volatile uint8_t* EcamAddress(uint64_t ecamBase, uint8_t bus, uint8_t device, uint8_t function, uint16_t offset) {
        uint64_t phys = ecamBase
            + ((uint64_t)bus << 20)
            + ((uint64_t)device << 15)
            + ((uint64_t)function << 12)
            + offset;
        return (volatile uint8_t*)Memory::HHDM(phys);
    }

    uint8_t EcamRead8(uint64_t ecamBase, uint8_t bus, uint8_t device, uint8_t function, uint16_t offset) {
        return *EcamAddress(ecamBase, bus, device, function, offset);
    }

    uint16_t EcamRead16(uint64_t ecamBase, uint8_t bus, uint8_t device, uint8_t function, uint16_t offset) {
        return *(volatile uint16_t*)EcamAddress(ecamBase, bus, device, function, offset);
    }

    uint32_t EcamRead32(uint64_t ecamBase, uint8_t bus, uint8_t device, uint8_t function, uint16_t offset) {
        return *(volatile uint32_t*)EcamAddress(ecamBase, bus, device, function, offset);
    }

    void EcamWrite8(uint64_t ecamBase, uint8_t bus, uint8_t device, uint8_t function, uint16_t offset, uint8_t value) {
        *EcamAddress(ecamBase, bus, device, function, offset) = value;
    }

    void EcamWrite16(uint64_t ecamBase, uint8_t bus, uint8_t device, uint8_t function, uint16_t offset, uint16_t value) {
        *(volatile uint16_t*)EcamAddress(ecamBase, bus, device, function, offset) = value;
    }

    void EcamWrite32(uint64_t ecamBase, uint8_t bus, uint8_t device, uint8_t function, uint16_t offset, uint32_t value) {
        *(volatile uint32_t*)EcamAddress(ecamBase, bus, device, function, offset) = value;
    }

    // -------------------------------------------------------------------------
    // Legacy PCI configuration space access (I/O ports 0xCF8/0xCFC)
    // -------------------------------------------------------------------------

    static uint32_t LegacyBuildAddress(uint8_t bus, uint8_t device, uint8_t function, uint8_t offset) {
        return (1u << 31)                       // Enable bit
             | ((uint32_t)bus << 16)
             | ((uint32_t)(device & 0x1F) << 11)
             | ((uint32_t)(function & 0x07) << 8)
             | (offset & 0xFC);
    }

    static uint32_t LegacyRead32Unlocked(uint8_t bus, uint8_t device,
                                         uint8_t function, uint8_t offset) {
        Io::Out32(LegacyBuildAddress(bus, device, function, offset), ConfigAddressPort);
        return Io::In32(ConfigDataPort);
    }

    static void LegacyWrite32Unlocked(uint8_t bus, uint8_t device,
                                      uint8_t function, uint8_t offset,
                                      uint32_t value) {
        Io::Out32(LegacyBuildAddress(bus, device, function, offset), ConfigAddressPort);
        Io::Out32(value, ConfigDataPort);
    }

    // Drivers reach config space through the Legacy* calls below. Without
    // an I/O address space (ARM) there is no CF8/CFC mechanism, so they use
    // ECAM there, at the access's own width (16/8-bit writes must not
    // read-modify-write a neighbouring RW1C register).
    static bool EcamCovers(uint8_t bus);

    static bool UseEcamForLegacy(uint8_t bus) {
        return !Io::HasPortIo && EcamCovers(bus);
    }

    uint32_t LegacyRead32(uint8_t bus, uint8_t device, uint8_t function, uint8_t offset) {
        if (UseEcamForLegacy(bus)) {
            return EcamRead32(g_ecamBase, bus, device, function, offset & 0xFC);
        }
        g_legacyConfigLock.Acquire();
        uint32_t value = LegacyRead32Unlocked(bus, device, function, offset);
        g_legacyConfigLock.Release();
        return value;
    }

    uint16_t LegacyRead16(uint8_t bus, uint8_t device, uint8_t function, uint8_t offset) {
        uint32_t val = LegacyRead32(bus, device, function, offset & 0xFC);
        return (uint16_t)(val >> ((offset & 2) * 8));
    }

    uint8_t LegacyRead8(uint8_t bus, uint8_t device, uint8_t function, uint8_t offset) {
        uint32_t val = LegacyRead32(bus, device, function, offset & 0xFC);
        return (uint8_t)(val >> ((offset & 3) * 8));
    }

    void LegacyWrite32(uint8_t bus, uint8_t device, uint8_t function, uint8_t offset, uint32_t value) {
        if (UseEcamForLegacy(bus)) {
            EcamWrite32(g_ecamBase, bus, device, function, offset & 0xFC, value);
            return;
        }
        g_legacyConfigLock.Acquire();
        LegacyWrite32Unlocked(bus, device, function, offset, value);
        g_legacyConfigLock.Release();
    }

    void LegacyWrite16(uint8_t bus, uint8_t device, uint8_t function, uint8_t offset, uint16_t value) {
        if (UseEcamForLegacy(bus)) {
            EcamWrite16(g_ecamBase, bus, device, function, offset & 0xFE, value);
            return;
        }
        g_legacyConfigLock.Acquire();
        uint32_t tmp = LegacyRead32Unlocked(bus, device, function, offset & 0xFC);
        int shift = (offset & 2) * 8;
        tmp &= ~(0xFFFF << shift);
        tmp |= ((uint32_t)value << shift);
        LegacyWrite32Unlocked(bus, device, function, offset & 0xFC, tmp);
        g_legacyConfigLock.Release();
    }

    void LegacyWrite8(uint8_t bus, uint8_t device, uint8_t function, uint8_t offset, uint8_t value) {
        if (UseEcamForLegacy(bus)) {
            EcamWrite8(g_ecamBase, bus, device, function, offset, value);
            return;
        }
        g_legacyConfigLock.Acquire();
        uint32_t tmp = LegacyRead32Unlocked(bus, device, function, offset & 0xFC);
        int shift = (offset & 3) * 8;
        tmp &= ~(0xFF << shift);
        tmp |= ((uint32_t)value << shift);
        LegacyWrite32Unlocked(bus, device, function, offset & 0xFC, tmp);
        g_legacyConfigLock.Release();
    }

    // -------------------------------------------------------------------------
    // Unified read helpers (use ECAM if available, else legacy)
    // -------------------------------------------------------------------------

    static uint16_t ReadConfig16(uint8_t bus, uint8_t device, uint8_t function, uint16_t offset) {
        if (g_useEcam) {
            return EcamRead16(g_ecamBase, bus, device, function, offset);
        }
        return LegacyRead16(bus, device, function, (uint8_t)offset);
    }

    static uint8_t ReadConfig8(uint8_t bus, uint8_t device, uint8_t function, uint16_t offset) {
        if (g_useEcam) {
            return EcamRead8(g_ecamBase, bus, device, function, offset);
        }
        return LegacyRead8(bus, device, function, (uint8_t)offset);
    }

    static bool EcamCovers(uint8_t bus) {
        return g_useEcam && bus >= g_ecamStartBus && bus <= g_ecamEndBus;
    }

    uint32_t ConfigRead(uint8_t bus, uint8_t device, uint8_t function, uint16_t offset, uint8_t width) {
        if (device > 31 || function > 7 || offset > 0xFFF || (offset & (width - 1)))
            return 0xFFFFFFFF;

        if (EcamCovers(bus)) {
            switch (width) {
                case 1: return EcamRead8(g_ecamBase, bus, device, function, offset);
                case 2: return EcamRead16(g_ecamBase, bus, device, function, offset);
                case 4: return EcamRead32(g_ecamBase, bus, device, function, offset);
                default: return 0xFFFFFFFF;
            }
        }

        if (offset > 0xFF) return 0xFFFFFFFF;
        switch (width) {
            case 1: return LegacyRead8(bus, device, function, (uint8_t)offset);
            case 2: return LegacyRead16(bus, device, function, (uint8_t)offset);
            case 4: return LegacyRead32(bus, device, function, (uint8_t)offset);
            default: return 0xFFFFFFFF;
        }
    }

    bool ConfigWrite(uint8_t bus, uint8_t device, uint8_t function, uint16_t offset, uint8_t width, uint32_t value) {
        if (device > 31 || function > 7 || offset > 0xFFF || (offset & (width - 1)))
            return false;

        if (EcamCovers(bus)) {
            switch (width) {
                case 1: EcamWrite8(g_ecamBase, bus, device, function, offset, (uint8_t)value); return true;
                case 2: EcamWrite16(g_ecamBase, bus, device, function, offset, (uint16_t)value); return true;
                case 4: EcamWrite32(g_ecamBase, bus, device, function, offset, value); return true;
                default: return false;
            }
        }

        if (offset > 0xFF) return false;
        switch (width) {
            case 1: LegacyWrite8(bus, device, function, (uint8_t)offset, (uint8_t)value); return true;
            case 2: LegacyWrite16(bus, device, function, (uint8_t)offset, (uint16_t)value); return true;
            case 4: LegacyWrite32(bus, device, function, (uint8_t)offset, value); return true;
            default: return false;
        }
    }

    // -------------------------------------------------------------------------
    // PCI capability list traversal
    // -------------------------------------------------------------------------

    uint8_t FindCapability(uint8_t bus, uint8_t device, uint8_t function, uint8_t capId) {
        // Check Status register bit 4 (Capabilities List present)
        uint16_t status = ReadConfig16(bus, device, function, RegStatus);
        if (!(status & (1 << 4))) {
            return 0;
        }

        // Read Capabilities Pointer (offset 0x34), mask to dword-aligned
        uint8_t offset = ReadConfig8(bus, device, function, 0x34) & 0xFC;

        // Walk the conventional 256-byte capability list (cap_id @ +0,
        // next_ptr @ +1). Firmware owns these pointers, so reject offsets
        // below the capability area and cycles instead of allowing a corrupt
        // list to spin forever during device initialization or resume.
        uint64_t visited = 0;
        while (offset != 0) {
            if (offset < 0x40) return 0;
            uint64_t bit = 1ULL << (offset >> 2);
            if (visited & bit) return 0;
            visited |= bit;

            uint8_t id = ReadConfig8(bus, device, function, offset);
            if (id == capId) {
                return offset;
            }
            offset = ReadConfig8(bus, device, function, offset + 1) & 0xFC;
        }

        return 0;
    }

    bool WriteMsiMessage(uint8_t bus, uint8_t device, uint8_t function, uint8_t cap, uint8_t irq) {
        Arch::MsiMessage msi;
        if (!Arch::ComposeMsi(irq, msi)) {
            return false;
        }
        uint16_t msgCtrl = LegacyRead16(bus, device, function, cap + 2);

        LegacyWrite32(bus, device, function, cap + 4, (uint32_t)msi.address);
        if (msgCtrl & (1 << 7)) {
            // 64-bit: Upper Address at cap+8, Data at cap+12
            LegacyWrite32(bus, device, function, cap + 8, (uint32_t)(msi.address >> 32));
            LegacyWrite16(bus, device, function, cap + 12, (uint16_t)msi.data);
        } else {
            // 32-bit: Data at cap+8
            LegacyWrite16(bus, device, function, cap + 8, (uint16_t)msi.data);
        }
        return true;
    }

    bool DisableInterruptDelivery(uint8_t bus, uint8_t device, uint8_t function) {
        uint16_t command = LegacyRead16(bus, device, function,
            (uint8_t)PCI_REG_COMMAND);
        if (!(command & PCI_CMD_INTX_DISABLE)) {
            LegacyWrite16(bus, device, function, (uint8_t)PCI_REG_COMMAND,
                command | PCI_CMD_INTX_DISABLE);
        }

        bool disabledMessageDelivery = false;
        uint8_t msi = FindCapability(bus, device, function, PCI_CAP_MSI);
        if (msi != 0) {
            uint16_t control = LegacyRead16(bus, device, function, msi + 2);
            if (control & 1) {
                LegacyWrite16(bus, device, function, msi + 2, control & ~1u);
                disabledMessageDelivery = true;
            }
        }

        uint8_t msix = FindCapability(bus, device, function, PCI_CAP_MSIX);
        if (msix != 0) {
            uint16_t control = LegacyRead16(bus, device, function, msix + 2);
            if (control & (1u << 15)) {
                // Function-mask before clearing MSI-X Enable so no vector can
                // escape during the ownership transition.
                uint16_t disabled = (control | (1u << 14)) & ~(1u << 15);
                LegacyWrite16(bus, device, function, msix + 2, disabled);
                disabledMessageDelivery = true;
            }
        }
        return disabledMessageDelivery;
    }

    // -------------------------------------------------------------------------
    // PCI class code names
    // -------------------------------------------------------------------------

    const char* GetClassName(uint8_t classCode, uint8_t subClass) {
        switch (classCode) {
            case 0x00:
                return "Unclassified";
            case 0x01:
                switch (subClass) {
                    case 0x00: return "SCSI Bus Controller";
                    case 0x01: return "IDE Controller";
                    case 0x02: return "Floppy Disk Controller";
                    case 0x05: return "ATA Controller";
                    case 0x06: return "SATA Controller";
                    case 0x08: return "NVM Controller";
                    default:   return "Mass Storage Controller";
                }
            case 0x02:
                switch (subClass) {
                    case 0x00: return "Ethernet Controller";
                    case 0x80: return "Other Network Controller";
                    default:   return "Network Controller";
                }
            case 0x03:
                switch (subClass) {
                    case 0x00: return "VGA Compatible Controller";
                    case 0x01: return "XGA Controller";
                    case 0x02: return "3D Controller";
                    default:   return "Display Controller";
                }
            case 0x04:
                return "Multimedia Controller";
            case 0x05:
                return "Memory Controller";
            case 0x06:
                switch (subClass) {
                    case 0x00: return "Host Bridge";
                    case 0x01: return "ISA Bridge";
                    case 0x04: return "PCI-to-PCI Bridge";
                    case 0x80: return "Other Bridge";
                    default:   return "Bridge Device";
                }
            case 0x07:
                return "Simple Communication Controller";
            case 0x08:
                return "Base System Peripheral";
            case 0x09:
                return "Input Device Controller";
            case 0x0A:
                return "Docking Station";
            case 0x0B:
                return "Processor";
            case 0x0C:
                switch (subClass) {
                    case 0x03: return "USB Controller";
                    case 0x05: return "SMBus Controller";
                    default:   return "Serial Bus Controller";
                }
            case 0x0D:
                return "Wireless Controller";
            case 0x0E:
                return "Intelligent Controller";
            case 0x0F:
                return "Satellite Communication Controller";
            case 0x10:
                return "Encryption Controller";
            case 0x11:
                return "Signal Processing Controller";
            case 0xFF:
                return "Unassigned";
            default:
                return "Unknown";
        }
    }

    // -------------------------------------------------------------------------
    // Device enumeration
    // -------------------------------------------------------------------------

    static void EnumerateFunction(uint8_t bus, uint8_t device, uint8_t function) {
        uint16_t vendorId = ReadConfig16(bus, device, function, RegVendorId);
        if (vendorId == 0xFFFF) {
            return;
        }

        PciDevice dev{};
        dev.Segment = 0;
        dev.Bus = bus;
        dev.Device = device;
        dev.Function = function;
        dev.VendorId = vendorId;
        dev.DeviceId = ReadConfig16(bus, device, function, RegDeviceId);
        dev.ClassCode = ReadConfig8(bus, device, function, RegClassCode);
        dev.SubClass = ReadConfig8(bus, device, function, RegSubClass);
        dev.ProgIf = ReadConfig8(bus, device, function, RegProgIf);
        dev.RevisionId = ReadConfig8(bus, device, function, RegRevisionId);
        dev.HeaderType = ReadConfig8(bus, device, function, RegHeaderType);

        g_devices.push_back(dev);

        KernelLogStream(DEBUG, "PCI") << base::hex
            << (uint64_t)bus << ":" << (uint64_t)device << "." << (uint64_t)function
            << " " << (uint64_t)vendorId << ":" << (uint64_t)dev.DeviceId
            << " " << GetClassName(dev.ClassCode, dev.SubClass)
            << " (class " << (uint64_t)dev.ClassCode << "." << (uint64_t)dev.SubClass << ")";
    }

    static void EnumerateDevice(uint8_t bus, uint8_t device) {
        uint16_t vendorId = ReadConfig16(bus, device, 0, RegVendorId);
        if (vendorId == 0xFFFF) {
            return;
        }

        // Always check function 0
        EnumerateFunction(bus, device, 0);

        // Check if multi-function device (bit 7 of header type)
        uint8_t headerType = ReadConfig8(bus, device, 0, RegHeaderType);
        if (headerType & 0x80) {
            for (uint8_t func = 1; func < 8; func++) {
                EnumerateFunction(bus, device, func);
            }
        }
    }

    static void EnumerateBus(uint8_t bus) {
        for (uint8_t device = 0; device < 32; device++) {
            EnumerateDevice(bus, device);
        }
    }

    static void EnumerateAll() {
        uint8_t startBus = g_useEcam ? g_ecamStartBus : 0;
        uint8_t endBus = g_useEcam ? g_ecamEndBus : 255;

        for (uint32_t bus = startBus; bus <= (uint32_t)endBus; bus++) {
            EnumerateBus((uint8_t)bus);
        }
    }

    // -------------------------------------------------------------------------
    // Public API
    // -------------------------------------------------------------------------

    kcp::vector<PciDevice>& GetDevices() {
        return g_devices;
    }

    // -------------------------------------------------------------------------
    // Shared PCI utilities
    // -------------------------------------------------------------------------

    uint64_t ReadBar0(uint8_t bus, uint8_t device, uint8_t function) {
        uint32_t bar0Low = LegacyRead32(bus, device, function, (uint8_t)PCI_REG_BAR0);

        // Bit 0 == 1 means this is an I/O-space BAR, not memory. The low bits of
        // an I/O BAR are a port number, not a physical address, and must never be
        // fed to MapMMIO (it would panic on the non-page-aligned value). Callers
        // treat 0 as "no usable MMIO BAR". (e.g. VirtualBox AC'97 8086:2415.)
        if (bar0Low & 0x1)
            return 0;

        uint64_t addr = bar0Low & 0xFFFFFFF0u;

        // Check for 64-bit BAR (type field bits 2:1 == 0b10)
        if ((bar0Low & 0x06) == 0x04) {
            uint32_t bar0High = LegacyRead32(bus, device, function, (uint8_t)(PCI_REG_BAR0 + 4));
            addr |= ((uint64_t)bar0High << 32);
        }

        return addr;
    }

    void EnableBusMaster(uint8_t bus, uint8_t device, uint8_t function) {
        uint16_t cmd = LegacyRead16(bus, device, function, (uint8_t)PCI_REG_COMMAND);
        cmd |= PCI_CMD_MEM_SPACE | PCI_CMD_BUS_MASTER;
        LegacyWrite16(bus, device, function, (uint8_t)PCI_REG_COMMAND, cmd);
    }

    bool IsBridge(const PciDevice& dev) {
        return (dev.HeaderType & 0x7F) == 0x01;
    }

    // PMCSR lives at PM capability + 4; its low two bits are the D-state.
    static constexpr uint16_t PmcsrOffset = 4;

    int GetPowerState(uint8_t bus, uint8_t device, uint8_t function) {
        if (ReadConfig16(bus, device, function, RegVendorId) == 0xFFFF) return -1;
        uint8_t pm = FindCapability(bus, device, function, PCI_CAP_PM);
        if (pm == 0) return -1;
        return (int)(ReadConfig16(bus, device, function, pm + PmcsrOffset) & 0x3);
    }

    bool SetPowerState(uint8_t bus, uint8_t device, uint8_t function, int state) {
        if (state < 0 || state > 3) return false;
        uint8_t pm = FindCapability(bus, device, function, PCI_CAP_PM);
        if (pm == 0) return false;

        // D1/D2 are optional; PMC bits 25/26 advertise them.
        uint32_t pmc = ConfigRead(bus, device, function, pm, 4);
        if (state == 1 && !(pmc & (1u << 25))) return false;
        if (state == 2 && !(pmc & (1u << 26))) return false;

        uint16_t pmcsr = (uint16_t)ConfigRead(bus, device, function, pm + PmcsrOffset, 2);
        int current = pmcsr & 0x3;
        if (current == state) return true;

        // Keep PME_Status (bit 15, write-1-to-clear) from being cleared by
        // accident; only the state field changes.
        pmcsr = (uint16_t)((pmcsr & ~0x8003u) | (uint16_t)state);
        ConfigWrite(bus, device, function, pm + PmcsrOffset, 2, pmcsr);

        // PCI PM 1.2 table 5-6: 10 ms into or out of D3hot, 200 us for D2.
        if (state == 3 || current == 3)      Timekeeping::BusyWaitMs(10);
        else if (state == 2 || current == 2) Timekeeping::BusyWaitUs(200);

        return GetPowerState(bus, device, function) == state;
    }

    // -------------------------------------------------------------------------
    // PCI driver matching
    // -------------------------------------------------------------------------

    static bool MatchDriver(const PciDriverDesc& drv, const PciDevice& dev) {
        if (drv.VendorId != 0 && drv.VendorId != dev.VendorId)
            return false;
        if (drv.ClassCode != 0xFF && drv.ClassCode != dev.ClassCode)
            return false;
        if (drv.SubClass != 0xFF && drv.SubClass != dev.SubClass)
            return false;
        if (drv.ProgIf != 0xFF && drv.ProgIf != dev.ProgIf)
            return false;

        if (drv.DeviceIdCount > 0 && drv.DeviceIds != nullptr) {
            bool found = false;
            for (uint16_t i = 0; i < drv.DeviceIdCount; i++) {
                if (drv.DeviceIds[i] == dev.DeviceId) {
                    found = true;
                    break;
                }
            }
            if (!found) return false;
        }

        return true;
    }

    void ProbeAll(const PciDriverDesc* table, uint16_t count, ProbePhase phase) {
        for (uint64_t d = 0; d < g_devices.size(); d++) {
            PciDevice& dev = g_devices[d];
            if (dev.Claimed || dev.Reserved)
                continue;

            for (uint16_t i = 0; i < count; i++) {
                if (table[i].Phase != phase)
                    continue;
                if (!MatchDriver(table[i], dev))
                    continue;

                KernelLogStream(INFO, "PCI") << "Probing " << table[i].Name
                    << " for device " << base::hex
                    << (uint64_t)dev.Bus << ":" << (uint64_t)dev.Device
                    << "." << (uint64_t)dev.Function;

                if (table[i].Probe(dev)) {
                    dev.Claimed = true;
                    KernelLogStream(OK, "PCI") << table[i].Name << " claimed device "
                        << base::hex << (uint64_t)dev.Bus << ":"
                        << (uint64_t)dev.Device << "." << (uint64_t)dev.Function;
                    break; // device claimed, move to next device
                }
            }
        }
    }

    // -------------------------------------------------------------------------
    // Public API
    // -------------------------------------------------------------------------

    void Initialize(Hal::ACPI::CommonSDTHeader* xsdt) {
        KernelLogStream(INFO, "PCI") << "Initializing PCI subsystem";

        // A plain ECAM walk of a quirky config space reads (and can fault
        // on) host-bridge registers. Its devices are described by ACPI
        // instead (Drivers/Init.cpp's platform devices).
        char quirk[32];
        if (RootBridgeEcamQuirk(quirk, sizeof(quirk))) {
            KernelLogStream(WARNING, "PCI") << "Root bridge config space is not ECAM (quirk \""
                << quirk << "\"): not enumerating PCI";
            return;
        }

        // Try to parse MCFG for ECAM access; fall back to legacy if unavailable
        ParseMcfg(xsdt);

        if (g_useEcam) {
            KernelLogStream(OK, "PCI") << "Using ECAM (memory-mapped) config access";
        } else {
            KernelLogStream(INFO, "PCI") << "Using legacy I/O port config access";
        }

        // Enumerate all PCI devices
        EnumerateAll();

        KernelLogStream(OK, "PCI") << "Enumeration complete: " << base::dec << (uint64_t)g_devices.size() << " devices found";
    }
};