// ldr_image.c - PE32+ loader and export resolver
#include <stdint.h>
#include <string.h>
#include <ldr_image.h>
#include <boot_info.h>
#include <ke_process.h>
#include <io_file.h>
#include <mm_virt.h>
#include <ke_debug.h>

int _stricmp(const char* a, const char* b);

// ---------------------------------------------------------------------------
// PE32+ structures (just what we need)
// ---------------------------------------------------------------------------
#define PE_SIGNATURE        0x00004550UL   // 'PE\0\0'
#define IMAGE_FILE_MACHINE_AMD64  0x8664
#define IMAGE_FILE_MACHINE_I386   0x014c
#define IMAGE_NT_OPTIONAL_HDR32_MAGIC 0x10b
#define IMAGE_NT_OPTIONAL_HDR64_MAGIC 0x20b

#pragma pack(push, 1)
typedef struct {
    uint16_t Magic;
    uint8_t  MajorLinkerVersion;
    uint8_t  MinorLinkerVersion;
    uint32_t SizeOfCode;
    uint32_t SizeOfInitializedData;
    uint32_t SizeOfUninitializedData;
    uint32_t AddressOfEntryPoint;
    uint32_t BaseOfCode;
    uint64_t ImageBase;
    uint32_t SectionAlignment;
    uint32_t FileAlignment;
    uint16_t MajorOSVersion;
    uint16_t MinorOSVersion;
    uint16_t MajorImageVersion;
    uint16_t MinorImageVersion;
    uint16_t MajorSubsystemVersion;
    uint16_t MinorSubsystemVersion;
    uint32_t Win32VersionValue;
    uint32_t SizeOfImage;
    uint32_t SizeOfHeaders;
    uint32_t CheckSum;
    uint16_t Subsystem;
    uint16_t DllCharacteristics;
    uint64_t SizeOfStackReserve;
    uint64_t SizeOfStackCommit;
    uint64_t SizeOfHeapReserve;
    uint64_t SizeOfHeapCommit;
    uint32_t LoaderFlags;
    uint32_t NumberOfRvaAndSizes;
    struct { uint32_t VirtualAddress; uint32_t Size; } DataDirectory[16];
} OPT_HEADER64;

typedef struct {
    uint32_t Signature;
    uint16_t Machine;
    uint16_t NumberOfSections;
    uint32_t TimeDateStamp;
    uint32_t PointerToSymbolTable;
    uint32_t NumberOfSymbols;
    uint16_t SizeOfOptionalHeader;
    uint16_t Characteristics;
    OPT_HEADER64 OptionalHeader;
} PE_HEADER;

typedef struct {
    char     Name[8];
    uint32_t VirtualSize;
    uint32_t VirtualAddress;
    uint32_t SizeOfRawData;
    uint32_t PointerToRawData;
    uint32_t PointerToRelocations;
    uint32_t PointerToLinenumbers;
    uint16_t NumberOfRelocations;
    uint16_t NumberOfLinenumbers;
    uint32_t Characteristics;
} SECTION_HEADER;

typedef struct {
    uint32_t Characteristics;
    uint32_t TimeDateStamp;
    uint16_t MajorVersion;
    uint16_t MinorVersion;
    uint32_t Name;
    uint32_t Base;
    uint32_t NumberOfFunctions;
    uint32_t NumberOfNames;
    uint32_t AddressOfFunctions;
    uint32_t AddressOfNames;
    uint32_t AddressOfNameOrdinals;
} EXPORT_DIR;

typedef struct {
    uint32_t OriginalFirstThunk;
    uint32_t TimeDateStamp;
    uint32_t ForwarderChain;
    uint32_t Name;
    uint32_t FirstThunk;
} IMPORT_DESC;

typedef struct {
    uint32_t VirtualAddress;
    uint32_t SizeOfBlock;
} BASE_RELOC;
#pragma pack(pop)

#define IMAGE_REL_BASED_DIR64 10

// ---------------------------------------------------------------------------
// Module registry
// ---------------------------------------------------------------------------
#define MAX_MODULES  32

typedef struct {
    char      Name[32];
    uintptr_t Base;
    size_t    Size;
} MODULE_ENTRY;

static MODULE_ENTRY gModules[MAX_MODULES];
static uint32_t     gModuleCount = 0;


static const char* KiCanonicalModuleName(const char* name) {
    if (!name) return name;

    if (!_stricmp(name, "kernel32.dll") ||
        !_stricmp(name, "user32.dll")   ||
        !_stricmp(name, "gdi32.dll")    ||
        !_stricmp(name, "ole32.dll")    ||
        !_stricmp(name, "winmm.dll")) {
        return "win32.dll";
    }

    if (!_stricmp(name, "msvcrt.dll")) {
        return "posix.dll";
    }

    return name;
}

