// Kernel entry point. This file connects boot information to core subsystem initialization and scheduler start.

#include <stdint.h>
#include <string.h>
#include <intrin.h>
#pragma intrinsic(_disable)
#pragma intrinsic(_enable)
#include <ntlite/boot_info.h>
#include <kernel/mm_phys.h>
#include <kernel/mm_pfn.h>
#include <kernel/mm_virt.h>
#include <kernel/mm_swap.h>
#include <kernel/ob_object.h>
#include <kernel/ke_process.h>
#include <kernel/ke_thread.h>
#include <kernel/ke_sched.h>
#include <kernel/ke_cpu.h>
#include <kernel/ke_idt.h>
#include <kernel/ke_irq.h>
#include <kernel/ke_dpc.h>
#include <kernel/ke_syscall.h>
#include <ntlite/ddi/net_device.h>
#include <kernel/ke_debug.h>
#include <kernel/ke_util.h>
#include <kernel/ldr_image.h>
#include <ntlite/ddi/hal_api.h>
#include <kernel/io_file.h>
#include <kernel/ke_selftest.h>
#include <kernel/ke_smp.h>
#include <kernel/ke_input.h>
#include <kernel/ke_time.h>
#include <kernel/ke_display.h>
#include <kernel/ke_settings.h>
#include <kernel/ke_work.h>
#include <kernel/ke_system_root.h>
#include <common/ini_parser.h>

typedef struct {
    BOOT_INFO* bootInfo;
    const HAL_API* hal;
} SYSTEM_INIT_ARGS;

#define KI_SWAP_MIN_MB 4ull
#define KI_SWAP_MAX_MB 8192ull

// Parses a required unsigned decimal integer from settings.ini.
// Configuration errors are treated as boot-time corruption instead of guessed defaults.
static uint64_t KiParseRequiredUInt64Setting(const char* section, const char* key) {
    char value[64];
    KSTATUS s = KeGetSettingString(section, key, value, sizeof(value));
    if (!KSUCCESS(s)) {
        KePanic(KE_PANIC_CODE_INIT,
                "settings: missing required [%s] %s entry, status=0x%x\n",
                section,
                key,
                s);
    }

    uint64_t result = 0;
    const char* p = value;
    if (!*p) {
        KePanic(KE_PANIC_CODE_INIT,
                "settings: empty numeric value for [%s] %s\n",
                section,
                key);
    }

    while (*p) {
        if (*p < '0' || *p > '9') {
            KePanic(KE_PANIC_CODE_INIT,
                    "settings: invalid numeric value for [%s] %s: %s\n",
                    section,
                    key,
                    value);
        }

        uint64_t digit = (uint64_t)(*p - '0');
        if (result > (UINT64_MAX - digit) / 10ull) {
            KePanic(KE_PANIC_CODE_INIT,
                    "settings: numeric value for [%s] %s is too large: %s\n",
                    section,
                    key,
                    value);
        }

        result = result * 10ull + digit;
        ++p;
    }

    return result;
}

// Initializes swap according to settings.ini.
// A size of zero is the explicit opt-out. Any missing or invalid size is a broken configuration.
static void KiInitializeSwapFromSettings(void) {
    uint64_t swapMb = KiParseRequiredUInt64Setting("vm", "swap_size_mb");
    if (swapMb == 0) {
        DbgPrint("swap: disabled by settings\n");
        return;
    }

    if (swapMb < KI_SWAP_MIN_MB || swapMb > KI_SWAP_MAX_MB) {
        KePanic(KE_PANIC_CODE_INIT,
                "settings: [vm] swap_size_mb=%llu is outside the allowed range %llu..%llu MB\n",
                (unsigned long long)swapMb,
                (unsigned long long)KI_SWAP_MIN_MB,
                (unsigned long long)KI_SWAP_MAX_MB);
    }

    char pagefilePath[1024];
    if (!KeBuildSystemRootPath("pagefile.sys", pagefilePath, sizeof(pagefilePath))) {
        KePanic(KE_PANIC_CODE_INIT,
                "swap: failed to build system-root pagefile path\n");
    }

    uint64_t swapBytes = swapMb * 1024ull * 1024ull;
    KSTATUS s = MmSwapInitialize(SysSvFromCstr(pagefilePath), swapBytes);
    if (!KSUCCESS(s)) {
        KePanic(KE_PANIC_CODE_INIT,
                "swap: failed to initialize %llu MB pagefile at %s, status=0x%x\n",
                (unsigned long long)swapMb,
                pagefilePath,
                s);
    }
}

