#include "pci.h"
#include "arch/x86_64/io.h"
#include <stdio.h>

#define PCI_CFG_ADDR 0x0CF8
#define PCI_CFG_DATA 0x0CFC

static uint32_t pci_pack_addr(pci_address_t addr, uint16_t offset)
{
    return (1u << 31) |
           ((uint32_t)addr.bus << 16) |
           ((uint32_t)addr.device << 11) |
           ((uint32_t)addr.function << 8) |
           (offset & 0xFCu);
}

uint8_t pci_read8(pci_address_t addr, uint16_t offset)
{
    x86_64_outl(PCI_CFG_ADDR, pci_pack_addr(addr, offset));
    return (uint8_t)(x86_64_inl(PCI_CFG_DATA) >> ((offset & 3) * 8));
}

uint16_t pci_read16(pci_address_t addr, uint16_t offset)
{
    x86_64_outl(PCI_CFG_ADDR, pci_pack_addr(addr, offset));
    return (uint16_t)(x86_64_inl(PCI_CFG_DATA) >> ((offset & 2) * 8));
}

uint32_t pci_read32(pci_address_t addr, uint16_t offset)
{
    x86_64_outl(PCI_CFG_ADDR, pci_pack_addr(addr, offset));
    return x86_64_inl(PCI_CFG_DATA);
}

void pci_write8(pci_address_t addr, uint16_t offset, uint8_t val)
{
    uint32_t tmp = pci_read32(addr, offset & ~3u);
    uint8_t shift = (offset & 3) * 8;
    tmp = (tmp & ~(0xFFu << shift)) | ((uint32_t)val << shift);
    x86_64_outl(PCI_CFG_ADDR, pci_pack_addr(addr, offset));
    x86_64_outl(PCI_CFG_DATA, tmp);
}

void pci_write16(pci_address_t addr, uint16_t offset, uint16_t val)
{
    uint32_t tmp = pci_read32(addr, offset & ~3u);
    uint8_t shift = (offset & 2) * 8;
    tmp = (tmp & ~(0xFFFFu << shift)) | ((uint32_t)val << shift);
    x86_64_outl(PCI_CFG_ADDR, pci_pack_addr(addr, offset));
    x86_64_outl(PCI_CFG_DATA, tmp);
}

void pci_write32(pci_address_t addr, uint16_t offset, uint32_t val)
{
    x86_64_outl(PCI_CFG_ADDR, pci_pack_addr(addr, offset));
    x86_64_outl(PCI_CFG_DATA, val);
}

static void pci_read_bars(pci_device_t* dev)
{
    for (int i = 0; i < 6; i++) {
        uint16_t off = 0x10 + i * 4;
        dev->bar[i] = pci_read32(dev->addr, off);
        // If it's a 64-bit BAR we skip the next one (simple version)
        if ((dev->bar[i] & 0x7) == 0x4)
            i++;
    }
}

void pci_init(void)
{
    printf("[PCI] Scanning bus 0...\n");
    pci_print_all();
}

void pci_print_all(void)
{
    for (uint8_t bus = 0; bus < 1; bus++) {        // start with bus 0 only
        for (uint8_t dev = 0; dev < 32; dev++) {
            for (uint8_t func = 0; func < 8; func++) {
                pci_address_t addr = {bus, dev, func};

                uint16_t vendor = pci_read16(addr, 0x00);
                if (vendor == 0xFFFF) {
                    if (func == 0) break;   // no device
                    continue;
                }

                uint16_t device_id = pci_read16(addr, 0x02);
                uint16_t class_sub  = pci_read16(addr, 0x0A);
                uint8_t  prog_if   = pci_read8 (addr, 0x09);
                uint8_t  revision  = pci_read8 (addr, 0x08);

                pci_device_t d = {
                    .addr = addr,
                    .vendor_id = vendor,
                    .device_id = device_id,
                    .class_code = class_sub,
                    .prog_if = prog_if,
                    .revision = revision,
                };
                pci_read_bars(&d);

                printf("[PCI] %02x:%02x.%x  %04x:%04x  class %04x  rev %02x\n",
                       bus, dev, func, vendor, device_id, class_sub, revision);

                if (func == 0 && (pci_read8(addr, 0x0E) & 0x80) == 0)
                    break; // not multifunction
            }
        }
    }
}

bool pci_find_hda(pci_device_t* out_dev)
{
    // Same scan as above, but stop at first HDA device
    for (uint8_t bus = 0; bus < 1; bus++) {
        for (uint8_t dev = 0; dev < 32; dev++) {
            for (uint8_t func = 0; func < 8; func++) {
                pci_address_t addr = {bus, dev, func};
                uint16_t vendor = pci_read16(addr, 0x00);
                if (vendor == 0xFFFF) continue;

                uint16_t class_sub = pci_read16(addr, 0x0A);
                if (class_sub == (PCI_CLASS_AUDIO << 8) | PCI_SUBCLASS_HDA) {
                    out_dev->addr = addr;
                    out_dev->vendor_id = vendor;
                    out_dev->device_id = pci_read16(addr, 0x02);
                    out_dev->class_code = class_sub;
                    pci_read_bars(out_dev);
                    return true;
                }
            }
        }
    }
    return false;
}