//
// Created by rick on 02-02-21.
//
// https://wiki.osdev.org/PCI

#include <sys/types.h>
#include <stdio.h>

#include <myke/drivers/pci/pci.h>
#include <myke/drivers/ports.h>
#include <myke/libk/libk.h>

#ifdef ENABLE_PCIPP

#include <myke/drivers/pci/pci_devices.h>

#endif

#define PCI_CONFIG_ENABLE (1 << 31)

#define PCI_CONFIG_SHIFT_BUS_NUMBER 16
#define PCI_CONFIG_SHIFT_DEV_NUMBER 11
#define PCI_CONFIG_SHIFT_FUNC_NUMBER 8

#define PCI_CONFIG_LINE_0               0x00
#define PCI_CONFIG_LINE_1               0x01
#define PCI_CONFIG_LINE_2               0x08
#define PCI_CONFIG_LINE_3               0x0C


#define MAX_PCI_DRIVERS 64
#define MAX_PCI_DEVICES 64

#define MASK_BAR_IOSPACE 0xFFFFFFFC
#define MASK_BAR_MEMSPACE 0xFFFFFFF0

//const pci_driver *pci_drivers[MAX_PCI_DRIVERS];
extern struct pci_driver __start_pci_driver[];
extern struct pci_driver __stop_pci_driver[];
#define NUM_DRIVERS ((size_t)(__stop_pci_driver - __start_pci_driver))
#define DRIVER(i) ((__start_pci_driver) + (i))
int last_pci_device_index = 0;
pci_device pci_devices[MAX_PCI_DEVICES];

void pci_check_bus(uint8_t bus);

uint8_t used pci_secondary_bus_use(const pci_device *device) {
    uint8_t secondary_bus = pci_config_read_byte(device->bus, device->slot, device->func,
                                                 PCI_CONFIG_SECONDARY_BUS_NUMBER);
    pci_check_bus(secondary_bus);
    return PCI_USE_OK;
}

uint32_t pci_config_address(uint8_t bus, uint8_t slot, uint8_t func, uint8_t offset) {
    return PCI_CONFIG_ENABLE
           | ((uint32_t) bus << PCI_CONFIG_SHIFT_BUS_NUMBER)
           | ((uint32_t) slot << PCI_CONFIG_SHIFT_DEV_NUMBER)
           | ((uint32_t) func << PCI_CONFIG_SHIFT_FUNC_NUMBER)
           | (offset & 0xFC);
}

uint32_t pci_config_read_double_word(uint8_t bus, uint8_t slot, uint8_t func, uint8_t offset) {
    uint32_t address = pci_config_address(bus, slot, func, offset);
    port_double_word_out(PORT_PCI_CONFIG_ADDRESS, address);
    return port_double_word_in(PORT_PCI_CONFIG_DATA);
}

uint16_t pci_config_read_word(uint8_t bus, uint8_t slot, uint8_t func, uint8_t offset) {
    uint32_t address = pci_config_address(bus, slot, func, offset);
    port_double_word_out(PORT_PCI_CONFIG_ADDRESS, address);
    return port_double_word_in(PORT_PCI_CONFIG_DATA) >> ((offset & 2) * 8) & 0xFFFF;
}

uint8_t pci_config_read_byte(uint8_t bus, uint8_t slot, uint8_t func, uint8_t offset) {
    uint32_t address = pci_config_address(bus, slot, func, offset);
    port_double_word_out(PORT_PCI_CONFIG_ADDRESS, address);
    return port_double_word_in(PORT_PCI_CONFIG_DATA) >> ((offset & 0b11) * 8) & 0xFF;
}

void pci_config_write_double_word(uint8_t bus, uint8_t slot, uint8_t func, uint8_t offset, uint32_t value) {
    uint32_t address = pci_config_address(bus, slot, func, offset);
    port_double_word_out(PORT_PCI_CONFIG_ADDRESS, address);
    port_double_word_out(PORT_PCI_CONFIG_DATA, value);
}

void pci_config_write_word(uint8_t bus, uint8_t slot, uint8_t func, uint8_t offset, uint16_t value) {
    uint32_t address = pci_config_address(bus, slot, func, offset);
    port_double_word_out(PORT_PCI_CONFIG_ADDRESS, address);
    port_word_out(PORT_PCI_CONFIG_DATA, value);
}

void pci_config_write_byte(uint8_t bus, uint8_t slot, uint8_t func, uint8_t offset, uint8_t value) {
    uint32_t address = pci_config_address(bus, slot, func, offset);
    port_double_word_out(PORT_PCI_CONFIG_ADDRESS, address);
    port_byte_out(PORT_PCI_CONFIG_DATA, value);
}

pci_command_register_t pci_get_config(pci_device *device) {
    pci_command_register_t status;
    status.value = pci_config_read_word(device->bus, device->slot, device->func, PCI_CONFIG_COMMAND);
    return status;
}

void pci_set_status(pci_device *device, pci_status_register_t status) {
    pci_config_write_word(device->bus, device->slot, device->func, PCI_CONFIG_STATUS, status.value);
}