typedef KSTATUS (*DRIVER_ENTRY_FN)(const HAL_API *);

// HAL pointer cached at boot for subsystems that still need platform services.
static const HAL_API* gKmainHal;

// Implements the page In Bootloader Region operation.
static int KiPageInBootloaderRegion(const BOOT_INFO* bootInfo, PHYSICAL_ADDRESS page) {
    for (uint32_t i = 0; i < bootInfo->MemoryMap.RegionCount; i++) {
        const MEMORY_REGION* r = &bootInfo->MemoryMap.Regions[i];
        if (r->Type != MemoryBootloader) continue;

        uintptr_t start = r->Base & PAGE_MASK;
        uintptr_t end   = (r->Base + r->Length + PAGE_SIZE - 1) & PAGE_MASK;
        if ((uintptr_t)page >= start && (uintptr_t)page < end) {
            return 1;
        }
    }
    return 0;
}

// Implements the reclaim Bootloader Memory operation.
static void KiReclaimBootloaderMemory(const BOOT_INFO* bootInfo) {
    // Reclaim only the raw buffers backing boot-loaded .sys modules.
    // The kernel/HAL images, page tables, and bootloader code remain live.
    for (uint32_t i = 2; i < bootInfo->ModuleCount; i++) {
        const BOOT_MODULE* mod = &bootInfo->Modules[i];
        if (!mod->PhysBase || !mod->ImageSize) continue;

        uintptr_t base = (uintptr_t)(mod->PhysBase & PAGE_MASK);
        uintptr_t end  = (uintptr_t)((mod->PhysBase + mod->ImageSize + PAGE_SIZE - 1) & PAGE_MASK);

        for (uintptr_t page = base; page < end; page += PAGE_SIZE) {
            if (!KiPageInBootloaderRegion(bootInfo, (PHYSICAL_ADDRESS)page)) continue;
            MmPfnFree((PHYSICAL_ADDRESS)page);
        }
    }
}

// Phase Helper: Resolve Imports for core modules.
static void KiResolveSystemImports(BOOT_INFO* bootInfo) {
    // Only resolve imports for kernel.exe and hal.dll (typically indices 0 and 1).
    // The bootloader maps these, but drivers are now passed as raw physical files.
    uint32_t coreModuleCount = (bootInfo->ModuleCount > 2) ? 2 : bootInfo->ModuleCount;
    
    for (uint32_t i = 0; i < coreModuleCount; i++) {
        if (bootInfo->Modules[i].VirtBase != 0) {
            LdrResolveImports(bootInfo->Modules[i].VirtBase);
        }
    }
}


// Phase Helper: Initialize dynamically loaded .sys drivers.
static void KiDriversIniCb(const char* section, const char* key,
                          const char* val, void* ud) {
    (void)val;
    if (!strcmp(section, "runtime")) {
        char*** pArray = (char***)ud;
        size_t len = strlen(key) + 1;
        char* keyCopy = MmAllocatePool(len);
        if (keyCopy) {
            memcpy(keyCopy, key, len);
            arr_push(*pArray, keyCopy);
        }
    }
}

