#include <IndustryStandard/Pci.h>
#include <Uefi.h>
#include <Uefi/UefiSpec.h>

typedef struct {
    BOOLEAN isUsing;
    BOOLEAN isIDE;  // T:IDE F:AHCI
    UINT8 bus;
    UINT8 device;
    UINT8 function;
} DiskController;

DiskController controllers[5];

void init_pci() {
    for (UINT8 bus = 0; bus < 256; bus++) {
        for (UINT8 device = 0; device < 32; device++) {
            for (UINT8 function = 0; function < 8; function++) {
                UINT32 ConfigData =
                    PciCf8Read32(PCI_CF8_LIB_ADDRESS(bus, device, function, 0));
                if (ConfigData != 0xffffffff) {
                    ConfigData = PciCf8Read32(
                        PCI_CF8_LIB_ADDRESS(bus, device, function, 8));
                    if ((ConfigData >> 16) == 0x0101) {  // IDE Controller
                        int n = 0;
                        while (controllers[n].isUsing) {
                            n++;
                        }
                        controllers[n].isUsing = TRUE;
                        controllers[n].isIDE = TRUE;
                        controllers[n].bus = bus;
                        controllers[n].device = device;
                        controllers[n].function = function;
                        return;
                    } else if ((ConfigData >> 8) ==
                               0x010601) {  // AHCI Controller
                        int n = 0;
                        while (controllers[n].isUsing) {
                            n++;
                        }
                        controllers[n].isUsing = TRUE;
                        controllers[n].isIDE = FALSE;
                        controllers[n].bus = bus;
                        controllers[n].device = device;
                        controllers[n].function = function;
                    }
                }
            }
        }
    }
}