// ke_debug.c - debugging utilities
#include <stdarg.h>
#include <stdint.h>
#include <string.h>
#include <ntlite/ddi/hal_api.h>
#include <kernel/ke_debug.h>
#include <kernel/ke_spinlock.h>
#include <kernel/ke_process.h>
#include <kernel/ke_thread.h>
#include <kernel/ke_sched.h>
#include <kernel/mm_virt.h>
#include <common/printf.h>
int _fltused; // fuckass

int gDebugEnabled = 0;
void (*PrintChar)(char, uint32_t, uint32_t);

void KiKernelPutChar(char c) {
    PrintChar(c, 0x00FF00, 0x000000);
}

void KeInitDebug(void (*PrintPtr)(char, uint32_t, uint32_t), int v) {
    gDebugEnabled = v;
    PrintChar = PrintPtr;
    printf_set_putchar(KiKernelPutChar);
}

__declspec(dllexport) void DbgPrint(const char* fmt, ...) {
    if (!gDebugEnabled) return;

    static SPINLOCK sl = {0};
    uintptr_t flags = KeAcquireSpinlock(&sl);
    va_list va;
    va_start(va, fmt);
    vprintf(fmt, va);
    va_end(va);
    KeReleaseSpinlock(&sl, flags);
}

// KePanic
__declspec(dllexport) _Noreturn void KePanic(int code, const char* fmt, ...) {
    // Silence! Stop the world!
    _disable();

    // Print a bright-red banner so the framebuffer is unmistakable.
    // We call PrintChar directly so we can set a red foreground.
    // (DbgPrint uses green; we want this visually distinct.)
    static const char banner[] = "\n*** KERNEL PANIC ***\n";
    for (const char* p = banner; *p; p++)
        PrintChar(*p, 0xFF2222, 0x000000);

    // Print the panic code first, then the caller's message through the normal vprintf path.
    printf("Panic code: %d\n", code);

    va_list va;
    va_start(va, fmt);
    vprintf(fmt, va);
    va_end(va);

    // Die!
    while (true) __halt();
}

// KePrintConsole - write user-visible console text in white-on-black.
// Used by SYS_WRITE_CONSOLE so user programs get a distinct color from
// kernel debug output (green) and panic output (red).
__declspec(dllexport) void KePrintConsole(const char* buf, size_t len) {
    if (!buf || !PrintChar) return;
    for (size_t i = 0; i < len; i++)
        PrintChar(buf[i], 0xFFFFFF, 0x000000);
}

// ---------------------------------------------------------------------------
// KeBugCheckProcess - user-mode fault handler
//
// Resolves the faulting RIP to a module name, prints a yellow crash banner,
// kills the offending process, then exits the current thread so the scheduler
// can continue running everything else.
// ---------------------------------------------------------------------------

// Return the short exception name for a CPU exception vector.
// Keep this table complete for vectors 0..31 so faults never collapse into
// a vague "general protection" bucket just because the IDT entry existed
// first or the crash printer lacked a name.
static const char* KiFaultVectorName(int vector) {
    switch (vector) {
        case 0:  return "#DE Divide Error";
        case 1:  return "#DB Debug";
        case 2:  return "NMI Interrupt";
        case 3:  return "#BP Breakpoint";
        case 4:  return "#OF Overflow";
        case 5:  return "#BR Bound Range Exceeded";
        case 6:  return "#UD Invalid Opcode";
        case 7:  return "#NM Device Not Available";
        case 8:  return "#DF Double Fault";
        case 9:  return "Coprocessor Segment Overrun";
        case 10: return "#TS Invalid TSS";
        case 11: return "#NP Segment Not Present";
        case 12: return "#SS Stack-Segment Fault";
        case 13: return "#GP General Protection";
        case 14: return "#PF Page Fault";
        case 15: return "Reserved Exception 15";
        case 16: return "#MF x87 Floating-Point Exception";
        case 17: return "#AC Alignment Check";
        case 18: return "#MC Machine Check";
        case 19: return "#XM SIMD Floating-Point Exception";
        case 20: return "#VE Virtualization Exception";
        case 21: return "#CP Control Protection Exception";
        case 22: return "Reserved Exception 22";
        case 23: return "Reserved Exception 23";
        case 24: return "Reserved Exception 24";
        case 25: return "Reserved Exception 25";
        case 26: return "Reserved Exception 26";
        case 27: return "Reserved Exception 27";
        case 28: return "#HV Hypervisor Injection Exception";
        case 29: return "#VC VMM Communication Exception";
        case 30: return "#SX Security Exception";
        case 31: return "Reserved Exception 31";
        default: return "Unknown Exception";
    }
}