pci_status_register_t pci_get_status(pci_device *device) {
    pci_status_register_t status;
    status.value = pci_config_read_word(device->bus, device->slot, device->func, PCI_CONFIG_STATUS);
    return status;
}

uint16_t pci_get_vendor_id(uint8_t bus, uint8_t slot, uint8_t func) {
    return pci_config_read_word(bus, slot, func, PCI_CONFIG_VENDOR_ID);
}

uint8_t pci_get_header_type(uint8_t bus, uint8_t slot, uint8_t func) {
    return pci_config_read_byte(bus, slot, func, PCI_CONFIG_HEADER_TYPE);
}

void pci_pick_driver(pci_device *device) {
    for (size_t i = 0; i < NUM_DRIVERS; ++i) {
        struct pci_driver *driver = DRIVER(i);
        if (driver == NULL) {
            continue;
        }
        if (driver->mask.class && driver->match.class != device->class) {
            continue;
        }
        if (driver->mask.subclass && driver->match.subclass != device->subclass) {
            continue;
        }
        if (driver->mask.interface && driver->match.interface != device->programInterface) {
            continue;
        }
        if (driver->mask.vendor && driver->match.vendor != device->vendorId) {
            continue;
        }
        if (driver->mask.device && driver->match.device != device->deviceId) {
            continue;
        }
        if (driver->direct_use) {
            if (driver->use(device) != PCI_USE_OK) {
                continue;
            }
        } else if (driver->validatable) {
            if (driver->validate(device) != PCI_VALIDATE_OK) {
                continue;
            }
        } else {
            continue;
        }
        device->pci_driver = DRIVER(i);
    }
}

void pci_check_function(uint8_t bus, uint8_t slot, uint8_t func) {
    pci_devices[last_pci_device_index].bus = bus;
    pci_devices[last_pci_device_index].slot = slot;
    pci_devices[last_pci_device_index].func = func;
    pci_devices[last_pci_device_index].config_line_0 = pci_config_read_double_word(bus, slot, func, PCI_CONFIG_LINE_0);
    pci_devices[last_pci_device_index].config_line_2 = pci_config_read_double_word(bus, slot, func, PCI_CONFIG_LINE_2);
    pci_devices[last_pci_device_index].config_line_3 = pci_config_read_double_word(bus, slot, func, PCI_CONFIG_LINE_3);
    pci_devices[last_pci_device_index].device_state.present = 1;
    last_pci_device_index++;

    pci_pick_driver(&pci_devices[last_pci_device_index - 1]);

    // todo do something with the function
}

void pci_check_device(uint8_t bus, uint8_t device) {
    uint8_t function = 0;
    uint16_t vendor_id = pci_get_vendor_id(bus, device, function);
    if (vendor_id == 0xFFFF) return;
    pci_check_function(bus, device, function);
    uint8_t header_type = pci_get_header_type(bus, device, function);
    if (header_type & PCI_HEADER_TYPE_MULTI_FUNC) {
        for (function = 1; function < 8; ++function) {
            if (pci_get_vendor_id(bus, device, function) != 0xFFFF) {
                pci_check_function(bus, device, function);
            }
        }
    }
}

void pci_check_bus(uint8_t bus) {
    for (int device = 0; device < 32; ++device) {
        pci_check_device(bus, device);
    }
}

void pci_scan() {
    if (last_pci_device_index != 0) {
        k_panics("Can only scan once!");
    }
    uint8_t header_type = pci_get_header_type(0, 0, 0);
    if (header_type & PCI_HEADER_TYPE_MULTI_FUNC) {
        // multiple pci host controllers
        for (int function = 0; function < 8; ++function) {
            pci_check_bus(function);
        }
    } else {
        pci_check_bus(0);
    }
}

void pci_init_drivers() {
    for (int device_index = 0; device_index < MAX_PCI_DEVICES; ++device_index) {
        if (!pci_devices[device_index].device_state.present) {
            continue;
        }
        if (pci_devices[device_index].driver_state.initialized) {
            continue; // already done
        }
        if (pci_devices[device_index].pci_driver == NULL) {
            continue; // no driver found
        }
        if (pci_devices[device_index].pci_driver->initialisable) {
            pci_devices[device_index].pci_driver->initialize(&pci_devices[device_index]);
        }
        pci_devices[device_index].driver_state.initialized = 1;
    }
}

void pci_print_info() {
    for (int i = 0; i < last_pci_device_index; ++i) {
        printf("PCI BSF: %2x/%2x/%2x, CSI: %2x/%2x/%2x, V/D: %4x/%4x, driver: %s\n",
               pci_devices[i].bus, pci_devices[i].slot, pci_devices[i].func,
               pci_devices[i].class, pci_devices[i].subclass, pci_devices[i].programInterface,
               pci_devices[i].vendorId, pci_devices[i].deviceId,
               (pci_devices[i].pci_driver == NULL ? "none" : pci_devices[i].pci_driver->name));
    }
}

