#include "boot.hpp"

#include <bit>
#include <cstdio>

#include "os/draw.hpp"
#include "os/serial.hpp"
#include "os/tty.hpp"
#include "drivers/pic.hpp"
#include "drivers/ps2.hpp"
#include "./x86_64.hpp"

extern "C" int gKernelEnd;

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_DIVISION_ERROR = 0x00;
inline constexpr std::uint8_t INTERRUPT_OVERFLOW = 0x04;
inline constexpr std::uint8_t INTERRUPT_BOUND_RANGE_EXCEEDED = 0x05;
inline constexpr std::uint8_t INTERRUPT_INVALID_OPCODE = 0x06;
inline constexpr std::uint8_t INTERRUPT_DEVICE_NOT_AVAILABLE = 0x07;
inline constexpr std::uint8_t INTERRUPT_DOUBLE_FAULT = 0x08;
inline constexpr std::uint8_t INTERRUPT_INVALID_TSS = 0x0A;
inline constexpr std::uint8_t INTERRUPT_SEGMENT_NOT_PRESENT = 0x0B;
inline constexpr std::uint8_t INTERRUPT_STACK_SEGMENT_FAULT = 0x0C;
inline constexpr std::uint8_t INTERRUPT_GENERAL_PROTECTION_FAULT = 0x0D;
inline constexpr std::uint8_t INTERRUPT_PAGE_FAULT = 0x0E;
inline constexpr std::uint8_t INTERRUPT_X87_FP_EXCEPTION = 0x10;
inline constexpr std::uint8_t INTERRUPT_ALIGNMENT_CHECK = 0x11;
inline constexpr std::uint8_t INTERRUPT_MACHINE_CHECK = 0x12;
inline constexpr std::uint8_t INTERRUPT_SIMD_FP_EXCEPTION = 0x13;
inline constexpr std::uint8_t INTERRUPT_VIRTUALIZATION_EXCEPTION = 0x14;
inline constexpr std::uint8_t INTERRUPT_CONTROL_PROTECTION_EXCEPTION = 0x15;
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_CODE,
    SEGMENT_USER_DATA,
    SEGMENT_NULL, // TSS lower
    SEGMENT_NULL  // TSS higher
};

inline constexpr std::uint16_t SEGIDX_KERNEL_CODE = 1 << 3;
inline constexpr std::uint16_t SEGIDX_KERNEL_DATA = 2 << 3;
inline constexpr std::uint16_t SEGIDX_USER_CODE   = 3 << 3;
inline constexpr std::uint16_t SEGIDX_USER_DATA   = 4 << 3;
inline constexpr std::uint16_t SEGIDX_TSS         = 5 << 3;

constinit std::array<InterruptDescriptor, 256> IDT alignas(16);

void initHeapFromEfiMemoryMap(const EfiMemoryMap& memoryMap) noexcept
{
    const uint64_t minAddr = reinterpret_cast<uint64_t>(&gKernelEnd);
    for (const EFI_MEMORY_DESCRIPTOR& descriptor : memoryMap)
    {
        if (!isEfiMemoryTypeUsable(descriptor.Type))
        {
            continue;
        }

        uint64_t addr = descriptor.PhysicalStart;
        uint64_t len = descriptor.NumberOfPages * 4096; // TODO: is this the correct page size?
        // anything before the kernel we ignore
        if (addr < minAddr)
        {
            // if the entire entry is before the kernel, continue
            if (addr + len <= minAddr)
            {
                continue;
            }
            // otherwise shrink it and use the rest
            len -= (minAddr - addr);
            addr = minAddr;
        }
        __ba_registerAllocatableMemory(reinterpret_cast<void*>(addr), len);
    }
#if 0
    const uint64_t minAddr = reinterpret_cast<uint64_t>(&gKernelEnd);
    for (uint32_t addr = firstEntryAddr; addr < firstEntryAddr + bufferLength;)
    {
        multiboot_memory_map_t entry = *reinterpret_cast<multiboot_memory_map_t*>(addr);

        if(entry.type == MULTIBOOT_MEMORY_AVAILABLE)
        {
            std::printf("Start Addr: %X | Length: %b\n",
                        static_cast<unsigned>(entry.addr), static_cast<size_t>(entry.len));

            // anything before the kernel we ignore
            if (entry.addr < minAddr)
            {
                // if the entire entry is before the kernel, continue
                if (entry.addr + entry.len <= minAddr)
                {
                    addr += sizeof(entry.size) + entry.size;
                    continue;
                }
                // otherwise shrink it and use the rest
                entry.len -= (minAddr - entry.addr);
                entry.addr = minAddr;
            }
            __ba_registerAllocatableMemory(reinterpret_cast<void*>(entry.addr), entry.len);
        }
        addr += sizeof(entry.size) + entry.size;
    }
#endif
}

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 __setIDT(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)());

