// Synthetic callstack frames through kernel32!BaseThreadInitThunk and
// ntdll!RtlUserThreadStart. Parses .pdata unwind codes to compute per-function
// frame sizes, asm trampoline pushes frames with correct offsets.
const std = @import("std");
const resolve = @import("resolve.zig");
const syscall = @import("syscall.zig");

extern var qFakeRtlUserThreadStart: usize;
extern var qFakeBaseThreadInitThunk: usize;
extern var qRtlUserFrameSize: usize;
extern var qBaseThreadFrameSize: usize;

const UWOP_PUSH_NONVOL: u8 = 0;
const UWOP_ALLOC_LARGE: u8 = 1;
const UWOP_ALLOC_SMALL: u8 = 2;
const UWOP_SET_FPREG: u8 = 3;
const UWOP_SAVE_NONVOL: u8 = 4;
const UWOP_SAVE_NONVOL_FAR: u8 = 5;
const UNW_FLAG_CHAININFO: u8 = 0x04;
const RBP_OP_INFO: u8 = 5;

fn calc_frame_size(mod_base: [*]const u8, unwind_rva: u32, out_size: *usize) bool {
    const info = @as(*align(1) const extern struct {
        VersionAndFlags: u8,
        SizeOfProlog: u8,
        CountOfCodes: u8,
        FrameRegister: u8,
        FrameOffset: u8,
        UnwindCode: [1]u16,
    }, @ptrCast(mod_base + unwind_rva));

    var idx: usize = 0;
    var sz: usize = 0;
    var rbp_pushed: bool = false;

    while (idx < info.CountOfCodes) : (idx += 1) {
        const code = info.UnwindCode[idx];
        const op = @as(u8, @truncate(code & 0xF));
        const info_field = @as(u8, @truncate((code >> 4) & 0xF));

        switch (op) {
            UWOP_PUSH_NONVOL => {
                if (info_field == 2 and !rbp_pushed) return false;
                if (info_field == RBP_OP_INFO) rbp_pushed = true;
                sz += 8;
            },
            UWOP_ALLOC_SMALL => sz += 8 * (@as(usize, info_field) + 1),
            UWOP_ALLOC_LARGE => {
                idx += 1;
                var frame_offset: usize = info.UnwindCode[idx];
                if (info_field == 0) {
                    frame_offset *= 8;
                } else {
                    idx += 1;
                    frame_offset += @as(usize, info.UnwindCode[idx]) << 16;
                }
                sz += frame_offset;
            },
            UWOP_SET_FPREG => {
                rbp_pushed = true;
                const fp_off: isize = -0x10 * @as(isize, @intCast(info.FrameOffset));
                sz = @as(usize, @intCast(@as(isize, @intCast(sz)) + fp_off));
            },
            UWOP_SAVE_NONVOL => {
                if (info_field == RBP_OP_INFO or info_field == 2) return false;
                idx += 1;
            },
            UWOP_SAVE_NONVOL_FAR => {
                if (info_field == RBP_OP_INFO or info_field == 2) return false;
                idx += 2;
            },
            else => {},
        }
    }

    // Chained unwind info
    if ((info.VersionAndFlags & UNW_FLAG_CHAININFO) != 0) {
        idx = info.CountOfCodes;
        if ((idx & 1) != 0) idx += 1;
        const chain_rva = @as(*align(1) const u32, @ptrCast(&info.UnwindCode[idx])).*;
        var chain_sz: usize = 0;
        _ = calc_frame_size(mod_base, chain_rva, &chain_sz);
        sz += chain_sz;
    }

    out_size.* = sz + 8;
    return true;
}

// Binary search .pdata for unwind entry covering the given RVA.
fn find_runtime_entry(target_rva: u32) ?struct { begin: u32, unwind: u32 } {
    const exc_begin = syscall.g_exc_begin;
    const exc_count = syscall.g_exc_count;
    if (exc_begin == 0 or exc_count == 0) return null;

    const funcs = @as([*]align(1) const syscall.RUNTIME_FUNCTION, @ptrFromInt(exc_begin));
    var lo: usize = 0;
    var hi: usize = exc_count;
    while (lo < hi) {
        const mid = lo + (hi - lo) / 2;
        const entry = funcs[mid];
        if (target_rva < entry.BeginAddress) {
            hi = mid;
        } else if (target_rva >= entry.EndAddress) {
            lo = mid + 1;
        } else {
            return .{ .begin = entry.BeginAddress, .unwind = entry.UnwindInfoAddress };
        }
    }
    return null;
}

fn export_rva(mod: [*]const u8, name: []const u8) ?u32 {
    const dos = @as(*align(1) extern struct { e_magic: u16, pad: [58]u8, e_lfanew: u32 }, @ptrCast(@constCast(mod)));
    if (dos.e_magic != 0x5A4D) return null;

    const nt = @as(*align(1) extern struct {
        Signature: u32,
        FileHeader: extern struct { Machine: u16, NumberOfSections: u16, pad: [16]u8 },
        OptionalHeader: extern struct { Magic: u16, pad: [110]u8, DataDirectory: [16]extern struct { VirtualAddress: u32, Size: u32 } },
    }, @ptrCast(@constCast(mod + @as(usize, @intCast(dos.e_lfanew)))));
    if (nt.Signature != 0x00004550) return null;

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

    const exp = @as(*align(1) extern struct {
        Characteristics: u32,
        TimeDateStamp: u32,
        MajorVersion: u16,
        MinorVersion: u16,
        Name: u32,
        Base: u32,
        NumberOfFunctions: u32,
        NumberOfNames: u32,
        AddressOfFunctions: u32,
        AddressOfNames: u32,
        AddressOfNameOrdinals: u32,
    }, @ptrCast(@constCast(mod + @as(usize, @intCast(exp_dir.VirtualAddress)))));

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

    var i: u32 = 0;
    while (i < exp.NumberOfNames) : (i += 1) {
        const n = @as([*:0]const u8, @ptrCast(@alignCast(@constCast(mod + @as(usize, @intCast(names[i]))))));
        if (std.mem.eql(u8, std.mem.sliceTo(n, 0), name)) {
            if (ords[i] < exp.NumberOfFunctions) {
                return funcs[ords[i]];
            }
        }
    }
    return null;
}

pub fn init_frames() void {
    if (qFakeRtlUserThreadStart != 0) return;

    const ntdll = resolve.get_module_by_hash(resolve.hash_ror13("ntdll.dll")) orelse return;
    const k32 = resolve.get_module_by_hash(resolve.hash_ror13("kernel32.dll")) orelse return;
    const ntdll_bytes = @as([*]const u8, @ptrCast(ntdll));
    const k32_bytes = @as([*]const u8, @ptrCast(k32));

    const btit_rva = export_rva(k32_bytes, "BaseThreadInitThunk") orelse return;
    if (find_runtime_entry(btit_rva)) |entry| {
        var fs: usize = 0;
        if (calc_frame_size(k32_bytes, entry.unwind, &fs)) {
            qFakeBaseThreadInitThunk = @intFromPtr(k32) + entry.begin;
            qBaseThreadFrameSize = fs;
        }
    }

    const rust_rva = export_rva(ntdll_bytes, "RtlUserThreadStart") orelse return;
    if (find_runtime_entry(rust_rva)) |entry| {
        var fs: usize = 0;
        if (calc_frame_size(ntdll_bytes, entry.unwind, &fs)) {
            qFakeRtlUserThreadStart = @intFromPtr(ntdll) + entry.begin;
            qRtlUserFrameSize = fs;
        }
    }
}