// PEB-based module and function resolution. No static imports — everything resolved at runtime.
// 1. Loader structs — LIST_ENTRY → PEB_LDR_DATA → LDR_DATA_TABLE_ENTRY (DllBase, name, size).
// 2. PEB (GS:[0x60]) + hash_ror13 — case-insensitive, same algo as Hell's Gate manual-map tooling.
// 3. get_peb() — inline asm to read GS:[0x60], returns PEB pointer.
// 4. get_module_by_hash — walk InMemoryOrderModuleList, match by ror13 hash.
// 5. resolve_forward_string — chase "DLL.func" forwarded exports through chain of modules.
// 6. get_func_by_hash — parse IMAGE_EXPORT_DIRECTORY, handle forwarded exports.
// 7. resolve_api — three-tier: specific module → kernel32 cache → ntdll cache → full PEB walk.
const std = @import("std");
const win = @import("win32.zig");
const pe = @import("pe.zig");

// Doubly-linked list node used by Windows loader structures.
// PEB_LDR_DATA chains these to track loaded modules in three orderings.
pub const LIST_ENTRY = extern struct {
    Flink: *LIST_ENTRY,
    Blink: *LIST_ENTRY,
};

// Loader data block hung off the PEB. Each module list is a circular linked list
// of LDR_DATA_TABLE_ENTRY nodes. InMemoryOrder is the safest to walk — InLoadOrder
// can race during DLL load/unload on other threads.
pub const PEB_LDR_DATA = extern struct {
    Length: win.ULONG,
    Initialized: win.BYTE,
    Padding: [3]win.BYTE,
    SsHandle: win.LPVOID,
    InLoadOrderModuleList: LIST_ENTRY,
    InMemoryOrderModuleList: LIST_ENTRY,
    InInitializationOrderModuleList: LIST_ENTRY,
};

// One entry per loaded module. Contains the DLL base, image size, and its name.
// The InMemoryOrderLinks list_entry is embedded at a known offset from the struct start —
// we subtract that offset to get back to the full LDR_DATA_TABLE_ENTRY.
pub const LDR_DATA_TABLE_ENTRY = extern struct {
    InLoadOrderLinks: LIST_ENTRY,
    InMemoryOrderLinks: LIST_ENTRY,
    InInitializationOrderLinks: LIST_ENTRY,
    DllBase: win.LPVOID,
    EntryPoint: win.LPVOID,
    SizeOfImage: win.ULONG,
    FullDllName: win.UNICODE_STRING,
    BaseDllName: win.UNICODE_STRING,
};

// The Process Environment Block lives at GS:0x60 on x64.
// Only the fields we actually touch are defined — the real struct is much larger.
pub const PEB = extern struct {
    Reserved1: [2]win.BYTE,
    BeingDebugged: win.BYTE,
    Reserved2: [1]win.BYTE,
    Reserved3: [2]win.LPVOID,
    Ldr: *PEB_LDR_DATA,
};

// ror13 hash of a module or function name. Case-insensitive — uppercases each byte before mixing.
// Same algorithm used by Hell's Gate and other manual-map tooling for consistency.
pub fn hash_ror13(input: []const u8) u32 {
    var hash: u32 = 0;
    for (input) |c| {
        var c_upper = c;
        if (c_upper >= 'a' and c_upper <= 'z') c_upper -= 32;
        hash = (hash >> 13) | (hash << 19);
        hash +%= c_upper;
    }
    return hash;
}

// Reads GS:[0x60] to return the PEB pointer.
pub fn get_peb() *PEB {
    const ptr: usize = asm volatile ("mov %%gs:0x60, %[ptr]"
        : [ptr] "=r" (-> usize),
    );
    return @as(*PEB, @ptrFromInt(ptr));
}

// Walks InMemoryOrderModuleList looking for a DLL whose ror13-hashed BaseDllName matches.
// Returns the module's base address or null if not found.
pub fn get_module_by_hash(hash: u32) ?*anyopaque {
    const peb_ptr = get_peb();
    const ldr = peb_ptr.Ldr;
    // Walk the circular linked list starting from InMemoryOrderModuleList
    const head = @as(*LIST_ENTRY, @ptrCast(&ldr.InMemoryOrderModuleList));
    var entry = head.Flink;

    while (entry != head) : (entry = entry.Flink) {
        // Subtract the offset of InMemoryOrderLinks to get back to the LDR_DATA_TABLE_ENTRY
        const le: *LDR_DATA_TABLE_ENTRY = @ptrCast(@alignCast(@as(*anyopaque, @ptrFromInt(@intFromPtr(entry) - @offsetOf(LDR_DATA_TABLE_ENTRY, "InMemoryOrderLinks")))));

        if (le.BaseDllName.Length == 0) continue;

        const buf = le.BaseDllName.Buffer;
        const len = le.BaseDllName.Length / 2;
        var mod_hash: u32 = 0;
        var i: usize = 0;
        while (i < len) : (i += 1) {
            var c = buf[i];
            if (c >= 'a' and c <= 'z') c -= 32;
            mod_hash = (mod_hash >> 13) | (mod_hash << 19);
            mod_hash +%= @as(u32, c);
        }

        if (mod_hash == hash) return le.DllBase;
    }

    return null;
}

