//! Module for loading persistent materials.

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

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

const Atom = @import("../engine/Atom.zig").Atom;
const Engine = @import("../engine/Engine.zig");
const Textures = @import("Textures.zig");

pub const Id = enum(u16) {
    // VOLATILE Synchronize explicit values with `init` implementation.

    /// A material ID that can be used as a "null" material. An object with this
    /// material ID must not be rendered. If this ID is rendered, an "error"
    /// replacement material will be used.
    empty,
    _,

    /// Cast an integer into an ID. This can produce an invalid ID.
    pub fn fromInt(value: u16) Id {
        return @enumFromInt(value);
    }

    /// Cast an index into an ID. This can produce an invalid ID. The caller
    /// asserts that the index is not greater than the max ID value.
    pub fn fromIndex(index: usize) Id {
        std.debug.assert(index < max_materials);
        return @enumFromInt(@as(u16, @intCast(index)));
    }

    /// Cast an index into an ID. This can produce an invalid ID. Returns an
    /// error if the index is greater than the max ID value.
    pub fn fromIndexSafe(index: usize) error{Overflow}!Id {
        if (index >= max_materials) return error.Overflow;
        return @enumFromInt(@as(u16, @intCast(index)));
    }

    /// Cast an ID into an integer.
    pub fn toInt(self: Id) u16 {
        return @intFromEnum(self);
    }
};

pub const Key = struct {
    /// Atom representing the filename of the material.
    filename: Atom,
};

/// Maps a key value to a material ID.
map: std.AutoHashMapUnmanaged(Key, Id),
/// Stores all material data in a single contiguous storage buffer. Use the
/// material ID as an index into this buffer.
material_buffer: shaders.MaterialBuffer,
/// The amount of materials currently loaded.
material_count: usize,

/// With `@sizeOf(Material) == 52` and `max_materials == 4096`, the material
/// storage buffer should take 208 kiB in VRAM.
pub const max_materials = 4096;

pub fn init(engine: *Engine, allocator: std.mem.Allocator) !Materials {
    var map: std.AutoHashMapUnmanaged(Key, Id) = .empty;
    errdefer map.deinit(allocator);
    try map.ensureTotalCapacity(allocator, max_materials);

    var material_buffer = try shaders.MaterialBuffer.init(engine, .{
        .usage = .storage,
        .target_queue = .graphics,
        .array_capacity = max_materials,
        .name = "Materials",
    });
    errdefer material_buffer.deinit(engine);

    // VOLATILE Synchronize with explicit values on top of `Id` type.

    try material_buffer.write(engine, .{
        .element_offset = Id.empty.toInt(),
        .elements = &.{
            .{
                .base_color = .zero,
                .emissive = .init(1, 0, 1),
                .ior = 1.45,
                .metallic = 0,
                .normal_scale = 1,
                .occlusion_texture_strength = 1,
                .roughness = 1,
                .base_color_texture = .empty_base_color,
                .emissive_texture = .empty_emissive,
                .normal_texture = .empty_normal,
                .occlusion_roughness_metallic_texture = .empty_occlusion_roughness_metallic,
            },
        },
    });

    return .{
        .map = map,
        .material_buffer = material_buffer,
        .material_count = @typeInfo(Id).@"enum".fields.len,
    };
}

pub fn deinit(self: *Materials, engine: *Engine, allocator: std.mem.Allocator) void {
    std.log.scoped(.deinit).debug("Deinitializing {*} with {*} and Allocator{{{*},{*}}}", .{ self, engine, allocator.ptr, allocator.vtable });

    self.material_buffer.deinit(engine);
    self.map.deinit(allocator);
    self.* = undefined;
}

