const VkAllocator = @This();
const std = @import("std");

const ctx = @import("../AppContext.zig");
const vk = @import("vulkan");

const Allocation = struct {
    len: usize,
    stack_trace: [stack_frames]usize,

    fn capture(len: usize) Allocation {
        return .{
            .len = len,
            .stack_trace = stack_trace_buf,
        };
    }

    fn dumpStackTrace(self: *Allocation) void {
        const st = self.getStackTrace();
        std.debug.dumpStackTrace(&st);
    }

    fn getStackTrace(self: *Allocation) std.debug.StackTrace {
        var stack_trace_len: usize = 0;
        while (stack_trace_len < self.stack_trace.len and self.stack_trace[stack_trace_len] != 0) {
            stack_trace_len += 1;
        }
        return .{
            .return_addresses = self.stack_trace[0..stack_trace_len],
            .skipped = if (stack_trace_len < self.stack_trace.len) .none else .unknown,
        };
    }
};

threadlocal var stack_trace_buf: [stack_frames]usize = undefined;

/// Uses `allocator_general`.
allocations: std.AutoHashMapUnmanaged(*anyopaque, Allocation),
mutex: std.Io.Mutex,
allocated_bytes: usize,

interface: vk.AllocationCallbacks,

// Allocation in Zig's standard library relies heavily on the alignment being
// comptime-known, but we only know the alignment at runtime. We could work
// around it, but it would add much complexity. Instead, we allocate everything
// with the same, relatively big alignment and hope for the best.
const actual_alignment: std.mem.Alignment = .@"16";
const debug = @import("builtin").mode == .debug;
const stack_frames: usize = if (debug and std.debug.sys_can_stack_trace) 6 else 0;

const log = std.log.scoped(.vk_allocator);

pub fn init() !*VkAllocator {
    const allocator_persistent = ctx.allocator_persistent;

    // NOTE We allocate the structure to pin its address.
    const vk_allocator = try allocator_persistent.create(VkAllocator);
    errdefer allocator_persistent.destroy(vk_allocator);

    vk_allocator.* = .{
        .allocations = .empty,
        .mutex = .init,
        .allocated_bytes = 0,
        .interface = .{
            .p_user_data = vk_allocator,
            .pfn_allocation = &allocationFunction,
            .pfn_reallocation = &reallocationFunction,
            .pfn_free = &freeFunction,
        },
    };
    return vk_allocator;
}

pub fn deinit(self: *VkAllocator) void {
    const allocator_general = ctx.allocator_general;

    std.log.scoped(.deinit).debug("Deinitializing {*}", .{self});

    if (self.allocated_bytes > 0) {
        log.warn("{d} byte(s) still allocated while deinitializing", .{self.allocated_bytes});
    }

    if (self.allocations.size > 0) {
        log.warn("{d} allocation(s) still tracked while deinitializing", .{self.allocations.size});
        var it = self.allocations.iterator();
        var index: usize = 0;
        while (it.next()) |entry| : (index += 1) {
            const ptr = entry.key_ptr.*;
            const allocation = entry.value_ptr;

            log.warn("Leaked allocation ({d}/{d}) at 0x{x} of {d} byte(s)", .{ index + 1, self.allocations.size, @intFromPtr(ptr), allocation.len });
            allocation.dumpStackTrace();

            const memory = @as([*]align(actual_alignment.toByteUnits()) u8, @ptrCast(@alignCast(ptr)))[0..allocation.len];
            allocator_general.free(memory);
        }
    }

    self.allocations.deinit(allocator_general);
    self.* = undefined;
}

pub noinline fn capture(self: *const VkAllocator) *const vk.AllocationCallbacks {
    const ret_addr = @returnAddress();
    const stack_trace = std.debug.captureCurrentStackTrace(.{ .first_address = ret_addr }, &stack_trace_buf);
    @memset(stack_trace_buf[@min(stack_trace.return_addresses.len, stack_trace_buf.len)..], 0);
    return &self.interface;
}