// Initializes drivers.
static void KiInitializeDrivers(BOOT_INFO* bootInfo, const HAL_API* hal) {
    DbgPrint("KiInitializeDrivers: Initializing drivers...\n");
    
    // boot drivers (provided by the bootloader).
    for (uint32_t i = 0; i < bootInfo->ModuleCount; i++) {
        BOOT_MODULE* mod = &bootInfo->Modules[i];
        
        // Skip modules that are already mapped (Kernel/HAL) 
        // or have no physical data (invalid).
        if (mod->VirtBase != 0 || mod->PhysBase == 0) continue;

        DbgPrint("KiInitializeDrivers: Loading boot driver %s\n", mod->Name);

        // Calculate virtual address of the raw file in the physical window.
        void* rawPeBuffer = (void*)(bootInfo->PhysicalMapBase + mod->PhysBase);

        // Map the PE image into kernel space.
        uintptr_t mappedBase = 0;
        KSTATUS st = LdrLoadKernelImage(rawPeBuffer, mod->ImageSize, &mappedBase);
        if (!KSUCCESS(st)) {
            DbgPrint("ERROR: Failed to map driver %s (0x%X)\n", mod->Name, st);
            continue;
        }

        // Update the struct so other modules can find it.
        mod->VirtBase = mappedBase;

        // Resolve imports while the image is still writable. Kernel driver images
        // stay writable until the loader grows real section-backed mappings; the
        // loader still clears NX for the whole image before returning it.
        if (!KSUCCESS(LdrResolveImports(mappedBase))) {
            DbgPrint("ERROR: Failed to resolve imports for %s\n", mod->Name);
            continue;
        }

        // Register with the global module list so subsequent drivers can import from this one.
        LdrRegisterModule(mod->Name, mod->VirtBase);

        // Execute Entry Point.
        DRIVER_ENTRY_FN entry = (DRIVER_ENTRY_FN)LdrGetEntryPoint(mappedBase);
        if (entry) {
            KSTATUS ds = entry(hal);
            if (!KSUCCESS(ds)) {
                DbgPrint("WARNING: Driver %s returned failure status 0x%X\n", mod->Name, ds);
            }
        }
    }

    // runtime drivers.
    KSTATUS s;
    FILE_OBJECT* f;

    s = KeInitializeSystemRootVolume();
    if (s != STATUS_SUCCESS) {
        KePanic(KE_PANIC_CODE_INIT, "KiInitializeDrivers: Failed to find system root volume (status %d)\n", s);
    }

    char root[4];
    root[0] = (char)('0' + KeGetSystemRootDrive());
    root[1] = ':';
    root[2] = '/';
    root[3] = '\0';
    s = IoSetCwd(SysSvFromCstr(root));
    if (s != STATUS_SUCCESS) {
        KePanic(KE_PANIC_CODE_INIT, "KiInitializeDrivers: Failed to set system CWD to %s (status %d)\n", root, s);
    }

    s = KeOpenSystemRootFile("drivers\\drivers.ini", &f);
    if (s != STATUS_SUCCESS) {
        KePanic(KE_PANIC_CODE_INIT, "KiInitializeDrivers: Failed to open drivers.ini on system root (status %d)\n", s);
    }

    uint64_t size = IoGetFileSize(f) + 1;
    char* mem = MmAllocatePool(size);
    size_t got;
    if (!KSUCCESS(IoReadFile(f, mem, size, &got))) {
        IoCloseFile(f);
        KePanic(KE_PANIC_CODE_INIT, "KiInitializeDrivers: Failed to read drivers.ini\n");
    }
    mem[got] = '\0';

    IoCloseFile(f);

    char** runtimeKeys = NULL;
    IniParse(mem, KiDriversIniCb, &runtimeKeys);

    for (uint32_t i = 0; i < arr_len(runtimeKeys); i++) {
        DbgPrint("Found runtime driver: %s\n", runtimeKeys[i]);

        char path[1024];
        char relative[512];
        size_t len = strlen(runtimeKeys[i]);
        if (len + 9 > sizeof(relative)) {
            DbgPrint("KiInitializeDrivers: runtime driver path too long: %s\n", runtimeKeys[i]);
            continue;
        }
        memcpy(relative, "drivers\\", 8);
        memcpy(relative + 8, runtimeKeys[i], len + 1);

        if (!KeBuildSystemRootPath(relative, path, sizeof(path))) {
            DbgPrint("KiInitializeDrivers: failed to build runtime driver path: %s\n", runtimeKeys[i]);
            continue;
        }

        s = IoOpenFile(SysSvFromCstr(path), &f);
        if (s != STATUS_SUCCESS) {
            DbgPrint("KiInitializeDrivers: Failed to open runtime driver: %s (%s)\n", runtimeKeys[i], path);
            continue;
        }

        uint64_t drvSize = IoGetFileSize(f);
        void* drvMem = MmAllocatePool(drvSize);
        if (!KSUCCESS(IoReadFile(f, drvMem, drvSize, &got))) {
            DbgPrint("KiInitializeDrivers: Failed to read runtime driver: %s\n", runtimeKeys[i]);
            IoCloseFile(f);
            continue;
        }
        IoCloseFile(f);

        uintptr_t mappedBase = 0;
        if (!KSUCCESS(LdrLoadKernelImage(drvMem, drvSize, &mappedBase))) { // Load into kernel space
            DbgPrint("KiInitializeDrivers: Failed to load runtime driver: %s\n", runtimeKeys[i]);
            MmFreePool(drvMem);
            continue;
        }
        MmFreePool(drvMem);

        if (!KSUCCESS(LdrResolveImports(mappedBase))) {
            DbgPrint("KiInitializeDrivers: Failed to resolve imports for runtime driver: %s\n", runtimeKeys[i]);
            continue;
        }
        // Runtime modules may also export DDIs to drivers loaded later.  Boot
        // drivers were already registered above; keep the module model uniform.
        LdrRegisterModule(runtimeKeys[i], mappedBase);

        DRIVER_ENTRY_FN entry = (DRIVER_ENTRY_FN)(LdrGetEntryPoint(mappedBase));

        KSTATUS ds = entry(hal);
        if (!KSUCCESS(ds)) {
            DbgPrint("KiInitializeDrivers: Driver %s failed to init!\n", runtimeKeys[i]);
        }
    }

    // Free the individual strings allocated via strdup in the callback.
    if (runtimeKeys) {
        for (uint32_t i = 0; i < arr_len(runtimeKeys); i++) {
            if (runtimeKeys[i]) {
                MmFreePool(runtimeKeys[i]); 
            }
        }
        // Free the dynamic array structure itself.
        arr_free(runtimeKeys);
    }

    // Free the file buffer.
    MmFreePool(mem);
}