/// Get the ID of a material given its filename (as a string). Returns `null` if
/// such material hasn't been loaded. Returns `.empty` when the filename is
/// `null`.
pub fn get(self: *const Materials, maybe_filename: ?[]const u8) ?Id {
    if (maybe_filename) |filename| {
        return self.map.get(.{
            // If the atom doesn't exist, then the material cannot possibly exist.
            .filename = .fromStringIfExists(filename) orelse return null,
        });
    } else {
        return .empty;
    }
}

/// Get the ID of a material given its filename (as an atom). Returns `null` if
/// such material hasn't been loaded. Returns `.empty` when the filename is
/// `.empty`.
pub fn getAtom(self: *const Materials, filename: Atom) ?Id {
    if (filename != .empty) {
        return self.map.get(.{
            .filename = filename,
        });
    } else {
        return .empty;
    }
}

/// Get the ID of a material given its filename (as a string). Returns either an
/// existing material ID or loads a new material along with its textures and
/// assigns a new ID, if necessary. Will not return any error if the material
/// already exists. Returns `.empty` when the filename is `null`.
///
/// When a material or its textures are being loaded, `temp_allocator` is used
/// for temporary allocations necessary to perform all operations. No memory
/// allocated with `temp_allocator` is retained, so the allocator can be
/// deinitialized or reset after this function returns. Note that during loading
/// the engine will make its own persistent allocations, so an out of memory
/// error is not necessarily related to `temp_allocator`.
pub fn getOrLoad(self: *Materials, engine: *Engine, textures: *Textures, maybe_filename: ?[]const u8, temp_allocator: std.mem.Allocator) !Id {
    if (maybe_filename) |filename| {
        const key: Key = .{
            // If the material already exists, then the atom must exist and the
            // following line will not return any error.
            .filename = try .fromString(filename),
        };

        // We don't use `getOrPutAssumeCapacity` method, because we might already be
        // at full capacity, in which case we should return `error.OutOfMaterials`.

        if (self.map.get(key)) |id| {
            return id;
        } else {
            const id = Id.fromIndexSafe(self.material_count) catch |err| switch (err) {
                error.Overflow => return error.OutOfMaterials,
            };
            try self.loadMaterial(engine, textures, filename, id.toInt(), temp_allocator);

            self.map.putAssumeCapacityNoClobber(key, id);
            self.material_count += 1;

            return id;
        }
    } else {
        return .empty;
    }
}

/// Get the ID of a material given its filename (as an atom). Returns either an
/// existing material ID or loads a new material along with its textures and
/// assigns a new ID, if necessary. Will not return any error if the material
/// already exists. Returns `.empty` when the filename is `.empty`.
///
/// When a material or its textures are being loaded, `temp_allocator` is used
/// for temporary allocations necessary to perform all operations. No memory
/// allocated with `temp_allocator` is retained, so the allocator can be
/// deinitialized or reset after this function returns. Note that during loading
/// the engine will make its own persistent allocations, so an out of memory
/// error is not necessarily related to `temp_allocator`.
pub fn getOrLoadAtom(self: *Materials, engine: *Engine, textures: *Textures, filename: Atom, temp_allocator: std.mem.Allocator) !Id {
    if (filename != .empty) {
        const key: Key = .{
            .filename = filename,
        };

        // We don't use `getOrPutAssumeCapacity` method, because we might already be
        // at full capacity, in which case we should return `error.OutOfMaterials`.

        if (self.map.get(key)) |id| {
            return id;
        } else {
            const id = Id.fromIndexSafe(self.material_count) catch |err| switch (err) {
                error.Overflow => return error.OutOfMaterials,
            };
            try self.loadMaterial(engine, textures, filename.toString(), id.toInt(), temp_allocator);

            self.map.putAssumeCapacityNoClobber(key, id);
            self.material_count += 1;

            return id;
        }
    } else {
        return .empty;
    }
}

