#include "boot.hpp"

#include <algorithm>
#include <bit>
#include <cassert>
#include <cstdio>

#include "os/draw.hpp"
#include "os/memory.hpp"
#include "os/segments.hpp"
#include "os/serial.hpp"
#include "os/syscall.hpp"
#include "os/tty.hpp"
#include "drivers/pci.hpp"
#include "drivers/pic.hpp"
#include "drivers/ps2.hpp"
#include "drivers/usb.hpp"
#include "./x86_64.hpp"

namespace
{
inline constexpr std::uint8_t MASTER_PIC_OFFSET = 0x20;
inline constexpr std::uint8_t SLAVE_PIC_OFFSET = 0x28;
static_assert(isValidPICOffset(MASTER_PIC_OFFSET));
static_assert(isValidPICOffset(SLAVE_PIC_OFFSET));

constexpr std::uint8_t irqToInterrupt(std::uint8_t irq) noexcept
{
    if (irq < 8) {
        return MASTER_PIC_OFFSET + irq;
    }
    return SLAVE_PIC_OFFSET + irq;
}
inline constexpr std::uint8_t INTERRUPT_KEYBOARD = irqToInterrupt(1);

TaskStateSegment tss alignas(16);

constinit std::array GDT alignas(16) = {
    SEGMENT_NULL,
    SEGMENT_KERNEL_CODE,
    SEGMENT_KERNEL_DATA,
    SEGMENT_USER_DATA,
    SEGMENT_USER_CODE,
    SEGMENT_NULL, // TSS lower
    SEGMENT_NULL  // TSS higher
};

constinit baos::InterruptDescriptorTable gInterruptDescriptorTable;

inline baos::MemoryType mapEfiMemoryType(UINT32 type) noexcept
{
    switch (type)
    {
        case EfiLoaderCode:
        case EfiLoaderData:
        case EfiBootServicesCode:
        case EfiBootServicesData:
        case EfiConventionalMemory:
            return baos::MemoryType::USABLE;
        case EfiRuntimeServicesCode:
        case EfiRuntimeServicesData:
        case EfiUnusableMemory:
        case EfiACPIReclaimMemory:
        case EfiMaxMemoryType:
        default:
            return baos::MemoryType::UNUSABLE;
        case EfiMemoryMappedIO:
        case EfiMemoryMappedIOPortSpace:
            return baos::MemoryType::MMIO;
        case EfiACPIMemoryNVS:
        case EfiPalCode:
        case EfiReservedMemoryType:
            return baos::MemoryType::RESERVED;
    }
}

std::vector<baos::MemoryRange> detectMemoryRangesFromEfi(const EfiMemoryMap& memoryMap) noexcept
{
    std::vector<baos::MemoryRange> ranges;

    ranges.reserve(memoryMap.mapSize / memoryMap.descriptorSize);
    for (const EFI_MEMORY_DESCRIPTOR& descriptor : memoryMap)
    {
        const baos::MemoryType type = mapEfiMemoryType(descriptor.Type);

        assert(descriptor.PhysicalStart % 4096 == 0);
        uint64_t pageBase = descriptor.PhysicalStart / 4096;
        uint64_t numPages = descriptor.NumberOfPages;
        ranges.push_back({.pageBase = pageBase, .numPages = numPages, .type = type});
    }
    std::sort(ranges.begin(), ranges.end());

    // merge adjacent ranges
    for (std::size_t idx = ranges.size() - 1; idx > 0; --idx)
    {
        baos::MemoryRange& left = ranges[idx-1];
        baos::MemoryRange& right = ranges[idx];

        if (left.type == right.type && left.pageBase + left.numPages == right.pageBase)
        {
            left.numPages += right.numPages;
            right.numPages = 0;
        }
    }
    std::erase_if(ranges, [](const baos::MemoryRange& range)
    {
        return range.numPages == 0;
    });
    return ranges;
}

using initfunc_t = void *();
extern "C" initfunc_t* start_ctors;
extern "C" initfunc_t* end_ctors;

void initGlobals()
{
    for (initfunc_t** initFunc = &start_ctors; initFunc != &end_ctors; ++initFunc)
    {
        // TODO: this confuses me
        if (*initFunc == reinterpret_cast<initfunc_t*>(0xffffffffffffffff)) {
            break;
        }
        (*initFunc)();
    }
}

extern "C" void __setGDT(uint16_t limit, void* base);
extern "C" void __reloadSegments(uint64_t code, uint64_t data);
extern "C" void __flushTSS(uint16_t segment);
extern "C" void __enterUsermode(uint16_t codeSegment, uint16_t dataSegment, void* rsp, void (*function)(), void* stackPtr);

void initializePCIDevice(const baos::pci::Header& header) noexcept
{
    using namespace baos;
    switch (header.baseClass)
    {
        case pci::BaseClass::SERIAL_BUS_CONTROLLER:
            switch (header.subClass)
            {
                case pci::SubClass::USB_CONTROLLER:
                    (void) usb::initController(header); // TODO: print something on error?
                    break;
            }
    }
}
}