// Phase Helper: Load and launch a process from the filesystem.
// Returns the new PROCESS* with one reference owned by the caller.
// Returns NULL on failure.
static PROCESS* KiLaunchUserProcess(const char* filePath, const char* args) {
    KSTATUS s;
    FILE_OBJECT* f;

    // Kernel-launched user images are rooted at the selected system volume.
    s = KeOpenSystemRootFile(filePath, &f);
    if (s != STATUS_SUCCESS) {
        DbgPrint("Failed to load system file: \"%s\" (status %d)!\n", filePath, s);
        return 0;
    }

    uint64_t size = IoGetFileSize(f);
    (void)size;

    // Create the process object.
    PROCESS* proc;
    THREAD* thread;

    s = KeCreateProcess(filePath, &proc);
    if (!KSUCCESS(s)) {
        DbgPrint("Failed to create process. System halted.\n");
        IoCloseFile(f);
        return 0;
    }
    DbgPrint("Created process (PID %u) with command line: %s\n", proc->ProcessId, (args && args[0]) ? args : filePath);

    s = KeSetupUserProcessEnvironment(proc,
                                      SysSvFromCstr(filePath),
                                      SysSvFromCstr(args),
                                      SysSvFromCstr(proc->CurrentDirectory));
    if (!KSUCCESS(s)) {
        DbgPrint("KiLaunchUserProcess: Failed to build user process environment. Status: %d\n", s);
        IoCloseFile(f);
        ObDereferenceObject(&proc->Header);
        return 0;
    }

    // Map the executable through an image section backed by the open file.
    uintptr_t imageBase = 0;
    s = LdrLoadImageFileForProcess(proc, f, &imageBase);
    IoCloseFile(f);
    if (!KSUCCESS(s)) {
        DbgPrint("KiLaunchUserProcess: LdrLoadImageFileForProcess failed with status: %d\n", s);
        ObDereferenceObject(&proc->Header);
        return 0;
    }

    s = LdrResolveImportsForProcess(proc, imageBase);
    if (!KSUCCESS(s)) {
        DbgPrint("KiLaunchUserProcess: Failed to resolve EXE imports for %s (status %d)\n", filePath, s);
        ObDereferenceObject(&proc->Header);
        return 0;
    }

    s = LdrProtectImageForProcess(proc, imageBase, true);
    if (!KSUCCESS(s)) {
        DbgPrint("KiLaunchUserProcess: Failed to protect EXE image for %s (status %d)\n", filePath, s);
        ObDereferenceObject(&proc->Header);
        return 0;
    }

    proc->ImageBase = imageBase;
    proc->ImageSize = LdrGetMappedImageSize(proc, imageBase);

    uintptr_t entryPoint = LdrGetEntryPoint(imageBase);

    // Create the main thread last. KeCreateThread queues it before returning.
    s = KeCreateThread(proc, (uintptr_t)entryPoint, 0, PriorityNormal, &thread);
    if (!KSUCCESS(s)) {
        DbgPrint("KiLaunchUserProcess: Failed to create thread. Status: %d\n", s);
        proc->ImageBase = 0;
        proc->ImageSize = 0;
        ObDereferenceObject(&proc->Header);
        return 0;
    }
    uint32_t tid = thread->ThreadId;
    DbgPrint("KiLaunchUserProcess: Thread created and queued! TID=%u\n", tid);
    ObDereferenceObject(&thread->Header);

    // Caller owns the reference returned by KeCreateProcess.
    return proc;
}