// Follows a forwarded export string ("module.function") to resolve the real target.
// Splits on the dot, hashes both halves, looks up the module, then the function within it.
fn resolve_forward_string(base_bytes: [*]u8, forward_rva: u32) ?*anyopaque {
    const forward_str = @as([*:0]u8, @ptrCast(@alignCast(base_bytes + @as(usize, @intCast(forward_rva)))));
    const str = std.mem.sliceTo(forward_str, 0);
    const dot_pos = std.mem.indexOfScalar(u8, str, '.') orelse return null;
    const dll_hash = hash_ror13(str[0..dot_pos]);
    const func_hash = hash_ror13(str[dot_pos + 1 ..]);
    const dll_base = get_module_by_hash(dll_hash) orelse return null;
    return get_func_by_hash(dll_base, func_hash);
}

// Parses a module's export directory, hashes each exported name, and returns the matching
// function address. Handles forwarded exports (looking-glass exports that redirect to another DLL).
pub fn get_func_by_hash(base: *anyopaque, func_hash: u32) ?*anyopaque {
    const base_bytes = @as([*]u8, @ptrCast(base));
    const dos = @as(*pe.IMAGE_DOS_HEADER, @ptrCast(@alignCast(base_bytes)));
    if (dos.e_magic != 0x5A4D) return null;

    const nt = @as(*pe.IMAGE_NT_HEADERS64, @ptrCast(@alignCast(base_bytes + @as(usize, @intCast(dos.e_lfanew)))));
    if (nt.Signature != 0x00004550) return null;

    const export_dir = nt.OptionalHeader.DataDirectory[0];
    if (export_dir.VirtualAddress == 0 or export_dir.Size == 0) return null;

    const exp = @as(*pe.IMAGE_EXPORT_DIRECTORY, @ptrCast(@alignCast(base_bytes + @as(usize, @intCast(export_dir.VirtualAddress)))));

    if (exp.NumberOfNames == 0 or exp.AddressOfFunctions == 0 or exp.AddressOfNames == 0 or exp.AddressOfNameOrdinals == 0) return null;

    const names = @as([*]u32, @ptrCast(@alignCast(base_bytes + @as(usize, @intCast(exp.AddressOfNames)))));
    const funcs = @as([*]u32, @ptrCast(@alignCast(base_bytes + @as(usize, @intCast(exp.AddressOfFunctions)))));
    const ords = @as([*]u16, @ptrCast(@alignCast(base_bytes + @as(usize, @intCast(exp.AddressOfNameOrdinals)))));

    var i: u32 = 0;
    while (i < exp.NumberOfNames) : (i += 1) {
        if (names[i] == 0) continue;

        const name_ptr = @as([*:0]u8, @ptrCast(@alignCast(base_bytes + @as(usize, @intCast(names[i])))));
        const name = std.mem.sliceTo(name_ptr, 0);

        if (hash_ror13(name) == func_hash) {
            const ordinal = ords[i];
            if (ordinal >= exp.NumberOfFunctions) continue;
            const rva = funcs[ordinal];
            if (rva >= export_dir.VirtualAddress and rva < export_dir.VirtualAddress + export_dir.Size) {
                return resolve_forward_string(base_bytes, rva);
            }
            return @as(*anyopaque, @ptrCast(base_bytes + @as(usize, @intCast(rva))));
        }
    }

    return null;
}

var cached_kernel32: ?*anyopaque = null;
var cached_ntdll: ?*anyopaque = null;

// Three-tier lookup: specific module by hash → kernel32 → ntdll → full PEB module walk.
// Falls back to scanning every loaded module if the requested module isn't found or isn't loaded.
pub fn resolve_api(module_hash: u32, func_hash: u32) ?*anyopaque {
    if (module_hash != 0) {
        if (get_module_by_hash(module_hash)) |base| {
            if (get_func_by_hash(base, func_hash)) |f| return f;
        }
    }

    if (cached_kernel32 == null) cached_kernel32 = get_module_by_hash(hash_ror13("kernel32.dll"));
    if (cached_kernel32) |base| {
        if (get_func_by_hash(base, func_hash)) |f| return f;
    }

    if (cached_ntdll == null) cached_ntdll = get_module_by_hash(hash_ror13("ntdll.dll"));
    if (cached_ntdll) |base| {
        if (get_func_by_hash(base, func_hash)) |f| return f;
    }

    const peb_ptr = get_peb();
    const ldr = peb_ptr.Ldr;
    const head = @as(*LIST_ENTRY, @ptrCast(&ldr.InMemoryOrderModuleList));
    var entry = head.Flink;

    while (entry != head) : (entry = entry.Flink) {
        const le: *LDR_DATA_TABLE_ENTRY = @ptrCast(@alignCast(@as(*anyopaque, @ptrFromInt(@intFromPtr(entry) - @offsetOf(LDR_DATA_TABLE_ENTRY, "InMemoryOrderLinks")))));

        if (cached_kernel32) |ck| if (le.DllBase == ck) continue;
        if (cached_ntdll) |cn| if (le.DllBase == cn) continue;

        if (get_func_by_hash(le.DllBase, func_hash)) |f| return f;
    }

    return null;
}