/// Scan the materials directory and load all materials that haven't been loaded
/// already. Will not traverse subdirectories. Any errors will be skipped and
/// logged.
///
/// When a material or its textures are being loaded, `temp_allocator` is used
/// for temporary allocations necessary to perform all operations. No memory
/// allocated with `temp_allocator` is retained, so the allocator can be
/// deinitialized or reset after this function returns. Note that during loading
/// the engine will make its own persistent allocations, so an out of memory
/// error is not necessarily related to `temp_allocator`.
pub fn loadAll(self: *Materials, engine: *Engine, textures: *Textures, temp_allocator: std.mem.Allocator) void {
    const cwd = std.fs.cwd();

    var dir = cwd.openDir("assets/materials", .{ .iterate = true }) catch |err| {
        std.log.err("Error while opening metarials directory: {s}", .{@errorName(err)});
        return;
    };
    defer dir.close();

    var it = dir.iterate();
    while (it.next() catch |err| {
        std.log.err("Error while iterating over materials directory: {s}", .{@errorName(err)});
        return;
    }) |entry| {
        if (entry.kind != .file) {
            std.log.warn("Skipping material entry {s}, which is not a file", .{entry.name});
            continue;
        }

        _ = self.getOrLoad(engine, textures, entry.name, temp_allocator) catch |err| {
            std.log.err("Error while loading material entry {s}: {s}", .{ entry.name, @errorName(err) });
        };
    }
}

fn loadMaterial(self: *Materials, engine: *Engine, textures: *Textures, filename: []const u8, index: u32, temp_allocator: std.mem.Allocator) !void {
    const MaterialJson = struct {
        baseColor: [3]f32 = .{ 1, 1, 1 },
        baseColorTexture: ?[]const u8 = null,
        emissive: [3]f32 = .{ 0, 0, 0 },
        emissiveTexture: ?[]const u8 = null,
        ior: f32 = 1.45,
        metallic: f32 = 1,
        normalScale: f32 = 1,
        normalTexture: ?[]const u8 = null,
        occlusionRoughnessMetallicTexture: ?[]const u8 = null,
        occlusionTextureStrength: f32 = 1,
        roughness: f32 = 1,
    };

    std.log.debug("Loading material \"{s}\"...", .{filename});

    const cwd = std.fs.cwd();

    var dir = try cwd.openDir("assets/materials", .{});
    defer dir.close();

    // NOTE Buffer size approximated based on expected JSON structure.
    var buffer: [512]u8 = undefined;

    const file = try dir.openFile(filename, .{});
    defer file.close();

    var file_reader = file.reader(&buffer);
    var json_reader = std.json.Reader.init(temp_allocator, &file_reader.interface);
    defer json_reader.deinit();

    const parsed: std.json.Parsed(MaterialJson) = try std.json.parseFromTokenSource(MaterialJson, temp_allocator, &json_reader, .{
        .duplicate_field_behavior = .@"error",
        .ignore_unknown_fields = false,
        .allocate = .alloc_if_needed,
    });
    defer parsed.deinit();

    const material_json = parsed.value;

    try self.material_buffer.write(engine, .{
        .element_offset = index,
        .elements = &.{
            .{
                .base_color = .initArray(material_json.baseColor),
                .emissive = .initArray(material_json.emissive),
                .ior = material_json.ior,
                .metallic = material_json.metallic,
                .normal_scale = material_json.normalScale,
                .occlusion_texture_strength = material_json.occlusionTextureStrength,
                .roughness = material_json.roughness,
                .base_color_texture = try textures.getOrLoad(engine, material_json.baseColorTexture, .base_color, temp_allocator),
                .emissive_texture = try textures.getOrLoad(engine, material_json.emissiveTexture, .emissive, temp_allocator),
                .normal_texture = try textures.getOrLoad(engine, material_json.normalTexture, .normal, temp_allocator),
                .occlusion_roughness_metallic_texture = try textures.getOrLoad(engine, material_json.occlusionRoughnessMetallicTexture, .occlusion_roughness_metallic, temp_allocator),
            },
        },
    });
}