// Implements the system Init Thread operation.
void SystemInitThread(void* arg) {
    SYSTEM_INIT_ARGS* args = (SYSTEM_INIT_ARGS*)arg;

    DbgPrint("SystemInitThread: Thread started. Igniting the system...\n");

    KiRunEarlyKernelSelfTests();

    KiInitializeDrivers(args->bootInfo, args->hal);
    KeTimeReloadSettings();

    KiInitializeSwapFromSettings();

    KiResolveSystemImports(args->bootInfo);
    KiReclaimBootloaderMemory(args->bootInfo);

    if (args->bootInfo->Debug & KE_DEBUG_SELFTEST)
        KiRunKernelSelfTests(args->hal);

    // Launch User-Mode init.exe.
    PROCESS* shellProc = KiLaunchUserProcess("system/init.exe", args->bootInfo->CommandLine);
    if (!shellProc) {
        DbgPrint("SystemInitThread: Failed to start system/init.exe.\n");
        KePanic(KE_PANIC_CODE_INIT, "SystemInitThread: Failed to start system/init.exe.\n");
    }

    KeWaitForEvent(shellProc->ExitEvent, 0);
    int exitCode = shellProc->ExitCode;
    if (exitCode != 0) {
        KePanic(KE_PANIC_CODE_INIT, "SystemInitThread: init.exe exited with code %d.\n", exitCode);
    }
    DbgPrint("SystemInitThread: init.exe exited with code %d.\n", exitCode);
    ObDereferenceObject(&shellProc->Header);

    if (args->bootInfo->Debug & (KE_DEBUG_CONSOLE | KE_DEBUG_FILE_LOG)) {
        char logPath[1024];
        KSTATUS logStatus = STATUS_NOT_FOUND;
        if (KeBuildSystemRootPath("kernel.log", logPath, sizeof(logPath))) {
            logStatus = KeFlushDebugLogToFile(SysSvFromCstr(logPath));
            if (!KSUCCESS(logStatus))
                DbgPrint("debuglog: shutdown flush to %s failed st=0x%x\n", logPath, logStatus);
        } else {
            DbgPrint("debuglog: shutdown flush skipped; system root path unavailable\n");
        }
    }
    gKmainHal->Power.Shutdown();
    MmFreePool(args);
}

