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

const image = @import("image.zig");
const vm = @import("vecmath");

allocator: std.mem.Allocator,
allocations: std.AutoHashMapUnmanaged(*anyopaque, usize) = .empty,
mutex: std.Thread.Mutex = .{},
allocated_bytes: usize = 0,

const alignment: std.mem.Alignment = .@"16";
const VoidPtr = ?*align(alignment.toByteUnits()) anyopaque;
const log = std.log.scoped(.stbi);

pub fn init(allocator: std.mem.Allocator) Self {
    return .{
        .allocator = allocator,
    };
}

pub fn deinit(self: *Self) void {
    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) {
            log.warn("Leaked allocation ({d}/{d}) at 0x{x} of {d} byte(s)", .{
                index + 1,
                self.allocations.size,
                @intFromPtr(entry.key_ptr.*),
                entry.value_ptr.*,
            });
            const memory = @as([*]align(alignment.toByteUnits()) u8, @ptrCast(@alignCast(entry.key_ptr.*)))[0..entry.value_ptr.*];
            self.allocator.free(memory);
        }
    }

    self.allocations.deinit(self.allocator);
}

const import = struct {
    extern fn stbi_load_from_memory(buffer: [*]const u8, len: i32, x: ?*i32, y: ?*i32, channels_in_file: ?*i32, desired_channels: i32) ?[*]align(alignment.toByteUnits()) u8;
    extern fn stbi_load_16_from_memory(buffer: [*]const u8, len: i32, x: ?*i32, y: ?*i32, channels_in_file: ?*i32, desired_channels: i32) ?[*]align(alignment.toByteUnits()) u16;
    extern fn stbi_loadf_from_memory(buffer: [*]const u8, len: i32, x: ?*i32, y: ?*i32, channels_in_file: ?*i32, desired_channels: i32) ?[*]align(alignment.toByteUnits()) f32;

    extern fn stbi_load_from_callbacks(callbacks: *const IoCallbacks, ctx: ?*anyopaque, x: ?*i32, y: ?*i32, channels_in_file: ?*i32, desired_channels: i32) ?[*]align(alignment.toByteUnits()) u8;
    extern fn stbi_load_16_from_callbacks(callbacks: *const IoCallbacks, ctx: ?*anyopaque, x: ?*i32, y: ?*i32, channels_in_file: ?*i32, desired_channels: i32) ?[*]align(alignment.toByteUnits()) u16;
    extern fn stbi_loadf_from_callbacks(callbacks: *const IoCallbacks, ctx: ?*anyopaque, x: ?*i32, y: ?*i32, channels_in_file: ?*i32, desired_channels: i32) ?[*]align(alignment.toByteUnits()) f32;

    extern fn stbi_info_from_memory(buffer: [*]const u8, len: i32, x: ?*i32, y: ?*i32, channels_in_file: ?*i32) i32;
    extern fn stbi_is_16_bit_from_memory(buffer: [*]const u8, len: i32) i32;
    extern fn stbi_is_hdr_from_memory(buffer: [*]const u8, len: i32) i32;

    extern fn stbi_info_from_callbacks(callbacks: *const IoCallbacks, ctx: ?*anyopaque, x: ?*i32, y: ?*i32, channels_in_file: ?*i32) i32;
    extern fn stbi_is_16_bit_from_callbacks(callbacks: *const IoCallbacks, ctx: ?*anyopaque) i32;
    extern fn stbi_is_hdr_from_callbacks(callbacks: *const IoCallbacks, ctx: ?*anyopaque) i32;
};

pub fn loadStaticBuf(self: *Self, comptime W: u32, comptime H: u32, buf: []const u8) !image.Static(W, H) {
    current_self = self;
    defer current_self = undefined;

    var x: i32 = undefined;
    var y: i32 = undefined;

    const res = import.stbi_load_from_memory(buf.ptr, @intCast(buf.len), &x, &y, null, 4) orelse return error.StbiError;
    defer castle_media_stbi_free(res);

    if (x != W or y != H) return error.WrongDimensions;

    return .{ .data = @as(*const [W * H]vm.Color, @ptrCast(@alignCast(res))).* };
}