static int KiFaultHasErrorCode(int vector) {
    switch (vector) {
        case 8:   // #DF
        case 10:  // #TS
        case 11:  // #NP
        case 12:  // #SS
        case 13:  // #GP
        case 14:  // #PF
        case 17:  // #AC
        case 21:  // #CP
        case 29:  // #VC
        case 30:  // #SX
            return 1;
        default:
            return 0;
    }
}

// Walk the process module list to find which image owns an address.
// Returns the module name string, or "unknown" if the address is outside all
// known image ranges.  Base/size are optional and are useful for offset dumps.
static const char* KiResolveModuleEx(PROCESS* proc, uintptr_t addr,
                                     uintptr_t* baseOut, size_t* sizeOut) {
    if (baseOut) *baseOut = 0;
    if (sizeOut) *sizeOut = 0;

    if (!proc) return "unknown";

    if (proc->ImageBase && proc->ImageSize &&
        addr >= proc->ImageBase && addr < proc->ImageBase + proc->ImageSize) {
        if (baseOut) *baseOut = proc->ImageBase;
        if (sizeOut) *sizeOut = proc->ImageSize;
        return proc->Name;
    }

    for (uint32_t i = 0; i < proc->ModuleCount; i++) {
        uintptr_t base = proc->Modules[i].Base;
        size_t    size = proc->Modules[i].Size;
        if (base && size && addr >= base && addr < base + size) {
            if (baseOut) *baseOut = base;
            if (sizeOut) *sizeOut = size;
            return proc->Modules[i].Name;
        }
    }

    return "unknown";
}

static const char* KiResolveModule(PROCESS* proc, uintptr_t addr) {
    return KiResolveModuleEx(proc, addr, (uintptr_t*)0, (size_t*)0);
}

static void KiPrintAddress(PROCESS* proc, const char* label, uintptr_t addr) {
    uintptr_t base = 0;
    size_t size = 0;
    const char* mod = KiResolveModuleEx(proc, addr, &base, &size);

    if (base) {
        printf("%s%016llx <%s+0x%llx>\n",
               label,
               (unsigned long long)addr,
               mod,
               (unsigned long long)(addr - base));
    } else {
        printf("%s%016llx <unknown>\n",
               label,
               (unsigned long long)addr);
    }
}

static int KiUserReadU64(uintptr_t addr, uint64_t* out) {
    if (!out) return 0;
    if (!addr) return 0;
    if (!KSUCCESS(MmProbeUserRead((const void*)addr, sizeof(uint64_t))))
        return 0;
    *out = *(const uint64_t*)addr;
    return 1;
}

static int KiUserReadU8(uintptr_t addr, uint8_t* out) {
    if (!out) return 0;
    if (!addr) return 0;
    if (!KSUCCESS(MmProbeUserRead((const void*)addr, sizeof(uint8_t))))
        return 0;
    *out = *(const uint8_t*)addr;
    return 1;
}