template<typename THandler>
void setupInterrupt(std::uint8_t index, THandler handler) noexcept
{
    const std::uint64_t offset = std::bit_cast<std::uint64_t>(handler);
    IDT[index] = InterruptDescriptor{
        .offsetLow = static_cast<unsigned>(offset & 0xFFFF),
        .selectorPrivilegeLevel = 0,
        .selectorIndex = 1, // kernel code
        .interruptStackTable = 0,
        .gateType = InterruptGateType::INTERRUPT_GATE,
        .privilegeLevel = 0,
        .present = true,
        .offsetHigh = (offset >> 16)
    };
}
}

extern "C"
{
EfiBootInfo* gBootInfo;

void main();

void kernel_main()
{
    using namespace baos;

    setupInterrupt(INTERRUPT_DIVISION_ERROR, &baos::isrDivisionError);
    setupInterrupt(INTERRUPT_OVERFLOW, &baos::isrOverflow);
    setupInterrupt(INTERRUPT_BOUND_RANGE_EXCEEDED, &baos::isrBoundRangeExceeded);
    setupInterrupt(INTERRUPT_INVALID_OPCODE, &baos::isrInvalidOpCode);
    setupInterrupt(INTERRUPT_DEVICE_NOT_AVAILABLE, &baos::isrDeviceNotAvailable);
    setupInterrupt(INTERRUPT_DOUBLE_FAULT, &baos::isrDoubleFault);
    setupInterrupt(INTERRUPT_INVALID_TSS, &baos::isrInvalidTSS);
    setupInterrupt(INTERRUPT_SEGMENT_NOT_PRESENT, &baos::isrSegmentNotPresent);
    setupInterrupt(INTERRUPT_STACK_SEGMENT_FAULT, &baos::isrStackSegmentFault);
    setupInterrupt(INTERRUPT_GENERAL_PROTECTION_FAULT, &baos::isrGeneralProtectionFault);
    setupInterrupt(INTERRUPT_PAGE_FAULT, &baos::isrPageFault);
    setupInterrupt(INTERRUPT_X87_FP_EXCEPTION, &baos::isrX87FPException);
    setupInterrupt(INTERRUPT_ALIGNMENT_CHECK, &baos::isrAlignmentCheck);
    setupInterrupt(INTERRUPT_MACHINE_CHECK, &baos::isrMachineCheck);
    setupInterrupt(INTERRUPT_SIMD_FP_EXCEPTION, &baos::isrSimdFpException);
    setupInterrupt(INTERRUPT_VIRTUALIZATION_EXCEPTION, &baos::isrVirtualizationException);
    setupInterrupt(INTERRUPT_CONTROL_PROTECTION_EXCEPTION, &baos::isrControlProtectionException);
    setupInterrupt(INTERRUPT_KEYBOARD, &ps2::isrKeyboard);

    makeTSSSegment(&tss, &GDT[5]);
    __setGDT(sizeof(GDT), &GDT);
    __setIDT(sizeof(IDT), &IDT);
    __reloadSegments(SEGIDX_KERNEL_CODE, SEGIDX_KERNEL_DATA);
    __flushTSS(SEGIDX_TSS);

    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");

    initGlobals();

    // init the heap (required for the double buffer)
    // initHeapFromEfiMemoryMap(gBootInfo->memoryMap);

    // 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

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

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

    std::puts("This is BadAppleOS and everything is fine!\n");

    __enterUsermode(SEGIDX_USER_CODE, SEGIDX_USER_DATA, &tss.rsp0, &main);
    while (true)
    {
        ps2::KeyEvent event;
        while (!ps2::readKey(event));
        std::printf(
            "Key event: scancode=%s(0x%X), down=%s, repeat=%s\n",
            ps2::keyName(event.scancode),
            static_cast<unsigned>(event.scancode),
            event.down ? "true" : "false",
            event.repeat ? "true" : "false"
        );
    }


}
} // extern "C"