//
// Created by rick on 02-02-21.
//

#include "pci.h"

#include <types.h>
#include <drivers/ports.h>
#include <libc/libc.h>
#include <kprint.h>
#include <libk.h>

#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_VENDOR_ID            0x00
#define PCI_CONFIG_DEVICE_ID            0x02
#define PCI_CONFIG_LINE_1               0x01
#define PCI_CONFIG_COMMAND              0x04
#define PCI_CONFIG_STATUS               0x06
#define PCI_CONFIG_LINE_2               0x08
#define PCI_CONFIG_REVISION_ID          0x08
#define PCI_CONFIG_PROG_IF              0x09
#define PCI_CONFIG_SUBCLASS             0x0A
#define PCI_CONFIG_CLASS_CODE           0x0B
#define PCI_CONFIG_LINE_3               0x0C
#define PCI_CONFIG_CACHE_LINE_SIZE      0x0C
#define PCI_CONFIG_LATENCY_TIMER        0x0D
#define PCI_CONFIG_HEADER_TYPE          0x0E
#define PCI_CONFIG_BIST                 0x0F
#define PCI_CONFIG_BAR0                 0x10
#define PCI_CONFIG_BAR1                 0x14
#define PCI_CONFIG_BAR2                 0x18
#define PCI_CONFIG_BAR3                 0x1C
#define PCI_CONFIG_BAR4                 0x20
#define PCI_CONFIG_BAR5                 0x24
#define PCI_CONFIG_CARDBUS_CIS_P        0x28
#define PCI_CONFIG_SUBSYSTEM_VENDOR_ID  0x2C
#define PCI_CONFIG_SUBSYSTEM_ID         0x2E
#define PCI_CONFIG_EXPANSION_ROM_ADDR   0x30
#define PCI_CONFIG_CAP_POINTER          0x34
#define PCI_CONFIG_INTERRUPT_LINE       0x3C
#define PCI_CONFIG_INTERRUPT_PIN        0x3D
#define PCI_CONFIG_MAX_GRANT            0x3E
#define PCI_CONFIG_MAX_LATENCY          0x3F

#define PCI_HEADER_TYPE_MULTI_FUNC 0x80

#define MAX_PCI_DRIVERS 64
#define MAX_PCI_DEVICES 64

typedef struct {
    pci_driver driver;
    const char *driver_name;
    u8 pci_class;
    u8 pci_subclass;
    union {
        u8 flags;
        struct {
            u8 pci_use_subclass: 1;
        };
    };
} pci_driver_info;

typedef struct {
    u8 bus;
    u8 slot;
    u8 func;
    union {
        struct {
            u16 vendorId;
            u16 deviceId;
        };
        u32 config_line_0;
    };
    union {
        struct {
            u8 revisionId;
            u8 programInterface;
            u8 subclass;
            u8 class;
        };
        u32 config_line_2;
    };
    union {
        struct {
            u8 cacheLineSize;
            u8 latencyTimer;
            u8 headerType;
            u8 bist;
        };
        u32 config_line_3;
    };
    pci_driver_info *driver_info;
} pci_device;

pci_driver_info pci_drivers[MAX_PCI_DRIVERS];
int last_pci_device_index = 0;
pci_device pci_devices[MAX_PCI_DEVICES];

u32 pci_register_driver(u8 pci_class, u8 pci_subclass, u8 flags, pci_driver driver, const char *driver_name) {
    for (int i = 0; i < MAX_PCI_DRIVERS; ++i) {
        if (pci_drivers[i].driver == NULL) {
            pci_drivers[i].driver = driver;
            pci_drivers[i].pci_class = pci_class;
            pci_drivers[i].pci_subclass = pci_subclass;
            pci_drivers[i].flags = flags;
            pci_drivers[i].driver_name = driver_name;
            return PCI_REGISTER_OK;
        }
    }
    return PCI_REGISTER_ERR_FULL;
}


u32 pci_config_read_double_word(u8 bus, u8 slot, u8 func, u8 offset) {
    u32 address = PCI_CONFIG_ENABLE
                  | ((u32) bus << PCI_CONFIG_SHIFT_BUS_NUMBER)
                  | ((u32) slot << PCI_CONFIG_SHIFT_DEV_NUMBER)
                  | ((u32) func << PCI_CONFIG_SHIFT_FUNC_NUMBER)
                  | (offset & 0xFC);
    port_double_word_out(PORT_PCI_CONFIG_ADDRESS, address);
    return port_double_word_in(PORT_PCI_CONFIG_DATA);
}