extern "C"
{
EfiBootInfo* gBootInfo;

void usermode_main();

void kernel_main()
{
    using namespace baos;

    gInterruptDescriptorTable.setupErrorInterrupts();
    gInterruptDescriptorTable.setupInterrupt(INTERRUPT_KEYBOARD, &ps2::isrKeyboard);

    makeTSSSegment(&tss, &GDT[5]);
    __setGDT(sizeof(GDT), &GDT);
    gInterruptDescriptorTable.install();
    __reloadSegments(SEGIDX_KERNEL_CODE, SEGIDX_KERNEL_DATA);
    __flushTSS(SEGIDX_TSS);

    setupSyscall();

    const bool picsInited = initPICs(MASTER_PIC_OFFSET, SLAVE_PIC_OFFSET);
    unmaskIRQ(1);

    // initialize the hardware
    const bool ps2Inited = ps2::initialize();

    // done initializing OS stuff, enable interrupts
    __asm__ __volatile__("sti");

    // setup the "initial heap", required so we can use std::vector for detecting the memory ranges and stuff
    __ba_initInitialHeap();

    // init the heap (required for the double buffer)
    const std::vector<MemoryRange> ranges = detectMemoryRangesFromEfi(gBootInfo->memoryMap);
    if (!setupPaging({ranges.begin(), ranges.end()}))
    {
        panic("Error setting up paging.");
    }

    // make sure boot info and the frame buffer are accessible
    identityMapRegion(gBootInfo, sizeof(*gBootInfo));
    identityMapRegion(gBootInfo->displayInfo.frameBufferBase, gBootInfo->displayInfo.frameBufferPitch * gBootInfo->displayInfo.frameBufferHeight);

    initGlobals();

    // initialize the framebuffer
    draw::initializeDefaultFramebuffer({
        /* base   = */ static_cast<draw::Pixel*>(gBootInfo->displayInfo.frameBufferBase),
        /* width  = */ gBootInfo->displayInfo.frameBufferWidth,
        /* height = */ gBootInfo->displayInfo.frameBufferHeight,
        /* pitch  = */ gBootInfo->displayInfo.frameBufferPitch / 4
    });

	// initialize terminal interface
	tty::initialize();

    if (!initSerialPort(PORT_COM1))
    {
        tty::write("Error initializing serial port.\n");
        return;
    }
    serialWriteString(PORT_COM1, "\r\n\r\n"); // write newlines to separate from UEFI output

    for (const MemoryRange& range : ranges)
    {
        std::printf("Start: 0x%lX, Pages: %ld, Type: %s\n", range.pageBase << 12, range.numPages, memoryTypeName(range.type));
    }

    if (!picsInited)
    {
        std::puts("Error initializing the PICs!\n");
        return;
    }

    if (!ps2Inited)
    {
        std::puts("Error initializing the PS2 controller!\n");
        return;
    }

    for (const pci::Header header : pci::enumerateDevices())
    {
        initializePCIDevice(header);
    }

    static std::array<std::uint8_t, 16384> usermodeStack alignas(16);
    __enterUsermode(SEGIDX_USER_CODE, SEGIDX_USER_DATA, &tss.rsp0, &usermode_main, &usermodeStack.back() + 1);
}
} // extern "C"