static uintptr_t KiFindModuleBaseInTable(const MODULE_ENTRY* modules, uint32_t count, const char* name) {
    for (uint32_t i = 0; i < count; i++) {
        if (_stricmp(modules[i].Name, name) == 0) return modules[i].Base;
    }
    return 0;
}

static void KiRegisterModuleEntry(MODULE_ENTRY* entry, const char* name, uintptr_t base, size_t size) {
    strncpy(entry->Name, name, 31);
    entry->Name[31] = '\0';
    entry->Base = base;
    entry->Size = size;
}

void LdrRegisterModulesFromBootInfo(const BOOT_INFO* bootInfo) {
    for (uint32_t i = 0; i < bootInfo->ModuleCount && gModuleCount < MAX_MODULES; i++) {
        // Only register modules that have a virtual base (Kernel, HAL)
        if (bootInfo->Modules[i].VirtBase != 0) {
            KiRegisterModuleEntry(&gModules[gModuleCount], 
                                bootInfo->Modules[i].Name,
                                bootInfo->Modules[i].VirtBase, 
                                bootInfo->Modules[i].ImageSize);
            gModuleCount++;
        }
    }
}

void LdrRegisterModule(const char* name, uintptr_t base) {
    if (gModuleCount < MAX_MODULES) {
        KiRegisterModuleEntry(&gModules[gModuleCount], name, base, 0);
        gModuleCount++;
    }
}

static uintptr_t KiFindGlobalModuleBase(const char* name) {
    return KiFindModuleBaseInTable(gModules, gModuleCount, KiCanonicalModuleName(name));
}

static uintptr_t KiFindProcessModuleBaseInTable(
    const PROCESS_MODULE* modules,
    uint32_t count,
    const char* name
) {
    for (uint32_t i = 0; i < count; i++) {
        if (_stricmp(modules[i].Name, KiCanonicalModuleName(name)) == 0) {
            return modules[i].Base;
        }
    }
    return 0;
}

static uintptr_t KiFindProcessOrGlobalModuleBase(PROCESS* process, const char* name) {
    if (process) {
        uintptr_t base = KiFindProcessModuleBaseInTable(process->Modules, process->ModuleCount, name);
        if (base) return base;
    }
    return KiFindGlobalModuleBase(name);
}

static void* KiGetExportAtBase(uintptr_t base, const char* functionName) {
    if (!base) return (void *)0;

    uint32_t peOffset = *(uint32_t *)(base + 0x3C);
    PE_HEADER* pe = (PE_HEADER *)(base + peOffset);
    if (pe->Signature != PE_SIGNATURE) return (void *)0;

    OPT_HEADER64* opt = &pe->OptionalHeader;
    if (opt->NumberOfRvaAndSizes < 1) return (void *)0;

    uint32_t expRva = opt->DataDirectory[0].VirtualAddress;
    if (!expRva) return (void *)0;

    EXPORT_DIR* exp = (EXPORT_DIR *)(base + expRva);
    uint32_t* names     = (uint32_t *)(base + exp->AddressOfNames);
    uint16_t* ordinals  = (uint16_t *)(base + exp->AddressOfNameOrdinals);
    uint32_t* functions = (uint32_t *)(base + exp->AddressOfFunctions);

    for (uint32_t i = 0; i < exp->NumberOfNames; i++) {
        const char* name = (const char *)(base + names[i]);
        if (strcmp(name, functionName) == 0) {
            uint16_t ord = ordinals[i];
            return (void *)(base + functions[ord]);
        }
    }
    return (void *)0;
}

// ---------------------------------------------------------------------------
// Export lookup
// ---------------------------------------------------------------------------
void* LdrGetExport(const char* moduleName, const char* functionName) {
    uintptr_t base = KiFindGlobalModuleBase(KiCanonicalModuleName(moduleName));
    return KiGetExportAtBase(base, functionName);
}

void* LdrGetExportForProcess(PROCESS* process, const char* moduleName, const char* functionName) {
    uintptr_t base = KiFindProcessOrGlobalModuleBase(process, KiCanonicalModuleName(moduleName));
    return KiGetExportAtBase(base, functionName);
}

size_t LdrGetImageSize(uintptr_t base) {
    uint32_t peOffset = *(uint32_t *)(base + 0x3C);
    PE_HEADER* pe = (PE_HEADER *)(base + peOffset);
    if (pe->Signature != PE_SIGNATURE) return 0;
    return (size_t)pe->OptionalHeader.SizeOfImage;
}