static void KiDumpBytes(PROCESS* proc, const char* label, uintptr_t center) {
    if (!center) {
        printf("%s: <null>\n", label);
        return;
    }

    printf("%s around ", label);
    KiPrintAddress(proc, "", center);

    uintptr_t start = center >= 16 ? center - 16 : 0;
    for (int row = 0; row < 2; row++) {
        uintptr_t addr = start + (uintptr_t)row * 16u;
        printf("  %016llx:", (unsigned long long)addr);
        for (int i = 0; i < 16; i++) {
            uint8_t b = 0;
            if (KiUserReadU8(addr + (uintptr_t)i, &b))
                printf(" %02x", b);
            else
                printf(" ??");
        }
        printf("\n");
    }
}

static int KiUserReadI32(uintptr_t addr, int32_t* out) {
    if (!out) return 0;
    if (!addr) return 0;
    if (!KSUCCESS(MmProbeUserRead((const void*)addr, sizeof(int32_t))))
        return 0;
    *out = *(const int32_t*)addr;
    return 1;
}

static void KiDumpIndirectCallSlot(PROCESS* proc, uintptr_t ret) {
    if (!ret) return;

    // Most x64 import/function-pointer calls are one of:
    //   FF 15 disp32       call qword ptr [rip+disp32]
    //   48 FF 15 disp32    rex.w call qword ptr [rip+disp32]
    // The saved return address points immediately after the instruction.
    uintptr_t insn = 0;
    int32_t disp = 0;
    uint8_t b0 = 0, b1 = 0, b2 = 0;

    if (KiUserReadU8(ret - 6, &b0) &&
        KiUserReadU8(ret - 5, &b1) &&
        b0 == 0xff && b1 == 0x15 &&
        KiUserReadI32(ret - 4, &disp)) {
        insn = ret - 6;
    } else if (KiUserReadU8(ret - 7, &b0) &&
               KiUserReadU8(ret - 6, &b1) &&
               KiUserReadU8(ret - 5, &b2) &&
               b0 == 0x48 && b1 == 0xff && b2 == 0x15 &&
               KiUserReadI32(ret - 4, &disp)) {
        insn = ret - 7;
    } else {
        printf("No RIP-relative indirect call pattern found before return address.\n");
        return;
    }

    uintptr_t slot = ret + (intptr_t)disp;
    uint64_t target = 0;

    printf("Decoded null indirect call:\n");
    KiPrintAddress(proc, "  insn=", insn);
    KiPrintAddress(proc, "  slot=", slot);

    if (KiUserReadU64(slot, &target)) {
        KiPrintAddress(proc, "  *slot=", (uintptr_t)target);
    } else {
        printf("  *slot=<unreadable>\n");
    }

    KiDumpBytes(proc, "Bytes near indirect slot", slot);
}

static void KiDumpModules(PROCESS* proc) {
    if (!proc) return;

    printf("Images:\n");
    if (proc->ImageBase && proc->ImageSize) {
        printf("  %-31s %016llx-%016llx size=%llx\n",
               proc->Name,
               (unsigned long long)proc->ImageBase,
               (unsigned long long)(proc->ImageBase + proc->ImageSize),
               (unsigned long long)proc->ImageSize);
    }

    for (uint32_t i = 0; i < proc->ModuleCount; i++) {
        uintptr_t base = proc->Modules[i].Base;
        size_t size = proc->Modules[i].Size;
        printf("  %-31s %016llx-%016llx size=%llx\n",
               proc->Modules[i].Name,
               (unsigned long long)base,
               (unsigned long long)(base + size),
               (unsigned long long)size);
    }
}