fn allocationFunction(
    p_user_data: ?*anyopaque,
    size: usize,
    alignment: usize,
    _: vk.SystemAllocationScope,
) callconv(vk.vulkan_call_conv) ?*anyopaque {
    const allocator_general = ctx.allocator_general;
    const io = ctx.io;

    const self: *VkAllocator = @ptrCast(@alignCast(p_user_data.?));

    const desired_alignment = std.mem.Alignment.fromByteUnits(alignment);
    std.debug.assert(std.mem.Alignment.compare(actual_alignment, .gte, desired_alignment));

    self.mutex.lockUncancelable(io);
    defer self.mutex.unlock(io);

    self.allocations.ensureUnusedCapacity(allocator_general, 1) catch return null;
    const memory = allocator_general.alignedAlloc(u8, actual_alignment, size) catch return null;

    self.allocations.putAssumeCapacity(memory.ptr, .capture(size));
    self.allocated_bytes += size;
    //log.debug("Allocated {d} bytes(s) at 0x{x}", .{ size, @intFromPtr(memory.ptr) });

    return memory.ptr;
}

fn reallocationFunction(
    p_user_data: ?*anyopaque,
    maybe_p_original: ?*anyopaque,
    size: usize,
    alignment: usize,
    _: vk.SystemAllocationScope,
) callconv(vk.vulkan_call_conv) ?*anyopaque {
    const allocator_general = ctx.allocator_general;
    const io = ctx.io;

    const self: *VkAllocator = @ptrCast(@alignCast(p_user_data.?));

    const desired_alignment = std.mem.Alignment.fromByteUnits(alignment);
    std.debug.assert(std.mem.Alignment.compare(actual_alignment, .gte, desired_alignment));

    self.mutex.lockUncancelable(io);
    defer self.mutex.unlock(io);

    // NOTE If we were pedantic, we would consider the fact that we might not
    // need unused capacity if the memory doesn't get relocated.
    self.allocations.ensureUnusedCapacity(allocator_general, 1) catch return null;

    const old_memory = if (maybe_p_original) |p_original| blk_then: {
        const old_size = self.allocations.get(p_original).?.len;
        break :blk_then @as([*]align(actual_alignment.toByteUnits()) u8, @ptrCast(@alignCast(p_original)))[0..old_size];
    } else blk_else: {
        break :blk_else @as([]align(actual_alignment.toByteUnits()) u8, &.{});
    };

    const memory = allocator_general.realloc(old_memory, size) catch return null;
    std.debug.assert(std.mem.isAligned(@intFromPtr(memory.ptr), alignment));

    if (maybe_p_original) |p_original| {
        const old_size = self.allocations.fetchRemove(p_original).?.value.len;
        self.allocated_bytes -= old_size;
    }

    self.allocations.putAssumeCapacityNoClobber(memory.ptr, .capture(size));
    self.allocated_bytes += size;
    //log.debug("Reallocated into {d} bytes(s) at 0x{x} from 0x{x}", .{ size, @intFromPtr(memory.ptr), @intFromPtr(maybe_p_original) });

    return memory.ptr;
}

fn freeFunction(
    p_user_data: ?*anyopaque,
    maybe_p_memory: ?*anyopaque,
) callconv(vk.vulkan_call_conv) void {
    const allocator_general = ctx.allocator_general;
    const io = ctx.io;

    const self: *VkAllocator = @ptrCast(@alignCast(p_user_data.?));

    if (maybe_p_memory) |p_memory| {
        self.mutex.lockUncancelable(io);
        defer self.mutex.unlock(io);

        const size = self.allocations.fetchRemove(p_memory).?.value.len;
        self.allocated_bytes -= size;
        const memory = @as([*]align(actual_alignment.toByteUnits()) u8, @ptrCast(@alignCast(p_memory)))[0..size];

        allocator_general.free(memory);
    }
}