// ---------------------------------------------------------------------------
// Entry point lookup
// ---------------------------------------------------------------------------
uintptr_t LdrGetEntryPoint(uintptr_t base) {
    uint32_t pe_offset = *(uint32_t*)(base + 0x3C);
    uint32_t entry_rva = *(uint32_t*)(base + pe_offset + 0x28);
    return base + entry_rva;
}

// ---------------------------------------------------------------------------
// Base Relocation
// ---------------------------------------------------------------------------
static KSTATUS LdrApplyRelocations(uintptr_t base, uintptr_t delta) {
    uint32_t peOffset = *(uint32_t *)(base + 0x3C);
    PE_HEADER* pe = (PE_HEADER *)(base + peOffset);
    OPT_HEADER64* opt = &pe->OptionalHeader;

    if (opt->NumberOfRvaAndSizes < 6) return STATUS_SUCCESS;
    uint32_t relocRva = opt->DataDirectory[5].VirtualAddress;
    uint32_t relocSize = opt->DataDirectory[5].Size;
    if (!relocRva || !relocSize) return STATUS_SUCCESS;

    BASE_RELOC* block = (BASE_RELOC *)(base + relocRva);
    uintptr_t end = (uintptr_t)block + relocSize;

    while ((uintptr_t)block < end && block->SizeOfBlock > 0) {
        uint32_t count = (block->SizeOfBlock - sizeof(BASE_RELOC)) / sizeof(uint16_t);
        uint16_t* entries = (uint16_t *)(block + 1);

        for (uint32_t i = 0; i < count; i++) {
            int type = entries[i] >> 12;
            int offset = entries[i] & 0xFFF;

            // For 64-bit PEs, we mostly care about DIR64 (Type 10)
            if (type == IMAGE_REL_BASED_DIR64) {
                uint64_t* ptr = (uint64_t *)(base + block->VirtualAddress + offset);
                *ptr += delta;
            }
        }
        block = (BASE_RELOC *)((uint8_t *)block + block->SizeOfBlock);
    }
    return STATUS_SUCCESS;
}

// ---------------------------------------------------------------------------
// Import resolution
// ---------------------------------------------------------------------------
typedef void* (*EXPORT_LOOKUP_FN)(void* ctx, const char* moduleName, const char* functionName);

static void* KiGlobalExportLookup(void* ctx, const char* moduleName, const char* functionName) {
    (void)ctx;
    return LdrGetExport(moduleName, functionName);
}

static void* KiProcessExportLookup(void* ctx, const char* moduleName, const char* functionName) {
    return LdrGetExportForProcess((PROCESS*)ctx, moduleName, functionName);
}

static KSTATUS LdrResolveImportsCommon(
    PROCESS* process,
    uintptr_t imageBase,
    EXPORT_LOOKUP_FN lookup,
    void* ctx
) {
    uint32_t peOffset = *(uint32_t *)(imageBase + 0x3C);
    PE_HEADER* pe = (PE_HEADER *)(imageBase + peOffset);
    OPT_HEADER64* opt = &pe->OptionalHeader;

    if (opt->NumberOfRvaAndSizes < 2)
        return STATUS_SUCCESS;

    uint32_t impRva = opt->DataDirectory[1].VirtualAddress;
    if (!impRva)
        return STATUS_SUCCESS;

    IMPORT_DESC* imp = (IMPORT_DESC *)(imageBase + impRva);

    while (imp->Name) {
        const char* dllName = (const char *)(imageBase + imp->Name);

        //Compatibility remap.
        dllName = KiCanonicalModuleName(dllName);

        //
        // Ensure dependency is loaded into the process
        //
        if (process) {
            uintptr_t modBase =
                KiFindProcessOrGlobalModuleBase(process, dllName);

            if (!modBase) {
                KSTATUS s =
                    LdrLoadModuleForProcess(process, dllName, &modBase);

                if (!KSUCCESS(s)) {
                    DbgPrint(
                        "LdrResolveImports: Failed to load dependency %s (%d)\n",
                        dllName,
                        s
                    );
                    return s;
                }
            }
        }

        uint64_t* thunk =
            (uint64_t *)(imageBase + imp->FirstThunk);

        uint64_t* origThunk =
            (uint64_t *)(imageBase +
                (imp->OriginalFirstThunk
                    ? imp->OriginalFirstThunk
                    : imp->FirstThunk));

        while (*origThunk) {
            if (*origThunk & (1ULL << 63)) {
                DbgPrint(
                    "LdrResolveImports: Ordinal imports not supported! (%s)\n",
                    dllName
                );
                return STATUS_NOT_SUPPORTED;
            }

            const char* funcName =
                (const char *)(imageBase +
                    (*origThunk & 0xFFFFFFFF) + 2);

            void* addr = lookup(ctx, dllName, funcName);

            if (!addr) {
                DbgPrint(
                    "LdrResolveImports: Function not found! (%s!%s)\n",
                    dllName,
                    funcName
                );
                return STATUS_NOT_FOUND;
            }

            *thunk = (uint64_t)addr;

            thunk++;
            origThunk++;
        }

        imp++;
    }

    return STATUS_SUCCESS;
}

