/*
    * 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 <Api/Heap.hpp>
#include <Hal/SmpBoot.hpp>
#include <Timekeeping/ApicTimer.hpp>

namespace Hal {
    constexpr auto InterruptGate = 0x8E;

    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",
        "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;
    }

    static System::PanicFrame* GetExceptionRegs(uint8_t vector, System::PanicFrame* frame) {
        if (ExceptionHasErrorCode(vector)) {
            return (System::PanicFrame*)((uint8_t*)frame + sizeof(uint64_t));
        }
        return frame;
    }

    // Shared fatal-exception path: kill the faulting user process (with a
    // crash report) or panic the kernel. `frame` is the RAW interrupt frame
    // (error code at offset 0 for vectors that push one). Never returns.
    static void HandleFatalException(uint8_t i, 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) {
            // Interrupt gates arrive with IF clear. Full process teardown can
            // wait on device completions, sibling CPUs, and wall-clock-bounded
            // recovery paths, so leaving IF clear here can stop the BSP clock
            // and its local device IRQs indefinitely. GS is already the kernel
            // per-CPU base, and timer/IPI scheduling refuses to switch away
            // from a ring-0 frame, so nested hardware IRQs are safe now.
            asm volatile("sti" ::: "memory");

            auto* proc = Sched::GetCurrentProcessPtr();
            auto* regs = GetExceptionRegs(i, frame);
            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;
            int pn;
            for (pn = 0; pn < 63 && proc->name[pn]; pn++) rep.processName[pn] = proc->name[pn];
            rep.processName[pn] = '\0';
            rep.exceptionVector = i;
            int en;
            for (en = 0; en < 31 && ExceptionStrings[i][en]; en++) rep.exceptionName[en] = ExceptionStrings[i][en];
            rep.exceptionName[en] = '\0';
            rep.instructionPointer = regs->IP;
            rep.codeSegment = regs->CS;
            rep.flags = regs->Flags;
            rep.stackPointer = regs->SP;
            rep.stackSegment = regs->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(proc->pid);

            // Record a killed-by-signal exit code (256+sig) so a parent
            // blocked in SYS_WAITPID can tell a crash from a clean exit.
            {
                int sig = 11;                       /* SIGSEGV */
                if (i == 0x00) sig = 8;             /* #DE -> SIGFPE */
                else if (i == 0x06) sig = 4;        /* #UD -> SIGILL */
                else if (i == 0x10 || i == 0x13) sig = 8;  /* x87/SIMD FP */
                Sched::SetProcessExitCode(proc->pid, 256 + sig);
            }

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

    template<size_t i>
    __attribute__((interrupt)) void ExceptionHandler(System::PanicFrame* frame)
    {
        HandleFatalException(i, frame);
    }

    // Page faults get a dedicated handler with the proper error-code
    // signature (so GCC pops the error code before IRET) because, unlike
    // the generic handler, this one can RETURN: a non-present fault in the
    // user stack growth region maps a fresh zeroed page and retries the
    // faulting instruction. With the two-argument form, `frame` points past
    // the error code, directly at the saved IP.
    __attribute__((interrupt)) void PageFaultHandler(System::PanicFrame* frame, uint64_t errorCode)
    {
        bool fromUser = (frame->CS & 3) == 3;
        if (fromUser) asm volatile("swapgs");
        auto* cpu = Smp::TryGetCurrentCpuData();

        uint64_t cr2;
        asm volatile("mov %%cr2, %0" : "=r"(cr2));

        // Bit 0 of the error code: 0 = non-present page. Covers both user
        // pushes past the mapped stack and kernel accesses to not-yet-grown
        // user stack buffers passed into syscalls.
        if ((errorCode & 1) == 0 && cpu != nullptr && cpu->currentSlot >= 0
            && Sched::GetCurrentPid() >= 0) {
            if (montauk::abi::TryHandleAnonymousPageFault(cr2, errorCode) ||
                Sched::TryGrowUserStack(cr2)) {
                if (fromUser) asm volatile("swapgs");
                return;
            }
        }

        // Not a growable fault. Hand the RAW frame (error code at offset 0)
        // to the fatal path, which re-derives fromUser and swaps GS itself.
        if (fromUser) asm volatile("swapgs");
        HandleFatalException(0x0E, (System::PanicFrame*)((uint8_t*)frame - 8));
    }

    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;
            // Vector 14 uses the dedicated page fault handler (stack growth).
            void* handler = (I == 14) ? (void*)PageFaultHandler
                                      : (void*)ExceptionHandler<I>;
            IDTEncodeInterrupt(I, handler, InterruptGate, 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, 32>::run();

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

        LoadIDT(IDTR);

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

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