/*
    * IDT.cpp
    * Intel Interrupt Descriptor Table implementation
    * Copyright (c) 2025 Daniel Hammer
*/

#include "IDT.hpp"

#include <Memory/Heap.hpp>
#include <Common/Panic.hpp>
#include <Platform/Registers.hpp>
#include <CppLib/Stream.hpp>
#include <Memory/PageFrameAllocator.hpp>
#include <Sched/Scheduler.hpp>
#include <Sched/CrashReport.hpp>
#include <Timekeeping/ApicTimer.hpp>

namespace Hal {
    constexpr auto InterruptGate = 0x8E;
    constexpr auto TrapGate = 0x8F;

    InterruptDescriptor* IDT;
    IDTRStruct IDTR{};

    const char* ExceptionStrings[] = {
        "Division Error",
        "Debug",
        "Non-Maskable Interrupt",
        "Breakpoint",
        "Overflow",
        "Bound Rage Exceeded",
        "Invalid Opcode",
        "Device Not Available",
        "Double Fault",
        "Coprocessor Segment Overrun",
        "Invalid TSS",
        "Segment Not Present",
        "Stack-Segment Fault",
        "General Protection Fault",
        "Page Fault",
        "Reserved",
        "x87 Floating-Point Exception",
        "Alignment Check",
        "Machine Check",
        "SMID Floating-Point Exception",
        "Virtualization Exception",
        "Control Protection Exception",
        "Reserved",
        "Reserved",
        "Reserved",
        "Reserved",
        "Reserved",
        "Hypervisor Injection Exception",
        "VMM Communication Exception",
        "Security Exception",
        "Reserved"
    };

    // Exceptions that push a hardware error code before IP/CS/FLAGS/SP/SS
    static bool ExceptionHasErrorCode(uint8_t vector) {
        return vector == 8 || (vector >= 10 && vector <= 14)
            || vector == 17 || vector == 21 || vector == 29 || vector == 30;
    }

    // Extract CS from the interrupt frame. For error-code exceptions, the
    // error code sits at offset 0, shifting IP to +8 and CS to +16.
    static uint64_t GetExceptionCS(uint8_t vector, System::PanicFrame* frame) {
        if (ExceptionHasErrorCode(vector)) {
            return *(uint64_t*)((uint8_t*)frame + 16);
        }
        return frame->CS;
    }

    template<size_t i>
    __attribute__((interrupt)) void ExceptionHandler(System::PanicFrame* frame)
    {
        uint64_t cs = GetExceptionCS(i, frame);
        bool fromUser = (cs & 3) == 3;

        // SWAPGS: if from user mode, GS base is user-defined.
        // Swap to kernel per-CPU GS base so scheduler calls work.
        if (fromUser) asm volatile("swapgs");

        // If the fault originated in user-mode (ring 3), kill the process
        // instead of panicking the entire system.
        if (fromUser && Sched::GetCurrentPid() >= 0) {
            auto* proc = Sched::GetCurrentProcessPtr();
            Kt::KernelLogStream(Kt::ERROR, "Exception")
                << ExceptionStrings[i] << " in process \""
                << proc->name << "\" (pid " << proc->pid
                << ") - process terminated";

            // Capture crash report
            CrashReport::Report rep;
            rep.valid = true;
            rep.pid = proc->pid;
            for (int j = 0; j < 64 && proc->name[j]; j++) rep.processName[j] = proc->name[j];
            rep.processName[63] = '\0';
            rep.exceptionVector = i;
            for (int j = 0; j < 32 && ExceptionStrings[i][j]; j++) rep.exceptionName[j] = ExceptionStrings[i][j];
            rep.exceptionName[31] = '\0';
            rep.instructionPointer = frame->IP;
            rep.codeSegment = frame->CS;
            rep.flags = frame->Flags;
            rep.stackPointer = frame->SP;
            rep.stackSegment = frame->SS;

            if (i == 0x0E) {
                // Page fault: read CR2 for faulting address
                asm volatile("mov %%cr2, %0" : "=r"(rep.faultingAddress));
                auto* pf = (System::PageFaultPanicFrame*)frame;
                rep.pfPresent = pf->PageFaultError.Present;
                rep.pfWrite = pf->PageFaultError.Write;
                rep.pfUser = pf->PageFaultError.User;
                rep.pfReservedWrite = pf->PageFaultError.ReservedWrite;
                rep.pfInstructionFetch = pf->PageFaultError.InstructionFetch;
                rep.pfProtectionKey = pf->PageFaultError.ProtectionKey;
                rep.pfShadowStack = pf->PageFaultError.ShadowStack;
                rep.pfSGX = pf->PageFaultError.SGX;
            } else {
                rep.faultingAddress = 0;
                rep.pfPresent = 0;
                rep.pfWrite = 0;
                rep.pfUser = 0;
                rep.pfReservedWrite = 0;
                rep.pfInstructionFetch = 0;
                rep.pfProtectionKey = 0;
                rep.pfShadowStack = 0;
                rep.pfSGX = 0;
            }
            rep.timestampTick = Timekeeping::GetTicks();

            CrashReport::AddReport(rep);
            Sched::SpawnCrashPad();

            Sched::ExitProcess();
            __builtin_unreachable();
        } else {
            frame->InterruptVector = i;
            Panic(ExceptionStrings[i], frame);
        }

        // Unreachable in practice (user faults exit, kernel faults panic),
        // but balance the SWAPGS for correctness.
        if (fromUser) asm volatile("swapgs");
    }

