// main.c - kernel entry point
#include "boot_info.h"
#include "mm_phys.h"
#include "mm_virt.h"
#include "ob_object.h"
#include "ke_process.h"
#include "ke_thread.h"
#include "ke_sched.h"
#include "ke_idt.h"
#include "ke_irq.h"
#include "ke_syscall.h"
#include "ke_debug.h"
#include "ldr_image.h"
#include "hal_api.h"
#include "io_file.h"
#include "ini_parser.h"
#include <stdint.h>
#include <string.h>
#include <intrin.h>

// dynamic arrays bitch
typedef struct {
  size_t count;
  size_t capacity;
} Header;

#define ARR_INIT_CAPACITY 2

#define _arr_header(arr) ((Header*)(arr) - 1)

#define _arr_maybe_grow(arr) \
    do { \
        if (!(arr)) { \
            Header* h = MmAllocatePool(sizeof(Header) + sizeof(*(arr)) * ARR_INIT_CAPACITY); \
            h->count = 0; \
            h->capacity = ARR_INIT_CAPACITY; \
            (arr) = (void*)(h + 1); \
        } \
        Header* h = _arr_header(arr); \
        if (h->count >= h->capacity) { \
            h->capacity *= 2; \
            h = MmReallocatePool(h, sizeof(Header) + sizeof(*(arr)) * h->capacity); \
            (arr) = (void*)(h + 1); \
        } \
    } while (0)

#define arr_push(arr, x) \
    do { \
        _arr_maybe_grow(arr); \
        (arr)[_arr_header(arr)->count++] = (x); \
    } while (0)

#define arr_push_front(arr, x) \
    do { \
        _arr_maybe_grow(arr); \
        Header* h = _arr_header(arr); \
        memmove(&(arr)[1], &(arr)[0], sizeof(*(arr)) * h->count); \
        (arr)[0] = (x); \
        h->count++; \
    } while (0)

#define arr_remove(arr, index) \
    do { \
        if ((arr) && (index) < arr_len(arr)) { \
            Header* h = _arr_header(arr); \
            memmove(&(arr)[index], &(arr)[index + 1], sizeof(*(arr)) * (h->count - (index) - 1)); \
            h->count--; \
        } \
    } while (0)

#define arr_len(arr) ((Header*)(arr) - 1)->count

#define arr_free(arr) MmFreePool((Header*)(arr) - 1)

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

// HAL pointer cached at boot so kernel threads can reach it without
// threading a parameter through every call chain.
static const HAL_API* gKmainHal;

// ---------------------------------------------------------------------------
// Phase Helper: Resolve Imports for all registered modules
// ---------------------------------------------------------------------------
static void ResolveSystemImports(BOOT_INFO* bootInfo) {
    // Resolve imports for all loaded modules (drivers may import from hal.dll, win32.dll imports from syscall.dll, etc).
    // Starting at 2 to skip the kernel and HAL base images.
    for (uint32_t i = 2; i < bootInfo->ModuleCount; i++) {
        LdrResolveImports(bootInfo->Modules[i].VirtBase);
    }
}

// ---------------------------------------------------------------------------
// Phase Helper: Initialize dynamically loaded .sys drivers
// ---------------------------------------------------------------------------
static void DriversIniCb(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);
        }
    }
}