// Main Kernel Entry.
extern void KernelMain(BOOT_INFO* bootInfo) {

    // Physical & Virtual Memory initialization.

    if (!MmInitializePhysical(&bootInfo->MemoryMap, bootInfo->PhysicalMapBase)) {
        _disable(); while (true) __halt(); // there's nothing we can do
    }

    PHYSICAL_ADDRESS pml4 = (PHYSICAL_ADDRESS)__readcr3();
    // The early virtual allocator starts small because it only needs bootstrap page tables before the PMM is online.
    MmInitializeVirtual(pml4, 1 * 1024 * 1024);

    // Bootstrap CPU and architecture setup (GDT/IDT/Faults).

    KeInitializeBootstrapCpu();
    KeInitializeGdt();
    KeInitializeIdt();

    // Register boot modules (Allows LdrGetExport to work for drivers and the HAL)
    // Note: Kept to a single call here to avoid the duplicate in your original code.

    LdrRegisterModulesFromBootInfo(bootInfo);

    // HAL Initialization.

    HAL_INIT_PARAMS halParams = {
        .PhysicalWindowBase = bootInfo->PhysicalMapBase,
        .FramebufferBase    = bootInfo->FramebufferBase,
        .FramebufferWidth   = bootInfo->FramebufferWidth,
        .FramebufferHeight  = bootInfo->FramebufferHeight,
        .FramebufferStride  = bootInfo->FramebufferStride,
        .RsdpPhysical       = bootInfo->RsdpPhysical,
        .Debug              = bootInfo->Debug & 1, // bitmask 1 debug prints
        .AllocatePool       = MmAllocatePool
    };

    KSTATUS s = HalInitialize(&halParams);
    if (!KSUCCESS(s)) {
        _disable(); while (true) __halt(); // there's nothing we can do.
    }

    const HAL_API* hal = HalGetApi();
    gKmainHal = hal;

    KeInitDebug(hal->Screen.PrintChar, bootInfo->Debug);
    KeTimeInitialize(bootInfo, hal);
    KeDisplayInitialize(hal);
    KeSetDebugConsoleWriter(KeDisplayDebugPrintChar);

    // CPU startup intentionally lives after HAL initialization but before IRQs
    // and scheduler startup. HAL owns ACPI/MADT discovery and LAPIC INIT/SIPI;
    // the kernel owns KCPU records, per-CPU stacks, GDT/TSS/IDT state, and the
    // AP park loop. External device IRQs remain routed to the BSP during this
    // step; APs only become runnable after scheduler admission below.
    KeSetApplicationProcessorTrampoline((uint64_t)bootInfo->ApTrampolinePhys);
    KSTATUS smpStatus = KeStartApplicationProcessors();
    if (!KSUCCESS(smpStatus)) {
        KePanic(KE_PANIC_CODE_INIT, "SMP: failed to start application processors, status=0x%x\n", smpStatus);
    }

    KeIrqInitialize(hal);
    KeInitializeDpcSystem();

    KeInputInitialize(hal);

    KeIrqRegister(hal->Timer.IrqNumber, hal->Timer.Isr);

    HalRegisterScheduler(KeSchedulerTick);

    KeSyscallInit(hal);

    // Dynamic Linking & Driver Execution.

    KiResolveSystemImports(bootInfo);
    //KiInitializeDrivers(bootInfo, hal);.

    // Scheduler Prep & Status Update.

    KeInitializeScheduler();

    KSTATUS workStatus = KeStartWorkQueue();
    if (!KSUCCESS(workStatus)) {
        KePanic(KE_PANIC_CODE_INIT, "workq: failed to start passive work queue, status=0x%x\n", workStatus);
    }

    KSTATUS smpSchedStatus = KeAdmitApplicationProcessorsToScheduler();
    if (!KSUCCESS(smpSchedStatus)) {
        KePanic(KE_PANIC_CODE_INIT, "SMP: failed to admit APs to scheduler, status=0x%x\n", smpSchedStatus);
    }

    KiRunCoreHardeningSelfTests();

    // Pack the arguments for the init thread.
    SYSTEM_INIT_ARGS* initArgs = MmAllocatePool(sizeof(SYSTEM_INIT_ARGS));
    initArgs->bootInfo = bootInfo;
    initArgs->hal      = hal;

    // Spawn the Initialization Thread. No creator reference is needed; the
    // scheduler owns the live thread lifetime.
    KSTATUS st = KeCreateKernelThread((uintptr_t)SystemInitThread,
                                      (uintptr_t)initArgs,
                                      PriorityNormal, (THREAD**)0);
    
    if (!KSUCCESS(st)) {
        KePanic(KE_PANIC_CODE_INIT, "KernelMain: Failed to create SystemInitThread (status 0x%x)\n", st);
    }

    DbgPrint("KernelMain: Handing over to scheduler...\n");
    
    // Ignite the system!
    _enable();

    // The Idle Loop: The CPU rests here whenever all other threads are waiting on events.
    while (true) {
        __halt(); 
    }
}