const std = @import("std");
const builtin = @import("builtin");

const Hook = @import("Hook.zig");
const mem = @import("mem.zig");
const utils = @import("utils.zig");

const HookManager = @This();

hooks: std.ArrayList(Hook),
exec_page: []u8,

pub fn init(alloc: std.mem.Allocator) !HookManager {
    // create our exectuable page to store trampolines
    const page_size = std.heap.page_size_min;
    const exec_page: []u8 = try alloc.alignedAlloc(u8, .fromByteUnits(page_size), page_size);
    if (builtin.os.tag == .windows) {
        var oldprotect: u32 = undefined;
        try std.os.windows.VirtualProtect(exec_page.ptr, page_size, std.os.windows.PAGE_EXECUTE_READWRITE, &oldprotect);
    } else if (builtin.os.tag == .linux) {
        const PROT = std.os.linux.PROT;
        try std.os.linux.mprotect(exec_page.ptr, page_size, PROT.EXEC | PROT.WRITE | PROT.READ);
    } else {
        @compileError("unsupported os");
    }
    return HookManager{
        .hooks = std.ArrayList(Hook).init(alloc),
        .exec_page = exec_page[0..],
    };
}

pub fn deinit(self: *HookManager) usize {
    var count: usize = 0;
    for (self.hooks.items) |*hook| {
        hook.unhook() catch continue;
        count += 1;
    }

    self.hooks.deinit();
    return count;
}

pub fn findAndHook(self: *HookManager, T: type, module: []const u8, patterns: []const []const ?u8, target: *const anyopaque) !T {
    const match = mem.scanUniquePatterns(module, patterns) orelse {
        return error.PatternNotFound;
    };

    return self.hookDetour(T, match.ptr, target);
}

pub fn hookVMT(self: *HookManager, vt: [*]*const anyopaque, index: usize, target: anytype) !@TypeOf(target) {
    var hook = try Hook.hookVMT(vt, index, target);
    errdefer hook.unhook() catch {};

    try self.hooks.append(hook);

    return @ptrCast(hook.orig.?);
}

pub fn hookDetour(self: *HookManager, func: anytype, target: @TypeOf(func)) !@TypeOf(func) {
    var hook = try Hook.hookDetour(@constCast(func), target, self.exec_page);
    errdefer hook.unhook() catch {};

    try self.hooks.append(hook);

    self.exec_page = self.exec_page[hook.data.detour.trampoline.len..];

    return @ptrCast(hook.orig.?);
}

pub fn hookSymbol(self: *HookManager, module_name: []const u8, func: [:0]const u8, target: anytype) !@TypeOf(target) {
    comptime std.debug.assert(@typeInfo(@TypeOf(target)) == .pointer);

    var module = try std.DynLib.open(module_name);
    // this just decreases refcount by 1, we don't ever hook things that aren't
    // already loaded so it doesn't matter
    defer module.close();
    const func_ptr = module.lookup(@TypeOf(target), func) orelse return error.NoSuchSymbol;
    return self.hookDetour(func_ptr, target);
}