KSTATUS LdrResolveImports(uintptr_t imageBase) {
    return LdrResolveImportsCommon(
        (PROCESS*)0,
        imageBase,
        KiGlobalExportLookup,
        (void*)0
    );
}

KSTATUS LdrResolveImportsForProcess(PROCESS* process, uintptr_t imageBase) {
    return LdrResolveImportsCommon(
        process,
        imageBase,
        KiProcessExportLookup,
        process
    );
}

static void KiDebugLibtccAllocatorSlot(const char* phase, const char* filename, uintptr_t imageBase) {
    if (!filename || !imageBase) return;
    if (_stricmp(filename, "libtcc.dll") != 0) return;

    uint64_t slot = *(uint64_t*)(imageBase + 0x4A010);
    uint64_t expectedDefault = (uint64_t)(imageBase + 0x25F0);

    DbgPrint("LDR libtcc %s: base=%p slot[+4A010]=%p expected-default=%p\n",
             phase,
             (void*)imageBase,
             (void*)(uintptr_t)slot,
             (void*)(uintptr_t)expectedDefault);

    DbgPrint("LDR libtcc data[+4A000..+4A028]: %016llx %016llx %016llx %016llx %016llx %016llx\n",
             (unsigned long long)*(uint64_t*)(imageBase + 0x4A000),
             (unsigned long long)*(uint64_t*)(imageBase + 0x4A008),
             (unsigned long long)*(uint64_t*)(imageBase + 0x4A010),
             (unsigned long long)*(uint64_t*)(imageBase + 0x4A018),
             (unsigned long long)*(uint64_t*)(imageBase + 0x4A020),
             (unsigned long long)*(uint64_t*)(imageBase + 0x4A028));
}

static KSTATUS LdrRegisterProcessModule(PROCESS* process, const char* name, uintptr_t base) {
    if (!process || !name || !base) return STATUS_INVALID_PARAMETER;

    for (uint32_t i = 0; i < process->ModuleCount; i++) {
        if (_stricmp(process->Modules[i].Name, name) == 0) {
            return STATUS_SUCCESS;
        }
    }

    if (process->ModuleCount >= PROCESS_MODULE_MAX) return STATUS_NO_MEMORY;

    PROCESS_MODULE* mod = &process->Modules[process->ModuleCount++];
    strncpy(mod->Name, name, 31);
    mod->Name[31] = '\0';
    mod->Base = base;
    mod->Size = LdrGetImageSize(base);
    return STATUS_SUCCESS;
}

KSTATUS LdrLoadModuleForProcess(PROCESS* process, const char* filename, uintptr_t* baseOut) {
    if (!process || !filename) return STATUS_INVALID_PARAMETER;

    filename = KiCanonicalModuleName(filename);

    for (uint32_t i = 0; i < process->ModuleCount; i++) {
        if (_stricmp(process->Modules[i].Name, filename) == 0) {
            if (baseOut) *baseOut = process->Modules[i].Base;
            return STATUS_SUCCESS;
        }
    }

    FILE_OBJECT* f;
    KSTATUS s = IoOpenFile(filename, &f);
    if (!KSUCCESS(s)) {
        char path[260];
        memcpy(path, "0:/system/", 10);
        memcpy(path + 10, filename, strnlen(filename, 248) + 1);
        path[259] = '\0';
        s = IoOpenFile(path, &f);
        if (!KSUCCESS(s)) {
            DbgPrint("LdrLoadModuleForProcess: Failed to open %s (%d)\n", filename, s);
            return s;
        }
    }

    uint64_t fileSize = IoGetFileSize(f);
    void* mem = MmAllocatePool((size_t)fileSize);
    if (!mem) {
        IoCloseFile(f);
        return STATUS_NO_MEMORY;
    }

    size_t got = 0;
    s = IoReadFile(f, mem, (size_t)fileSize, &got);
    IoCloseFile(f);
    if (!KSUCCESS(s)) {
        MmFreePool(mem);
        return s;
    }

    uintptr_t imageBase = 0;
    s = LdrLoadImage(mem, (size_t)fileSize, &imageBase, true);
    MmFreePool(mem);
    if (!KSUCCESS(s)) {
        return s;
    }

    KiDebugLibtccAllocatorSlot("after load/reloc before imports", filename, imageBase);

    s = LdrResolveImportsForProcess(process, imageBase);
    if (!KSUCCESS(s)) {
        MmFreeUserMemory((void*)imageBase);
        return s;
    }

    KiDebugLibtccAllocatorSlot("after imports", filename, imageBase);

    s = LdrRegisterProcessModule(process, filename, imageBase);
    if (!KSUCCESS(s)) {
        MmFreeUserMemory((void*)imageBase);
        return s;
    }

    if (baseOut) *baseOut = imageBase;
    return STATUS_SUCCESS;
}