void pci_dump_caps_internal(pci_device *device) {
    if (pci_devices->headerType >= 0x02) {
        printf("\tNot supported for PCI-to-CardBus bridge\n");
        return;
    }
    if (!pci_get_status(device).status.capabilities_list) {
        printf("\tNo caps\n");
        return;
    }
    uint8_t cap_ptr = pci_config_read_byte(device->bus, device->slot, device->func, PCI_CONFIG_CAP_POINTER);
    printf("\t%x\n", cap_ptr);
    // todo traverse
}

void pci_dump_caps() {
    for (int i = 0; i < last_pci_device_index; ++i) {
        printf("PCI BSF: %2x/%2x/%2x, CSI: %2x/%2x/%2x, V/D: %4x/%4x, driver: %s\n",
               pci_devices[i].bus, pci_devices[i].slot, pci_devices[i].func,
               pci_devices[i].class, pci_devices[i].subclass, pci_devices[i].programInterface,
               pci_devices[i].vendorId, pci_devices[i].deviceId,
               (pci_devices[i].pci_driver == NULL ? "none" : pci_devices[i].pci_driver->name));
        pci_dump_caps_internal(&pci_devices[i]);
    }
}

#ifdef ENABLE_PCIPP

void pci_pretty_print() {
    for (int i = 0; i < last_pci_device_index; ++i) {
        char *class = NULL;
        char *subclass = NULL;
        char *function = NULL;
        pci_device_info *class_info = NULL;
        pci_device_info *subclass_info = NULL;
        pci_device_info *function_info = NULL;
        pci_device_info *current = pci_root_info;
        while (current->name != NULL) {
            if (current->code == pci_devices[i].class) {
                class = current->name;
                class_info = current;
                break;
            }
            current = current + 1;
        }
        if (class_info != NULL && class_info->sub != NULL) {
            current = class_info->sub;
            while (current->name != NULL) {
                if (current->code == pci_devices[i].subclass) {
                    subclass = current->name;
                    subclass_info = current;
                    break;
                }
                current = current + 1;
            }
        }
        if (subclass_info != NULL && subclass_info->sub != NULL) {
            current = subclass_info->sub;
            while (current->name != NULL) {
                if (current->code == pci_devices[i].programInterface) {
                    function = current->name;
                    function_info = current;
                    break;
                }
                current = current + 1;
            }
        }
        printf("PCI BSF: %2x/%2x/%2x %s %s %s\n",
               pci_devices[i].bus, pci_devices[i].slot, pci_devices[i].func,
               class, subclass, function);
    }
}

#endif

void pci_init_bar(pci_device *device, uint8_t bar_index) {
    if (device->headerType != 0x00) {
        k_panics("Only header 0x00 supported for now");
        return;
    }
    uint8_t offset = 0;
    bar_info *bar;
    switch (bar_index) {
        case 0:
            offset = PCI_CONFIG_BAR0;
            bar = &device->bar0;
            break;
        case 1:
            offset = PCI_CONFIG_BAR1;
            bar = &device->bar1;
            break;
        case 2:
            offset = PCI_CONFIG_BAR2;
            bar = &device->bar2;
            break;
        case 3:
            offset = PCI_CONFIG_BAR3;
            bar = &device->bar3;
            break;
        case 4:
            offset = PCI_CONFIG_BAR4;
            bar = &device->bar4;
            break;
        case 5:
            offset = PCI_CONFIG_BAR5;
            bar = &device->bar5;
            break;
        default:
            k_panics("Bar index too high");
            return;
    }


    uint32_t original_address = pci_config_read_double_word(device->bus, device->slot, device->func, offset);
    pci_config_write_double_word(device->bus, device->slot, device->func, offset, 0xFFFFFFFF);
    uint32_t masked_size = pci_config_read_double_word(device->bus, device->slot, device->func, offset);

    if (original_address & 0x1) {
        // IO Space
        bar->size = ~(masked_size & MASK_BAR_IOSPACE) + 1;
        bar->address = original_address & MASK_BAR_IOSPACE;
        bar->is_io_space = 1;
    } else {
        // Memory space
        bar->size = ~(masked_size & MASK_BAR_MEMSPACE) + 1;
        bar->address = original_address & 0xFFFFFFF0;
        bar->is_io_space = 0;
        bar->type = (original_address & 0b110) >> 1;
        bar->prefetchable = (original_address & 0b1000) >> 3;
    }
    // todo identify conflicts, move to different address
    pci_config_write_double_word(device->bus, device->slot, device->func, offset, original_address);
    bar->present = 1;
}

// internal drivers
PCI_DRIVER(0) = {
        .name = "pci-secondary-bus",
        .description = "A PCI bus connected to the primary bus",
        .match.class = PCI_CLASS_BRIDGE,
        .match.subclass = PCI_SUB_CLASS_PCI_PCI_BRIDGE_4,
        .mask.class = true,
        .mask.subclass = true,
        .direct_use = true,
        .use = pci_secondary_bus_use,
};

// todo https://wiki.osdev.org/Universal_Serial_Bus if i dare