#define _DEFAULT_SOURCE

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <dirent.h>
#include <fcntl.h>
#include <unistd.h>
#include <errno.h>
#include <stdint.h>
#include <pciaccess.h>

static struct pci_device *pci_devices = NULL;
static int pci_device_count = 0;

static uint32_t read_dword(const uint8_t *data, size_t offset) {
    return (uint32_t)data[offset]
        | ((uint32_t)data[offset + 1] << 8)
        | ((uint32_t)data[offset + 2] << 16)
        | ((uint32_t)data[offset + 3] << 24);
}

static uint16_t read_word(const uint8_t *data, size_t offset) {
    return (uint16_t)data[offset]
        | ((uint16_t)data[offset + 1] << 8);
}

static int read_config(const char *path, uint8_t *data, size_t len) {
    int fd = open(path, O_RDONLY);
    if (fd < 0) {
        return -1;
    }
    size_t total = 0;
    while (total < len) {
        ssize_t n = read(fd, data + total, len - total);
        if (n < 0) {
            close(fd);
            return -1;
        }
        if (n == 0) {
            break;
        }
        total += n;
    }
    close(fd);
    return total == len ? 0 : -1;
}

static int probe_bar_size(int fd, unsigned int offset, int is_io, uint64_t *size) {
    uint32_t original, inverted;
    uint32_t mask = is_io ? 0xFFFFFFFC : 0xFFFFFFF0;

    if (lseek(fd, offset, SEEK_SET) < 0) {
        return -1;
    }
    if (read(fd, &original, sizeof(original)) != (ssize_t)sizeof(original)) {
        return -1;
    }

    uint32_t probe = 0xFFFFFFFF;
    if (lseek(fd, offset, SEEK_SET) < 0) {
        return -1;
    }
    if (write(fd, &probe, sizeof(probe)) != (ssize_t)sizeof(probe)) {
        lseek(fd, offset, SEEK_SET);
        write(fd, &original, sizeof(original));
        return -1;
    }

    if (lseek(fd, offset, SEEK_SET) < 0) {
        lseek(fd, offset, SEEK_SET);
        write(fd, &original, sizeof(original));
        return -1;
    }
    if (read(fd, &inverted, sizeof(inverted)) != (ssize_t)sizeof(inverted)) {
        lseek(fd, offset, SEEK_SET);
        write(fd, &original, sizeof(original));
        return -1;
    }

    lseek(fd, offset, SEEK_SET);
    write(fd, &original, sizeof(original));

    if ((inverted & mask) == 0) {
        *size = 0;
        return 0;
    }

    *size = ((uint64_t)(~(inverted & mask) & mask)) + 1;
    return 0;
}

static int parse_bars(int config_fd, const uint8_t *config, struct pci_device *dev) {
    unsigned int bar_idx = 0;
    unsigned int offset = 0x10;

    memset(dev->regions, 0, sizeof(dev->regions));

    while (bar_idx < 6 && offset <= 0x24) {
        uint32_t val = read_dword(config, offset);
        int is_io = (val & 0x1) != 0;

        if (is_io) {
            dev->regions[bar_idx].type = PCIREGION_TYPE_IO;
            dev->regions[bar_idx].is_IO = 1;
            dev->regions[bar_idx].bus_addr = (pciaddr_t)(val & 0xFFFFFFFC);
            if (probe_bar_size(config_fd, offset, 1, &dev->regions[bar_idx].size) != 0) {
                dev->regions[bar_idx].size = 0;
            }
            bar_idx += 1;
            offset += 4;
        } else {
            int is_64 = ((val >> 1) & 0x3) == 2;
            dev->regions[bar_idx].is_prefetchable = ((val >> 3) & 0x1) != 0;
            dev->regions[bar_idx].type = PCIREGION_TYPE_MEMORY;
            dev->regions[bar_idx].is_IO = 0;

            if (is_64) {
                if (offset + 4 > 0x24) {
                    break;
                }
                dev->regions[bar_idx].is_64 = 1;
                uint32_t val_hi = read_dword(config, offset + 4);
                dev->regions[bar_idx].bus_addr = ((pciaddr_t)val_hi << 32) | (val & 0xFFFFFFF0);

                uint64_t size_lo = 0, size_hi = 0;
                if (probe_bar_size(config_fd, offset, 0, &size_lo) == 0) {
                    probe_bar_size(config_fd, offset + 4, 0, &size_hi);
                }
                dev->regions[bar_idx].size = (size_hi << 32) | size_lo;
                bar_idx += 2;
                offset += 8;
            } else {
                dev->regions[bar_idx].bus_addr = (pciaddr_t)(val & 0xFFFFFFF0);
                if (probe_bar_size(config_fd, offset, 0, &dev->regions[bar_idx].size) != 0) {
                    dev->regions[bar_idx].size = 0;
                }
                bar_idx += 1;
                offset += 4;
            }
        }
    }

    return 0;
}