pub fn loadStaticIo(self: *Self, comptime W: u32, comptime H: u32, reader: *std.io.Reader) !image.Static(W, H) {
    current_self = self;
    defer current_self = undefined;

    var x: i32 = undefined;
    var y: i32 = undefined;

    const res = import.stbi_load_from_callbacks(&.std_io_reader_interface, reader, &x, &y, null, 4) orelse return error.StbiError;
    defer castle_media_stbi_free(res);

    if (x != W or y != H) return error.WrongDimensions;

    return .{ .data = @as(*const [W * H]vm.Color, @ptrCast(@alignCast(res))).* };
}

/// On success, must free memory by calling `freeDynamic` method.
pub fn loadDynamicBuf(self: *Self, buf: []const u8) !image.Dynamic {
    current_self = self;
    defer current_self = undefined;

    var x: i32 = undefined;
    var y: i32 = undefined;

    const res = import.stbi_load_from_memory(buf.ptr, @intCast(buf.len), &x, &y, null, 4) orelse return error.StbiError;

    const buffer_ptr: [*]vm.Color = @ptrCast(@alignCast(res));
    const ux: u32 = @intCast(x);
    const uy: u32 = @intCast(y);

    return .initBuffer(ux, uy, buffer_ptr[0 .. ux * uy]);
}

/// On success, must free memory by calling `freeDynamic` method.
pub fn loadDynamicIo(self: *Self, reader: *std.io.Reader) !image.Dynamic {
    current_self = self;
    defer current_self = undefined;

    var x: i32 = undefined;
    var y: i32 = undefined;

    const res = import.stbi_load_from_callbacks(&.std_io_reader_interface, reader, &x, &y, null, 4) orelse return error.StbiError;

    const buffer_ptr: [*]vm.Color = @ptrCast(@alignCast(res));
    const ux: u32 = @intCast(x);
    const uy: u32 = @intCast(y);

    return .initBuffer(ux, uy, buffer_ptr[0 .. ux * uy]);
}

/// On success, must free memory by calling `freeHdr` method.
pub fn loadHdrBuf(self: *Self, buf: []const u8) !image.Hdr {
    current_self = self;
    defer current_self = undefined;

    var x: i32 = undefined;
    var y: i32 = undefined;

    const res_f32 = import.stbi_loadf_from_memory(buf.ptr, @intCast(buf.len), &x, &y, null, 4) orelse return error.StbiError;
    defer castle_media_stbi_free(res_f32);

    const ux: u32 = @intCast(x);
    const uy: u32 = @intCast(y);
    const buffer_ptr_f32: [*]vm.Vector4 = @ptrCast(res_f32);
    const buffer_ptr_f16: [*]vm.ColorHdr = @ptrCast(castle_media_stbi_malloc(ux * uy * @sizeOf(vm.ColorHdr)) orelse return error.OutOfMemory);
    errdefer castle_media_stbi_free(buffer_ptr_f16);

    for (buffer_ptr_f16[0 .. ux * uy], buffer_ptr_f32[0 .. ux * uy]) |*sample_f16, sample_f32| {
        sample_f16.* = .init(
            std.math.clamp(@as(f16, @floatCast(sample_f32.x)), -std.math.floatMax(f16), std.math.floatMax(f16)),
            std.math.clamp(@as(f16, @floatCast(sample_f32.y)), -std.math.floatMax(f16), std.math.floatMax(f16)),
            std.math.clamp(@as(f16, @floatCast(sample_f32.z)), -std.math.floatMax(f16), std.math.floatMax(f16)),
            std.math.clamp(@as(f16, @floatCast(sample_f32.w)), -std.math.floatMax(f16), std.math.floatMax(f16)),
        );
    }

    return .initBuffer(ux, uy, buffer_ptr_f16[0 .. ux * uy]);
}