static void InitializeDrivers(BOOT_INFO* bootInfo, const HAL_API* hal) {
    DbgPrint("KernelMain: Initializing drivers...\n");
    
    // boot drivers (provided by the bootloader)
    for (uint32_t i = 0; i < bootInfo->ModuleCount; i++) {
        const char* name = bootInfo->Modules[i].Name;
        
        // Identify .sys files by extension
        size_t len = strlen(name);
        if (len < 4 || strcmp(name + len - 4, ".sys")) continue;

        DbgPrint("KernelMain: Executing DriverEntry for %s\n", name);

        DRIVER_ENTRY_FN entry = (DRIVER_ENTRY_FN)(LdrGetEntryPoint(bootInfo->Modules[i].VirtBase));

        KSTATUS ds = entry(hal);
        if (!KSUCCESS(ds)) {
            DbgPrint("KernelMain: Driver %s failed to init!\n", name);
        }
    }

    // runtime drivers
    KSTATUS s;
    FILE_OBJECT* f;

    s = IoOpenFile("drivers\\drivers.ini", &f);
    if (s != STATUS_SUCCESS) {
        DbgPrint("KernelMain: Failed to open drivers.ini\n");
        _disable(); while(1) __halt();
    }

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

    IoCloseFile(f);

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

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

        // Load the driver
        size_t len = strlen(runtimeKeys[i]);
        char* path = MmAllocatePool(len + 10);
        if (!path) {
            DbgPrint("KernelMain: Failed to allocate memory for driver path. How?\n");
            continue;
        }

        // memcpy manouvers
        memcpy(path, "drivers\\", 8);
        memcpy(path + 8, runtimeKeys[i], len + 1);

        s = IoOpenFile(path, &f);
        if (s != STATUS_SUCCESS) {
            DbgPrint("KernelMain: Failed to open runtime driver: %s\n", runtimeKeys[i]);
            continue; // Skip this driver but keep going with others
        }

        MmFreePool(path);

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

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

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

        KSTATUS ds = entry(hal);
        if (!KSUCCESS(ds)) {
            DbgPrint("KernelMain: 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 a DLL from disk and register it
// ---------------------------------------------------------------------------
static void LoadAndRegisterDll(const char* filename) {
    KSTATUS s;
    FILE_OBJECT* f;
    
    // Read the DLL from disk
    s = IoOpenFile(filename, &f);
    if (s != STATUS_SUCCESS) {
        DbgPrint("KernelMain: Failed to open DLL: %s\n", filename);
        return;
    }
    
    uint64_t size = IoGetFileSize(f);
    void* mem = MmAllocatePool(size); // Temporary pool for the raw file
    size_t got;
    if (!KSUCCESS(IoReadFile(f, mem, size, &got))) {
        DbgPrint("KernelMain: Failed to read %s\n", filename);
        IoCloseFile(f);
        return;
    }
    IoCloseFile(f);

    // Map the DLL into memory 
    // Passing 'true' for user-mode
    uintptr_t mappedBase = 0;
    s = LdrLoadImage(mem, size, &mappedBase, true); 
    if (!KSUCCESS(s)) {
        DbgPrint("KernelMain: Failed to map %s\n", filename);
        return;
    }

    MmFreePool(mem); // Free the raw file buffer, we only need the mapped image now

    LdrRegisterModule(filename, mappedBase); // Register the module
}

// ---------------------------------------------------------------------------
// 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* LaunchUserProcess(const char* filePath) {
    KSTATUS s;
    FILE_OBJECT* f;
    
    // Read the executable file into a raw buffer
    s = IoOpenFile(filePath, &f);
    DbgPrint("KernelMain: IoOpenFile: %d\n", s);
    if (s != STATUS_SUCCESS) {
        DbgPrint("Failed to load file: \"%s\"!\n", filePath);
        return (PROCESS*)0;
    }
    
    uint64_t size = IoGetFileSize(f);
    DbgPrint("KernelMain: IoGetFileSize: %llu\n", size);

    void* mem = MmAllocateUserMemory(size);
    DbgPrint("KernelMain: MmAllocateUserMemory: 0x%llX\n", mem);

    size_t got;
    s = IoReadFile(f, mem, size, &got);
    DbgPrint("KernelMain: IoReadFile: %d\n", s);
    IoCloseFile(f);
    DbgPrint("KernelMain: IoCloseFile\n");

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

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

    // Parse the raw PE file and map it into executable memory
    uintptr_t imageBase = 0;
    s = LdrLoadImage(mem, size, &imageBase, true);
    if (!KSUCCESS(s)) {
        DbgPrint("KernelMain: LdrLoadImage failed with status: %d\n", s);
    }
    uintptr_t entryPoint = LdrGetEntryPoint(imageBase);
    DbgPrint("KernelMain: PE Image mapped! Base 0x%llx, Entry 0x%llX\n", imageBase, entryPoint);

    // Create the main thread, pointing it to the PE Entry Point
    s = KeCreateThread(proc, (uintptr_t)entryPoint, 0, PriorityNormal, &thread);
    if (!KSUCCESS(s)) {
        DbgPrint("KernelMain: Failed to create thread. Status: %d\n", s);
    }
    DbgPrint("KernelMain: Thread created and queued!\n");

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

// ---------------------------------------------------------------------------
// Kernel thread: wait for a process to exit, then log a message and shut down.
// Takes ownership of the PROCESS* reference passed as arg; releases it on exit.
// ---------------------------------------------------------------------------
static void ProcessWatcherThread(void* arg) {
    PROCESS* proc = (PROCESS*)arg;
    KeWaitForEvent(proc->ExitEvent);
    DbgPrint("KernelMain: Process \"%s\" (PID %u) exited with code %d.\n",
             proc->Name, proc->ProcessId, proc->ExitCode);
    ObDereferenceObject(&proc->Header);

    DbgPrint("KernelMain: Shutting down.\n");
    gKmainHal->Power.Shutdown();  // does not return
}

static void PrintProcessList() {
    LIST_ENTRY* listHead = KeGetProcessList();
    LIST_ENTRY* current;

    DbgPrint("\nPID   | Name             | Threads | ExitCode\n");
    DbgPrint("---------------------------------------------\n");

    for (current = listHead->Fwd; current != listHead; current = current->Fwd) {
        // Translate the list entry back to the PROCESS structure
        PROCESS* proc = CONTAINING_RECORD(current, PROCESS, GlobalLink);

        DbgPrint("%-5u | %-16s | %-7u | %d\n", 
            proc->ProcessId, 
            proc->Name, 
            proc->ThreadCount, 
            proc->ExitCode);
    }
}

// ---------------------------------------------------------------------------
// Main Kernel Entry
// ---------------------------------------------------------------------------
extern void KernelMain(BOOT_INFO* bootInfo) {
    //
    // Physical & Virtual Memory initialization
    //
    MmInitializePhysical(&bootInfo->MemoryMap, bootInfo->PhysicalMapBase);
    PHYSICAL_ADDRESS pml4 = (PHYSICAL_ADDRESS)__readcr3();
    MmInitializeVirtual(pml4, 4 * 1024 * 1024 /* 4 MB initial pool */);

    //
    // CPU Architecture setup (GDT/IDT/Faults)
    //
    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,
        .DebugEnabled       = bootInfo->DebugEnabled
    };
    KSTATUS s = HalInitialize(&halParams);
    if (!KSUCCESS(s)) for (;;) __halt(); // there's nothing we can do.
    const HAL_API* hal = HalGetApi();
    gKmainHal = hal;

    KeInitDebug(hal->Screen.PrintChar, bootInfo->DebugEnabled); // Enable debug

    KeIrqInitialize(hal);

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

    HalRegisterScheduler(KeSchedulerTick);

    KeSyscallSetHal(hal);

    //
    // Dynamic Linking & Driver Execution
    //
    ResolveSystemImports(bootInfo);
    InitializeDrivers(bootInfo, hal);

    //
    // Scheduler Prep & Status Update
    //
    KeInitializeScheduler();
    hal->Screen.Clear(0x000000);

    DbgPrint("KernelMain: Loading system DLLs...\n");
    LoadAndRegisterDll("syscall.dll");
    LoadAndRegisterDll("win32.dll");
    LoadAndRegisterDll("bink2w64.dll");

    // Now that they are in memory and registered, we link them together
    // (e.g., win32.dll resolving its imports against syscall.dll)
    ResolveSystemImports(bootInfo);

    //
    // User-Mode Launch
    //
    PROCESS* shellProc = LaunchUserProcess(bootInfo->CommandLine);

    // Spawn a kernel thread that blocks on the process exit event and prints
    // when it fires.  The watcher takes ownership of the LaunchUserProcess
    // reference so we do not hold a dangling pointer in KernelMain.
    if (shellProc) {
        KSTATUS ws = KeCreateKernelThread((uintptr_t)ProcessWatcherThread,
                                          (uintptr_t)shellProc,
                                          PriorityNormal, (THREAD**)0);
        if (!KSUCCESS(ws)) {
            DbgPrint("KernelMain: Failed to create watcher thread (%d) -- leaking proc ref.\n", ws);
        }
        // Ownership of shellProc transferred to watcher; do not use it below.
    }

    //
    // Ignite the System
    //
    DbgPrint("KernelMain: Starting scheduler... Handing over control.\n");
    _enable();

    DbgPrint("Memory Status:\n");
    DbgPrint("  -> Free (PMM):   %.2f MB\n", (float)MmFreeFrameCount() * PAGE_SIZE / (1024.0f * 1024.0f));
    DbgPrint("  -> Kernel Pool:  %.2f MB\n", (float)MmGetKernelPoolUsage() / (1024.0f * 1024.0f));
    DbgPrint("  -> User Heap:    %.2f MB\n", (float)MmGetUserHeapUsage() / (1024.0f * 1024.0f));

    while (1) {
        KeYield(); // nothing to do
    }
}