static int parse_device(const char *name, struct pci_device *dev) {
    unsigned int domain, bus, device, func;
    if (sscanf(name, "%x--%x--%x.%x", &domain, &bus, &device, &func) != 4) {
        return -1;
    }
    if (device > 0x1F || func > 0x7) {
        return -1;
    }

    char config_path[256];
    snprintf(config_path, sizeof(config_path), "/scheme/pci/%s/config", name);

    uint8_t config[256];
    if (read_config(config_path, config, sizeof(config)) != 0) {
        return -1;
    }

    dev->domain = domain;
    dev->bus = bus;
    dev->dev = device;
    dev->func = func;
    dev->vendor_id = read_word(config, 0x00);
    dev->device_id = read_word(config, 0x02);
    dev->revision = config[0x08];
    dev->class_code = config[0x0B];
    dev->subclass = config[0x0A];
    dev->prog_if = config[0x09];

    int config_fd = open(config_path, O_RDWR);
    if (config_fd < 0) {
        config_fd = open(config_path, O_RDONLY);
    }
    if (config_fd >= 0) {
        parse_bars(config_fd, config, dev);
        close(config_fd);
    }

    dev->_internal = NULL;
    return 0;
}

int pci_system_init(void) {
    if (pci_devices) {
        return 0;
    }

    DIR *dir = opendir("/scheme/pci");
    if (!dir) {
        return -1;
    }

    int capacity = 0;
    struct dirent *entry;
    while ((entry = readdir(dir)) != NULL) {
        if (entry->d_name[0] == '.') {
            continue;
        }

        struct pci_device probe;
        memset(&probe, 0, sizeof(probe));
        if (parse_device(entry->d_name, &probe) != 0) {
            continue;
        }

        if (pci_device_count >= capacity) {
            capacity = capacity ? capacity * 2 : 8;
            struct pci_device *new_devices = realloc(pci_devices, capacity * sizeof(struct pci_device));
            if (!new_devices) {
                closedir(dir);
                return -1;
            }
            pci_devices = new_devices;
        }

        pci_devices[pci_device_count++] = probe;
    }

    closedir(dir);
    return 0;
}

void pci_system_cleanup(void) {
    if (pci_devices) {
        free(pci_devices);
        pci_devices = NULL;
        pci_device_count = 0;
    }
}

struct pci_device *pci_device_find_by_slot(unsigned int domain,
                                             unsigned int bus,
                                             unsigned int dev,
                                             unsigned int func) {
    for (int i = 0; i < pci_device_count; i++) {
        if (pci_devices[i].domain == domain &&
            pci_devices[i].bus == bus &&
            pci_devices[i].dev == dev &&
            pci_devices[i].func == func) {
            return &pci_devices[i];
        }
    }
    return NULL;
}

int pci_device_probe(struct pci_device *dev) {
    char config_path[256];
    snprintf(config_path, sizeof(config_path),
             "/scheme/pci/%04x--%02x--%02x.%x/config",
             dev->domain, dev->bus, dev->dev, dev->func);

    uint8_t config[256];
    if (read_config(config_path, config, sizeof(config)) != 0) {
        return -1;
    }

    int config_fd = open(config_path, O_RDWR);
    if (config_fd < 0) {
        config_fd = open(config_path, O_RDONLY);
    }
    if (config_fd >= 0) {
        parse_bars(config_fd, config, dev);
        close(config_fd);
        return 0;
    }
    return -1;
}