u16 pci_config_read_word(u8 bus, u8 slot, u8 func, u8 offset) {
    u32 address = PCI_CONFIG_ENABLE
                  | ((u32) bus << PCI_CONFIG_SHIFT_BUS_NUMBER)
                  | ((u32) slot << PCI_CONFIG_SHIFT_DEV_NUMBER)
                  | ((u32) func << PCI_CONFIG_SHIFT_FUNC_NUMBER)
                  | (offset & 0xFC);
    port_double_word_out(PORT_PCI_CONFIG_ADDRESS, address);
    return port_double_word_in(PORT_PCI_CONFIG_DATA) >> ((offset & 2) * 8) & 0xFFFF;
}

u8 pci_config_read_byte(u8 bus, u8 slot, u8 func, u8 offset) {
    u32 address = PCI_CONFIG_ENABLE
                  | ((u32) bus << PCI_CONFIG_SHIFT_BUS_NUMBER)
                  | ((u32) slot << PCI_CONFIG_SHIFT_DEV_NUMBER)
                  | ((u32) func << PCI_CONFIG_SHIFT_FUNC_NUMBER)
                  | (offset & 0xFC);
    port_double_word_out(PORT_PCI_CONFIG_ADDRESS, address);
    return port_double_word_in(PORT_PCI_CONFIG_DATA) >> ((offset & 0b11) * 8) & 0xFF;
}

u16 pci_get_vendor_id(u8 bus, u8 slot, u8 func) {
    return pci_config_read_word(bus, slot, func, PCI_CONFIG_VENDOR_ID);
}

u16 pci_get_device_id(u8 bus, u8 slot, u8 func) {
    return pci_config_read_word(bus, slot, func, PCI_CONFIG_DEVICE_ID);
}

u8 pci_get_prog_if(u8 bus, u8 slot, u8 func) {
    return pci_config_read_byte(bus, slot, func, PCI_CONFIG_PROG_IF);
}

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

u8 pci_get_class_code(u8 bus, u8 slot, u8 func) {
    return pci_config_read_byte(bus, slot, func, PCI_CONFIG_CLASS_CODE);
}

u8 pci_get_subclass_code(u8 bus, u8 slot, u8 func) {
    return pci_config_read_byte(bus, slot, func, PCI_CONFIG_SUBCLASS);
}

void print_u8(u8 val) {
    char buf[3];
    itoa(val, buf, 16);
    kprint(buf);
}

void print_u16(u16 val) {
    char buf[5];
    itoa(val, buf, 16);
    kprint(buf);
}

void pci_pick_driver(pci_device *device) {
    for (int i = 0; i < MAX_PCI_DRIVERS; ++i) {
        if (pci_drivers[i].driver == NULL) {
            continue;
        }
        if (device->class == pci_drivers[i].pci_class) {
            if (!pci_drivers[i].pci_use_subclass || device->subclass == pci_drivers[i].pci_subclass) {
                pci_drivers[i].driver(device->bus, device->slot, device->func);
                device->driver_info = &pci_drivers[i];
            }
        }
    }
}

void pci_check_function(u8 bus, u8 slot, u8 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_pick_driver(&pci_devices[last_pci_device_index]);
    last_pci_device_index++;

    // todo do something with the function
}

void pci_check_device(u8 bus, u8 device) {
    u8 function = 0;
    u16 vendor_id = pci_get_vendor_id(bus, device, function);
    if (vendor_id == 0xFFFF) return;
    pci_check_function(bus, device, function);
    u8 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(u8 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!");
    }
    u8 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_print_info() {
    for (int i = 0; i < last_pci_device_index; ++i) {
        kprint("PCI BSF: ");
        print_u8(pci_devices[i].bus);
        kprint("/");
        print_u8(pci_devices[i].slot);
        kprint("/");
        print_u8(pci_devices[i].func);
        kprint(", Class/Sub/If: ");
        print_u8(pci_devices[i].class);
        kprint("/");
        print_u8(pci_devices[i].subclass);
        kprint("/");
        print_u8(pci_devices[i].programInterface);

        kprint(", V/D: ");
        print_u16(pci_devices[i].vendorId);
        kprint("/");
        print_u16(pci_devices[i].deviceId);
        kprint(" driver_info: ");

        if (pci_devices[i].driver_info == NULL) {
            kprint("NULL");
        } else {
            kprint(pci_devices[i].driver_info->driver_name);
        }

        kprint("\n");
    }
}

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