    void LoadIDT(IDTRStruct& idtr) {
        asm("lidt %0" : : "m"(idtr));
    }

    InterruptDescriptor* GetInterruptDescriptor(size_t index) {
        InterruptDescriptor* descriptor = (InterruptDescriptor*)(IDTR.Base + index * sizeof(InterruptDescriptor));
        return descriptor;
    }

    uint64_t GetHandlerAddress(InterruptDescriptor* descriptor) {
        uint64_t result{};

        result |= (uint64_t)descriptor->Offset1;
        result |= (uint64_t)descriptor->Offset2 << 16;
        result |= (uint64_t)descriptor->Offset3 << 32;

        return result;
    }

    void IDTEncodeInterrupt(size_t i, void* handler, uint8_t type_attr, uint8_t ist) {
        uint64_t offset = (uint64_t)handler;
        auto ptr = GetInterruptDescriptor(i);

        *ptr = InterruptDescriptor {
            .Offset1 = (uint16_t)(offset & 0x000000000000ffff),

            .Selector = 0x08,
            .IST = ist,
            .TypeAttributes = type_attr,

            .Offset2 = (uint16_t)((offset & 0x00000000ffff0000) >> 16),
            .Offset3 = (uint32_t)((offset & 0xffffffff00000000) >> 32),

            .Zero = 0x00
        };
    }

    template<int I, int N>
    struct SetHandler {
        static void run() {
            // Use IST1 for NMI (2) and Double Fault (8) so they get a
            // known-good stack even if the kernel stack has overflowed.
            uint8_t ist = (I == 2 || I == 8) ? 1 : 0;
            IDTEncodeInterrupt(I, (void*)ExceptionHandler<I>, TrapGate, ist);
            SetHandler<I+1,N>::run();
        }
    };

    template<int N>
    struct SetHandler<N,N> {static void run() {}};

    void IDTInitialize() {
        IDT = (InterruptDescriptor*)Memory::g_pfa->Allocate();
        Kt::KernelLogStream(Kt::DEBUG, "IDT") << "Allocated IDT at " << base::hex << (uint64_t)IDT;
        IDTR.Limit = (256 * sizeof(InterruptDescriptor)) - 1;
        IDTR.Base = (uint64_t)IDT;
        Kt::KernelLogStream(Kt::DEBUG, "IDT") << "Set IDTR Base to " << base::hex << IDTR.Base << " and Limit to " << base::hex << IDTR.Limit;

        SetHandler<0, 31>::run();

        Kt::KernelLogStream(Kt::OK, "Hal") << "Created exception interrupt vectors";

        LoadIDT(IDTR);

        Kt::KernelLogStream(Kt::OK, "Hal") << "Loaded new IDT";
    }

    void IDTReload() {
        LoadIDT(IDTR);
    }
};