/// On success, must free memory by calling `freeHdr` method.
pub fn loadHdrIo(self: *Self, reader: *std.io.Reader) !image.Hdr {
    current_self = self;
    defer current_self = undefined;

    var x: i32 = undefined;
    var y: i32 = undefined;

    const res_f32 = import.stbi_loadf_from_callbacks(&.std_io_reader_interface, reader, &x, &y, null, 4) orelse return error.StbiError;
    defer castle_media_stbi_free(res_f32);

    const ux: u32 = @intCast(x);
    const uy: u32 = @intCast(y);
    const buffer_ptr_f32: [*]vm.Vector4 = @ptrCast(res_f32);
    const buffer_ptr_f16: [*]vm.ColorHdr = @ptrCast(castle_media_stbi_malloc(ux * uy * @sizeOf(vm.ColorHdr)) orelse return error.OutOfMemory);
    errdefer castle_media_stbi_free(buffer_ptr_f16);

    for (buffer_ptr_f16[0 .. ux * uy], buffer_ptr_f32[0 .. ux * uy]) |*sample_f16, sample_f32| {
        sample_f16.* = .init(
            std.math.clamp(@as(f16, @floatCast(sample_f32.x)), -std.math.floatMax(f16), std.math.floatMax(f16)),
            std.math.clamp(@as(f16, @floatCast(sample_f32.y)), -std.math.floatMax(f16), std.math.floatMax(f16)),
            std.math.clamp(@as(f16, @floatCast(sample_f32.z)), -std.math.floatMax(f16), std.math.floatMax(f16)),
            std.math.clamp(@as(f16, @floatCast(sample_f32.w)), -std.math.floatMax(f16), std.math.floatMax(f16)),
        );
    }

    return .initBuffer(ux, uy, buffer_ptr_f16[0 .. ux * uy]);
}

pub fn freeDynamic(self: *Self, img: image.Dynamic) void {
    current_self = self;
    defer current_self = undefined;

    castle_media_stbi_free(@ptrCast(@alignCast(img.data)));
}

pub fn freeHdr(self: *Self, img: image.Hdr) void {
    current_self = self;
    defer current_self = undefined;

    castle_media_stbi_free(@ptrCast(@alignCast(img.data)));
}

// --- IO INTERFACE ------------------------------------------------------------

pub const IoCallbacks = extern struct {
    /// Fill `data` with `size` bytes. Return number of bytes actually read.
    read: ?*const fn (ctx: ?*anyopaque, data: [*]u8, size: i32) callconv(.c) i32,
    /// Skip the next `n` bytes, or backtrack `-n` bytes if `n < 0`.
    skip: ?*const fn (ctx: ?*anyopaque, n: i32) callconv(.c) i32,
    /// Return non-zero value if at the end of file/data.
    eof: ?*const fn (cxt: ?*anyopaque) callconv(.c) i32,

    pub const std_io_reader_interface: IoCallbacks = .{
        .read = stdIoReader_ReadFn,
        .skip = stdIoReader_SkipFn,
        .eof = stdIoReader_EofFn,
    };

    pub fn stdIoReader_ReadFn(ctx: ?*anyopaque, data: [*]u8, size: i32) callconv(.c) i32 {
        const reader: *std.Io.Reader = @ptrCast(@alignCast(ctx.?));
        const bytes_read = reader.readSliceShort(data[0..@intCast(size)]) catch return 0;
        return @intCast(bytes_read);
    }

    pub fn stdIoReader_SkipFn(ctx: ?*anyopaque, n: i32) callconv(.c) i32 {
        const reader: *std.Io.Reader = @ptrCast(@alignCast(ctx.?));
        // NOTE stb_image.h actually discards the return value from this
        // callback. If an actual error occurs, we're cooked (but it will be
        // very likely caught as a parsing error later).
        _ = reader.discardAll(@intCast(n)) catch return 0;
        return 0;
    }

    pub fn stdIoReader_EofFn(ctx: ?*anyopaque) callconv(.c) i32 {
        const reader: *std.Io.Reader = @ptrCast(@alignCast(ctx.?));

        _ = reader.peekByte() catch return 1;
        return 0;
    }
};

// --- MALLOC INTERFACE --------------------------------------------------------

threadlocal var current_self: *Self = undefined;

export fn castle_media_stbi_malloc(size: usize) callconv(.c) VoidPtr {
    const self = current_self;

    self.mutex.lock();
    defer self.mutex.unlock();

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

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

    return memory.ptr;
}

export fn castle_media_stbi_realloc(maybe_ptr: VoidPtr, size: usize) callconv(.c) VoidPtr {
    const self = current_self;

    self.mutex.lock();
    defer self.mutex.unlock();

    // 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(self.allocator, 1) catch return null;

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

    const memory = self.allocator.realloc(old_memory, size) catch return null;

    if (maybe_ptr) |ptr| {
        const old_size = self.allocations.fetchRemove(ptr).?.value;
        self.allocated_bytes -= old_size;
    }

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

    return memory.ptr;
}

export fn castle_media_stbi_free(maybe_ptr: VoidPtr) callconv(.c) void {
    const self = current_self;

    if (maybe_ptr) |ptr| {
        self.mutex.lock();
        defer self.mutex.unlock();

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

        self.allocator.free(memory);
    }
}