static void KiDumpUserStack(PROCESS* proc, uintptr_t rsp) {
    printf("User stack @ %016llx:\n", (unsigned long long)rsp);
    if (!rsp) {
        printf("  <no user RSP captured>\n");
        return;
    }

    for (int i = 0; i < 12; i++) {
        uintptr_t slot = rsp + (uintptr_t)i * 8u;
        uint64_t value = 0;
        if (KiUserReadU64(slot, &value)) {
            uintptr_t base = 0;
            size_t size = 0;
            const char* mod = KiResolveModuleEx(proc, (uintptr_t)value, &base, &size);
            if (base) {
                printf("  [%02d] %016llx : %016llx <%s+0x%llx>\n",
                       i,
                       (unsigned long long)slot,
                       (unsigned long long)value,
                       mod,
                       (unsigned long long)((uintptr_t)value - base));
            } else {
                printf("  [%02d] %016llx : %016llx\n",
                       i,
                       (unsigned long long)slot,
                       (unsigned long long)value);
            }
        } else {
            printf("  [%02d] %016llx : <unreadable>\n",
                   i,
                   (unsigned long long)slot);
        }
    }
}

static void KiDumpFaultFrame(PROCESS* proc, const KFAULT_FRAME* f) {
    if (!f) return;

    printf("Regs:\n");
    printf("  RAX=%016llx RBX=%016llx RCX=%016llx RDX=%016llx\n",
           (unsigned long long)f->Rax, (unsigned long long)f->Rbx,
           (unsigned long long)f->Rcx, (unsigned long long)f->Rdx);
    printf("  RSI=%016llx RDI=%016llx RBP=%016llx RSP=%016llx\n",
           (unsigned long long)f->Rsi, (unsigned long long)f->Rdi,
           (unsigned long long)f->Rbp, (unsigned long long)f->Rsp);
    printf("  R8 =%016llx R9 =%016llx R10=%016llx R11=%016llx\n",
           (unsigned long long)f->R8,  (unsigned long long)f->R9,
           (unsigned long long)f->R10, (unsigned long long)f->R11);
    printf("  R12=%016llx R13=%016llx R14=%016llx R15=%016llx\n",
           (unsigned long long)f->R12, (unsigned long long)f->R13,
           (unsigned long long)f->R14, (unsigned long long)f->R15);
    printf("  CS=%04llx SS=%04llx RFLAGS=%016llx\n",
           (unsigned long long)f->Cs,
           (unsigned long long)f->Ss,
           (unsigned long long)f->Rflags);

    if (f->Rip == 0) {
        uint64_t ret = 0;
        if (KiUserReadU64((uintptr_t)f->Rsp, &ret)) {
            printf("Likely null indirect call. Return address from [RSP]:\n");
            KiPrintAddress(proc, "  ret=", (uintptr_t)ret);
            KiDumpBytes(proc, "Bytes near return address", (uintptr_t)ret);
            KiDumpIndirectCallSlot(proc, (uintptr_t)ret);
        } else {
            printf("Likely null indirect call, but [RSP] is unreadable.\n");
        }
    } else {
        KiDumpBytes(proc, "Bytes near RIP", (uintptr_t)f->Rip);
    }

    KiDumpUserStack(proc, (uintptr_t)f->Rsp);
}

static void KiPrintCrashChar(char c) {
    PrintChar(c, 0xFFAA00, 0x000000);   // amber - distinct from green debug / red panic
}

static void KiPrintCrashHeader(PROCESS* proc, int vector, uintptr_t rip,
                               uintptr_t errCode, uintptr_t cr2) {
    const char* faultName = KiFaultVectorName(vector);
    const char* procName  = proc ? proc->Name  : "unknown";
    uint32_t    pid       = proc ? proc->ProcessId : 0;
    const char* modName   = proc ? KiResolveModule(proc, rip) : "unknown";

    static const char banner[] = "\n*** APPLICATION CRASH ***\n";
    for (const char* p = banner; *p; p++)
        KiPrintCrashChar(*p);

    printf_set_putchar(KiPrintCrashChar);

    if (vector == 14)
        printf("Process: %s (PID %u) | Fault: %s\n"
               "Module:  %s | RIP: %016llx | Code: %04llx | CR2: %016llx\n",
               procName, pid, faultName,
               modName,
               (unsigned long long)rip,
               (unsigned long long)errCode,
               (unsigned long long)cr2);
    else
        printf("Process: %s (PID %u) | Fault: %s\n"
               "Module:  %s | RIP: %016llx | Code: %04llx\n",
               procName, pid, faultName,
               modName,
               (unsigned long long)rip,
               (unsigned long long)errCode);
}