// LdrLoadImage - load a PE from a raw buffer into virtual space.
KSTATUS LdrLoadImage(const void* buffer, size_t size, uintptr_t* baseOut, bool user) {
    if (!buffer || !size || !baseOut) return STATUS_INVALID_PARAMETER;
    if (size < 0x40) return STATUS_FAILURE;

    const uint8_t* bytes = (const uint8_t*)buffer;
    uint32_t peOffset = *(const uint32_t *)(bytes + 0x3C);
    if ((size_t)peOffset > size || size - (size_t)peOffset < 24) return STATUS_FAILURE;

    uint32_t signature = *(const uint32_t *)(bytes + peOffset);
    if (signature != PE_SIGNATURE) return STATUS_FAILURE;

    uint16_t machine = *(const uint16_t *)(bytes + peOffset + 4);
    uint16_t optionalSize = *(const uint16_t *)(bytes + peOffset + 20);
    if (optionalSize < sizeof(OPT_HEADER64) ||
        (size_t)peOffset + 24 + optionalSize > size) {
        return STATUS_FAILURE;
    }

    uint16_t optionalMagic = *(const uint16_t *)(bytes + peOffset + 24);
    if (machine == IMAGE_FILE_MACHINE_I386 || optionalMagic == IMAGE_NT_OPTIONAL_HDR32_MAGIC) {
        DbgPrint("LdrLoadImage: hard error: 32-bit PE images are not supported (machine=0x%X magic=0x%X)\n",
                 machine, optionalMagic);
        return STATUS_NOT_SUPPORTED;
    }

    if (machine != IMAGE_FILE_MACHINE_AMD64 || optionalMagic != IMAGE_NT_OPTIONAL_HDR64_MAGIC) {
        DbgPrint("LdrLoadImage: unsupported PE image (machine=0x%X magic=0x%X)\n",
                 machine, optionalMagic);
        return STATUS_NOT_SUPPORTED;
    }

    PE_HEADER* pe = (PE_HEADER *)(bytes + peOffset);
    OPT_HEADER64* opt = &pe->OptionalHeader;
    size_t imageSize = opt->SizeOfImage;

    void* mem = user ? MmAllocateUserMemory(imageSize) : MmAllocatePool(imageSize);
    if (!mem) return STATUS_NO_MEMORY;
    memset(mem, 0, imageSize);

    memcpy(mem, buffer, opt->SizeOfHeaders);
    SECTION_HEADER* sec = (SECTION_HEADER *)((uint8_t *)&pe->OptionalHeader + pe->SizeOfOptionalHeader);
    for (int i = 0; i < pe->NumberOfSections; i++) {
        void* dst = (uint8_t *)mem + sec[i].VirtualAddress;
        const void* src = (const uint8_t *)buffer + sec[i].PointerToRawData;
        if (sec[i].SizeOfRawData) {
            memcpy(dst, src, sec[i].SizeOfRawData < sec[i].VirtualSize ? sec[i].SizeOfRawData : sec[i].VirtualSize);
        }

        uint64_t prot = PAGE_PRESENT | PAGE_WRITE | (user ? PAGE_USER : 0);
        MmSetPageProtect((uintptr_t)mem + sec[i].VirtualAddress, sec[i].VirtualSize, prot);
    }

    uintptr_t delta = (uintptr_t)mem - opt->ImageBase;
    if (delta != 0) {
        LdrApplyRelocations((uintptr_t)mem, delta);
    }

    *baseOut = (uintptr_t)mem;
    return STATUS_SUCCESS;
}