static _Noreturn void KiFinishFaultedProcess(PROCESS* proc, THREAD* thread) {
    printf_set_putchar(KiKernelPutChar);

    if (proc)   KeExitProcess(proc, -1);
    if (thread) KeExitThread(thread, -1);

    while (1) __halt();
}

_Noreturn void KeBugCheckProcess(int vector, uintptr_t rip, uintptr_t errCode, uintptr_t cr2) {
    THREAD*  thread = KeGetCurrentThread();
    PROCESS* proc   = thread ? thread->Owner : (PROCESS*)0;

    KiPrintCrashHeader(proc, vector, rip, errCode, cr2);
    KiDumpModules(proc);

    KiFinishFaultedProcess(proc, thread);
}

__declspec(dllexport) _Noreturn void KeFaultDispatchFrame(const KFAULT_FRAME* frame) {
    THREAD*  thread = KeGetCurrentThread();
    PROCESS* proc   = thread ? thread->Owner : (PROCESS*)0;

    if (!frame || frame->Magic != 0x3146544C5541464BULL) {
        KePanic(KE_PANIC_CODE_FAULT, "PANIC: bad KFAULT_FRAME %016llx\n",
                (unsigned long long)(uintptr_t)frame);
    }

    if ((frame->Cs & 3) == 0) {
        int vector = (int)frame->Vector;
        const char* name = KiFaultVectorName(vector);

        if (vector == 14) {
            KePanic(KE_PANIC_CODE_FAULT, "PANIC: %s | RIP: %016llx | Code: %04llx | CR2: %016llx\n",
                    name,
                    (unsigned long long)frame->Rip,
                    (unsigned long long)frame->ErrorCode,
                    (unsigned long long)frame->Cr2);
        } else if (KiFaultHasErrorCode(vector)) {
            KePanic(KE_PANIC_CODE_FAULT, "PANIC: %s | RIP: %016llx | Code: %04llx\n",
                    name,
                    (unsigned long long)frame->Rip,
                    (unsigned long long)frame->ErrorCode);
        } else {
            KePanic(KE_PANIC_CODE_FAULT, "PANIC: %s | RIP: %016llx\n",
                    name,
                    (unsigned long long)frame->Rip);
        }
    }

    KiPrintCrashHeader(proc, (int)frame->Vector,
                       (uintptr_t)frame->Rip,
                       (uintptr_t)frame->ErrorCode,
                       (uintptr_t)frame->Cr2);
    KiDumpFaultFrame(proc, frame);
    KiDumpModules(proc);

    KiFinishFaultedProcess(proc, thread);
}

// ---------------------------------------------------------------------------
// KeFaultDispatch - single C entry point for all three IDT fault stubs.
//
// The asm stubs (ke_syscall.asm) collect the CPU exception frame values and
// call here.  We check the CS RPL to decide: kernel fault → panic and halt,
// user fault → KeBugCheckProcess (kill process, resume scheduler).
// ---------------------------------------------------------------------------
__declspec(dllexport) _Noreturn void KeFaultDispatch(int vector, uintptr_t rip,
                                                     uintptr_t errCode, uintptr_t cr2,
                                                     uintptr_t cs) {
    KFAULT_FRAME frame;
    memset(&frame, 0, sizeof(frame));
    frame.Magic = 0x3146544C5541464BULL;
    frame.Vector = (uint64_t)vector;
    frame.ErrorCode = errCode;
    frame.Cr2 = cr2;
    frame.Rip = rip;
    frame.Cs = cs;
    KeFaultDispatchFrame(&frame);
}