const std = @import("std");

fn printLog(comptime fmt: []const u8, args: anytype) void {
    //_ = fmt;
    // _ = args;
    std.debug.print(fmt ++ "\n", args);
}

fn printErr(comptime fmt: []const u8, args: anytype) void {
    std.debug.print(fmt ++ "\n", args);
}

const FieldEntry = struct {
    name: []const u8,
    typeName: []const u8,
    offset: u32 = 0,
    typeIndex: ?usize = null, // this is only set if it's a reference to a type in the typemap
};

const TypeEntry = struct {
    name: []u8,
    size: u32 = 0,
    subType: ?usize = null, // if set, this is a list or a uniform, and the underlying type is the real type that we want to generate code for.
    stride: u32 = 0,
    alignment: u32 = 0,
    fields: std.ArrayListUnmanaged(FieldEntry) = .{},
    fieldsArena: std.heap.ArenaAllocator,

    pub fn init(allocator: std.mem.Allocator, name: []const u8) !@This() {
        var arena = std.heap.ArenaAllocator.init(allocator);

        const new: @This() = .{
            .name = try arena.allocator().dupe(u8, name),
            .fieldsArena = arena,
        };

        return new;
    }

    pub fn fieldsAllocator(self: *@This()) std.mem.Allocator {
        return self.fieldsArena.allocator();
    }

    pub fn setName(self: *@This(), name: []const u8) !void {
        self.fieldsAllocator().free(self.name);
        self.name = try self.fieldsAllocator().dupe(u8, name);
    }

    pub fn pushField(self: *@This(), field: FieldEntry) !void {
        try self.fields.append(self.fieldsAllocator(), field);
    }

    pub fn deinit(self: *@This()) void {
        self.fields.deinit(self.fieldsAllocator());
        self.fieldsArena.deinit();
    }
};

var stringsArena: std.heap.ArenaAllocator = undefined;

const BuiltinType = struct {
    name: []const u8,
    size: u32 = 0,
    alignment: u32 = 0,
};

const builtinTypes: []const BuiltinType = &.{
    .{ .name = "int", .size = 4, .alignment = 4 },
    .{ .name = "uint", .size = 4, .alignment = 4 },
    .{ .name = "float", .size = 4, .alignment = 4 },
    .{ .name = "vec2", .size = 8, .alignment = 8 },
    .{ .name = "vec3", .size = 12, .alignment = 16 }, // vec3 has complex alignment rules
    .{ .name = "vec4", .size = 16, .alignment = 16 },
    .{ .name = "u8vec4", .size = 4, .alignment = 1 },
    .{ .name = "mat4", .size = 64, .alignment = 16 }, // mat4 alignment depends on layout
};

const TypeMap = struct {
    types: std.StringHashMapUnmanaged(usize) = .{},
    typeList: std.ArrayListUnmanaged(TypeEntry) = .{},
    allocator: std.mem.Allocator,

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

    pub fn pushType(self: *@This(), typeName: []const u8, entry: TypeEntry) !void {
        try self.types.put(self.allocator, try stringsArena.allocator().dupe(u8, typeName), self.typeList.items.len);
        try self.typeList.append(self.allocator, entry);
    }

    fn isBuiltinType(tag: []const u8) bool {
        const builtinTypeTags: []const []const u8 = &.{
            "vec2",
            "int",
            "uint",
            "u8vec4",
            "vec3",
            "vec4",
            "mat4",
            "float",
        };

        for (builtinTypeTags) |cmp| {
            if (std.mem.eql(u8, cmp, tag)) {
                return true;
            }
        }

        return false;
    }

    pub fn parseTypes(self: *@This(), parsed: std.json.Parsed(std.json.Value)) !void {
        const root = parsed.value.object;
        const typesObj = root.get("types") orelse return;

        var typesIterator = typesObj.object.iterator();
        while (typesIterator.next()) |typeEntry| {
            // const typeId = typeEntry.key_ptr.*;
            const typeData = typeEntry.value_ptr.*.object;

            const typeName = typeData.get("name").?.string;
            var entry = try TypeEntry.init(self.allocator, typeName);

            // Parse members if they exist
            if (typeData.get("members")) |membersValue| {
                const members = membersValue.array;
                for (members.items) |member| {
                    const memberObj = member.object;
                    const fieldName = memberObj.get("name").?.string;
                    const fieldTypeTag = memberObj.get("type").?.string;
                    const fieldOffset = @as(u32, @intCast(memberObj.get("offset").?.integer));

                    var field: FieldEntry = .{
                        .name = fieldName,
                        .typeName = fieldTypeTag,
                        .offset = fieldOffset,
                    };

                    if (!isBuiltinType(fieldTypeTag)) {
                        if (typesObj.object.get(fieldTypeTag)) |*typeRefObj| {
                            if (typeRefObj.object.get("name")) |nameObj| {
                                const fieldType = nameObj.string;

                                field = FieldEntry{
                                    .name = try entry.fieldsAllocator().dupe(u8, fieldName),
                                    .typeName = try entry.fieldsAllocator().dupe(u8, fieldType),
                                };
                            } else {
                                printErr("No field 'name' in type '{s}'", .{fieldTypeTag});
                            }
                        } else {
                            printErr("unable to find type object '{s}' in json", .{fieldTypeTag});
                        }
                    }

                    try entry.pushField(field);
                }
            }

            try self.pushType(typeName, entry);
        }

        // elaborate type indices
        for (self.typeList.items) |*t| {
            for (t.fields.items) |*f| {
                if (self.types.get(f.typeName)) |typeIndex| {
                    f.typeIndex = typeIndex;
                }
            }
        }

        // resolve structured buffer subtypes
        for (self.typeList.items) |*t| {
            const structuredBufferStr = "type.StructuredBuffer.";
            if (std.mem.startsWith(u8, t.name, structuredBufferStr)) {
                if (t.fields.items.len == 0) {
                    printErr("structuredBuffer type {s} has no subfields? that's odd..", .{t.name});
                    continue;
                }
                const subName = t.fields.items[0].typeName;

                t.subType = self.types.get(subName);
                if (t.subType == null) {
                    printErr("structuredBuffer type {s} refers to child type {s} which is not found in the json", .{ t.name, subName });
                }
            }
        }

        // update field sizes
        for (self.typeList.items) |*t| {
            var lastOffset: u32 = 0;
            for (t.fields.items) |field| {
                std.debug.assert(lastOffset <= field.offset);
                t.size = field.offset + self.getFieldSize(field);
                lastOffset = field.offset;
            }
        }
    }

    pub fn getFieldSize(self: *const @This(), field: FieldEntry) u32 {
        // First check if it's a builtin type
        for (builtinTypes) |builtinType| {
            if (std.mem.eql(u8, builtinType.name, field.typeName)) {
                return builtinType.size;
            }
        }

        // If it's a reference to another type in the typemap
        if (field.typeIndex) |typeIndex| {
            if (typeIndex < self.typeList.items.len) {
                return self.typeList.items[typeIndex].size;
            }
        }

        // Unknown type, return 0
        printErr("Unknown field type '{s}', returning size 0", .{field.typeName});
        return 0;
    }

    pub fn printAll(self: *@This()) void {
        printLog("TypeMap contains {} types:", .{self.typeList.items.len});

        var typesIterator = self.types.iterator();
        while (typesIterator.next()) |entry| {
            const typeId = entry.key_ptr.*;
            const typeIndex = entry.value_ptr.*;
            const typeEntry = &self.typeList.items[typeIndex];

            printLog("  Type ID: '{s}' -> Name: '{s}' index: {d}", .{ typeId, typeEntry.name, typeIndex });
            printLog("    Size: {}, Stride: {}, Alignment: {}", .{ typeEntry.size, typeEntry.stride, typeEntry.alignment });

            if (typeEntry.fields.items.len > 0) {
                printLog("    Fields ({}):", .{typeEntry.fields.items.len});
                for (typeEntry.fields.items) |field| {
                    printLog("      - {s}: {s}", .{ field.name, field.typeName });
                }
            } else {
                printLog("    No fields", .{});
            }
            printLog("", .{});
        }
    }

    pub fn deinit(self: *@This()) void {
        for (self.typeList.items) |*t| {
            t.deinit();
        }

        self.typeList.deinit(self.allocator);
        self.types.deinit(self.allocator);
    }
};

const UboEntry = struct {
    name: []u8,
    typeIndex: usize,
    set: u32,
    binding: u32,
    block_size: u32,
};

const UboList = struct {
    list: std.ArrayListUnmanaged(UboEntry) = .{},
    allocator: std.mem.Allocator,

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

    pub fn parseFromJson(self: *@This(), parsed: std.json.Parsed(std.json.Value), map: *const TypeMap) !void {
        const root = parsed.value.object;
        const list = root.get("ubos") orelse {
            printErr("missing ubos entry", .{});
            return;
        };

        if (list != .array) {
            printErr("invalid type for list 'ubos' in root, skipping", .{});
            return;
        }

        for (list.array.items) |i| {
            if (i != .object) {
                printErr("invalid type in ubos json list, expected an object", .{});
                continue;
            }

            const typeField = root.get("types").?.object.get(i.object.get("type").?.string).?.object.get("name").?.string;
            const nameField = i.object.get("name").?.string;
            const blockSizeField = @as(u32, @intCast(i.object.get("block_size").?.integer));
            const setField = @as(u32, @intCast(i.object.get("set").?.integer));
            const bindingField = @as(u32, @intCast(i.object.get("binding").?.integer));

            var entry = UboEntry{
                .name = try stringsArena.allocator().dupe(u8, nameField),
                .typeIndex = undefined,
                .block_size = blockSizeField,
                .set = setField,
                .binding = bindingField,
            };

            if (map.types.get(typeField)) |typeIndex| {
                entry.typeIndex = typeIndex;
            } else {
                printErr("Type for UBO '{s}' of type '{s}' not found in provided typemap", .{ nameField, typeField });
                return error.BadUniformBuffer;
            }

            try self.list.append(self.allocator, entry);
        }
    }

    pub fn printAll(self: *@This(), typeMap: *const TypeMap) void {
        printLog("UboList contains {} UBOs:", .{self.list.items.len});

        for (self.list.items, 0..) |ubo, i| {
            printLog("  UBO {}: '{s}'", .{ i, ubo.name });
            printLog("    Type Index: {}", .{ubo.typeIndex});
            if (ubo.typeIndex < typeMap.typeList.items.len) {
                const typeEntry = &typeMap.typeList.items[ubo.typeIndex];
                printLog("    Type Name: '{s}'", .{typeEntry.name});
            }
            printLog("    Set: {}, Binding: {}", .{ ubo.set, ubo.binding });
            printLog("    Block Size: {}", .{ubo.block_size});
            printLog("", .{});
        }
    }

    pub fn deinit(self: *@This()) void {
        self.list.deinit(self.allocator);
    }
};

const SsboEntry = struct {
    name: []u8,
    typeIndex: usize,
    readOnly: bool,
    set: u32,
    binding: u32,
    block_size: u32,
};

const SsboList = struct {
    list: std.ArrayListUnmanaged(SsboEntry) = .{},
    allocator: std.mem.Allocator,

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

    pub fn parseFromJson(self: *@This(), parsed: std.json.Parsed(std.json.Value), map: *const TypeMap) !void {
        const root = parsed.value.object;
        const list = root.get("ssbos") orelse {
            printErr("missing ssbos entry", .{});
            return;
        };

        if (list != .array) {
            printErr("invalid type for list 'ssbos' in root, skipping", .{});
            return;
        }

        for (list.array.items) |i| {
            if (i != .object) {
                printErr("invalid type in ssbos json list, expected an object", .{});
                continue;
            }

            const typeField = root.get("types").?.object.get(i.object.get("type").?.string).?.object.get("name").?.string;
            const nameField = i.object.get("name").?.string;
            const readOnlyField = i.object.get("readonly").?.bool;
            const blockSizeField = @as(u32, @intCast(i.object.get("block_size").?.integer));
            const setField = @as(u32, @intCast(i.object.get("set").?.integer));
            const bindingField = @as(u32, @intCast(i.object.get("binding").?.integer));

            var entry = SsboEntry{
                .name = try stringsArena.allocator().dupe(u8, nameField),
                .typeIndex = undefined,
                .readOnly = readOnlyField,
                .block_size = blockSizeField,
                .set = setField,
                .binding = bindingField,
            };

            if (map.types.get(typeField)) |typeIndex| {
                entry.typeIndex = typeIndex;
                if (map.typeList.items[typeIndex].subType) |subType| {
                    entry.typeIndex = subType;
                }
            } else {
                printErr("Type for SSBO '{s}' of type '{s}' not found in provided typemap", .{ nameField, typeField });
                return error.BadStructuredBuffer;
            }

            try self.list.append(self.allocator, entry);
        }
    }

    pub fn printAll(self: *@This(), typeMap: *const TypeMap) void {
        printLog("SsboList contains {} SSBOs:", .{self.list.items.len});

        for (self.list.items, 0..) |ssbo, i| {
            printLog("  SSBO {}: '{s}'", .{ i, ssbo.name });
            printLog("    Type Index: {}", .{ssbo.typeIndex});
            if (ssbo.typeIndex < typeMap.typeList.items.len) {
                const typeEntry = &typeMap.typeList.items[ssbo.typeIndex];
                printLog("    Type Name: '{s}'", .{typeEntry.name});
            }
            printLog("    Read Only: {}", .{ssbo.readOnly});
            printLog("    Set: {}, Binding: {}", .{ ssbo.set, ssbo.binding });
            printLog("    Block Size: {}", .{ssbo.block_size});
            printLog("", .{});
        }
    }

    pub fn deinit(self: *@This()) void {
        self.list.deinit(self.allocator);
    }
};

const SamplersEntry = struct {
    name: []const u8,
    set: u32,
    binding: u32,
};

const SamplersList = struct {
    list: std.ArrayListUnmanaged(SamplersEntry) = .{},
    allocator: std.mem.Allocator,

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

    pub fn parseSamplers(self: *@This(), parsed: std.json.Parsed(std.json.Value)) !void {
        const root = parsed.value.object;
        const samplersArray = root.get("separate_samplers") orelse return;

        for (samplersArray.array.items) |samplerItem| {
            const samplerObj = samplerItem.object;

            const name = try stringsArena.allocator().dupe(u8, samplerObj.get("name").?.string);
            const set = @as(u32, @intCast(samplerObj.get("set").?.integer));
            const binding = @as(u32, @intCast(samplerObj.get("binding").?.integer));

            const entry = SamplersEntry{
                .name = name,
                .set = set,
                .binding = binding,
            };

            try self.list.append(self.allocator, entry);
        }
    }

    pub fn printAll(self: *@This()) void {
        printLog("SamplersList contains {} samplers:", .{self.list.items.len});

        for (self.list.items, 0..) |sampler, i| {
            printLog("  Sampler {}: '{s}'", .{ i, sampler.name });
            printLog("    Set: {}, Binding: {}", .{ sampler.set, sampler.binding });
            printLog("", .{});
        }
    }

    pub fn deinit(self: *@This()) void {
        self.list.deinit(self.allocator);
    }
};

const OutputFile = struct {
    output: std.ArrayList(u8),
    typeMap: *const TypeMap,
    ssbos: *const SsboList,
    ubos: *const UboList,
    samplers: *const SamplersList,

    pub fn init(
        allocator: std.mem.Allocator,
        typeMap: *const TypeMap,
        ssbos: *const SsboList,
        ubos: *const UboList,
        samplers: *const SamplersList,
    ) !@This() {
        return .{
            .output = std.ArrayList(u8).init(allocator),
            .typeMap = typeMap,
            .ssbos = ssbos,
            .ubos = ubos,
            .samplers = samplers,
        };
    }

    pub fn finalize(self: *@This()) !void {
        try self.output.append(0);
    }

    pub fn writeOut(self: *@This(), path: []const u8) !void {
        const file = try std.fs.cwd().createFile(path, .{});
        defer file.close();

        var ast = try std.zig.Ast.parse(
            self.output.allocator,
            @as([:0]const u8, @ptrCast(self.output.items[0 .. self.output.items.len - 1])),
            .zig,
        );
        // std.debug.print("output=\n{s}", .{content.items[0 .. content.items.len - 1]});
        defer ast.deinit(self.output.allocator);
        const out = try ast.render(self.output.allocator);

        try file.writeAll(out);
    }

    pub fn generatePreamble(self: *@This()) !void {
        var writer = self.output.writer();
        try writer.print(
            \\pub const shaderTypes = @import("shaderTypes");
            \\
            \\pub const int = shaderTypes.int;
            \\pub const uint = shaderTypes.uint;
            \\
            \\pub const vec2 = shaderTypes.vec2;
            \\pub const u8vec4 = shaderTypes.u8vec4;
            \\pub const vec3 = shaderTypes.vec3;
            \\pub const vec4 = shaderTypes.vec4;
            \\pub const mat4 = shaderTypes.mat4;
            \\pub const float = shaderTypes.float;
            \\
            \\pub const BufferInfo = shaderTypes.BufferInfo;
            \\
        , .{});
    }

    pub fn generateSsbos(self: *@This()) !void {
        var writer = self.output.writer();
        for (self.ssbos.list.items) |ssbo| {
            const ssboType = self.typeMap.typeList.items[ssbo.typeIndex];
            try writer.print("pub const {s} =  extern struct {{\n", .{ssboType.name});

            var lastFieldEnd: u32 = 0;
            var paddingCount: u32 = 0;
            for (ssboType.fields.items) |field| {
                if (lastFieldEnd != field.offset) {
                    printErr("WARNING: Field {s} in ssbo {s} generated padding, expected offset {d} offset + struct size = {d}", .{ field.name, ssbo.name, field.offset, lastFieldEnd });
                    try writer.print("    pad{d}: [{d}]c_char = undefined,\n", .{ paddingCount, field.offset - lastFieldEnd });
                    paddingCount += 1;
                }
                try writer.print("    {s}: {s},\n", .{ field.name, field.typeName });
                lastFieldEnd = field.offset + self.typeMap.getFieldSize(field);
            }

            // generate buffer info

            try writer.print("\n    pub const Buffer: BufferInfo = .{{ .storage = {d} }};\n", .{ssbo.binding});

            try writer.print("\n    pub const FieldDetails: []shaderTypes.FieldDetail = &.{{ \n", .{});
            for (ssboType.fields.items) |field| {
                try writer.print(" " ** 8, .{});
                try writer.print(".{{ .name = \"{s}\", .offset = {d}, .size = {d} }},\n", .{ field.name, field.offset, self.typeMap.getFieldSize(field) });
            }
            try writer.print("    }};\n", .{});

            try writer.print("}};\n", .{});
        }

        try writer.print("\ncomptime {{\n", .{});
        for (self.ssbos.list.items) |ssbo| {
            const ssboType = self.typeMap.typeList.items[ssbo.typeIndex];
            try writer.print("    shaderTypes.ValidateGeneratedStruct({s});\n", .{ssboType.name});
        }
        try writer.print("}}\n", .{});
    }

    pub fn generateUbos(self: *@This()) !void {
        var writer = self.output.writer();
        for (self.ubos.list.items) |ubo| {
            const uboType = self.typeMap.typeList.items[ubo.typeIndex];
            try writer.print("pub const Uniforms =  extern struct {{\n", .{});

            var lastFieldEnd: u32 = 0;
            var paddingCount: u32 = 0;
            for (uboType.fields.items) |field| {
                if (lastFieldEnd != field.offset) {
                    printErr("WARNING: Field {s} in ubo {s} generated padding, expected offset {d} offset + struct size = {d}", .{ field.name, ubo.name, field.offset, lastFieldEnd });
                    try writer.print("    pad{d}: [{d}]c_char = undefined,\n", .{ paddingCount, field.offset - lastFieldEnd });
                    paddingCount += 1;
                }
                try writer.print("    {s}: {s},\n", .{ field.name, field.typeName });
                lastFieldEnd = field.offset + self.typeMap.getFieldSize(field);
            }

            // generate buffer info using .uniform instead of .storage
            try writer.print("\n    pub const Buffer: BufferInfo = .{{ .uniform = {d} }};\n", .{ubo.binding});

            try writer.print("\n    pub const FieldDetails: []shaderTypes.FieldDetail = &.{{ \n", .{});
            for (uboType.fields.items) |field| {
                try writer.print(" " ** 8, .{});
                try writer.print(".{{ .name = \"{s}\", .offset = {d}, .size = {d} }},\n", .{ field.name, field.offset, self.typeMap.getFieldSize(field) });
            }
            try writer.print("    }};\n", .{});

            try writer.print("}};\n", .{});
        }

        try writer.print("\ncomptime {{\n", .{});
        for (self.ubos.list.items) |ubo| {
            const uboType = self.typeMap.typeList.items[ubo.typeIndex];
            _ = uboType;
            try writer.print("    shaderTypes.ValidateGeneratedStruct({s});\n", .{"Uniforms"});
        }
        try writer.print("}}\n", .{});
    }

    pub fn generateLoadArguments(self: *@This()) !void {
        var writer = self.output.writer();

        try writer.print("pub const LoadArgs = shaderTypes.ShaderLoadArgs {{\n", .{});

        try writer.print(
            \\.num_samplers = {d},
            \\.num_storage_textures = {d},
            \\.num_storage_buffers = {d},
            \\.num_uniform_buffers = {d},
        , .{
            self.samplers.list.items.len,
            0, // storage textures to be added later
            self.ssbos.list.items.len,
            self.ubos.list.items.len,
        });
        try writer.print("}};\n", .{});
    }

    pub fn deinit(self: *@This()) void {
        self.output.deinit();
    }
};

pub fn main() !void {
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gpa.deinit();

    var arena = std.heap.ArenaAllocator.init(gpa.allocator());
    defer arena.deinit();

    const allocator = arena.allocator();
    stringsArena = std.heap.ArenaAllocator.init(gpa.allocator());
    defer stringsArena.deinit();

    const args = try std.process.argsAlloc(allocator);
    defer std.process.argsFree(allocator, args);

    if (args.len < 3) {
        printLog("Usage: {s} <jsonFilePath> <outputPath>", .{args[0]});
        return;
    }

    const jsonFilePath = args[1];

    const file = std.fs.cwd().openFile(jsonFilePath, .{}) catch |err| {
        printErr("Error opening file '{s}': {}", .{ jsonFilePath, err });
        return;
    };
    defer file.close();

    const fileSize = try file.getEndPos();
    const jsonContent = try allocator.alloc(u8, fileSize);
    _ = try file.readAll(jsonContent);

    var parsed = std.json.parseFromSlice(std.json.Value, allocator, jsonContent, .{}) catch |err| {
        printErr("Error parsing JSON: {}", .{err});
        return;
    };
    defer parsed.deinit();

    var types: TypeMap = TypeMap.init(allocator);
    try types.parseTypes(parsed);
    defer types.deinit();

    types.printAll();

    // create ssbolist and parse out the ssbolist from the typeslist
    var ssbos = SsboList.init(allocator);
    try ssbos.parseFromJson(parsed, &types);
    defer ssbos.deinit();
    ssbos.printAll(&types);

    // create ubolist and parse out the ubolist from the typeslist
    var ubos = UboList.init(allocator);
    try ubos.parseFromJson(parsed, &types);
    defer ubos.deinit();

    ubos.printAll(&types);

    // parse out the samplers
    var samplers = SamplersList.init(allocator);
    try samplers.parseSamplers(parsed);
    defer samplers.deinit();

    var outputfile = try OutputFile.init(allocator, &types, &ssbos, &ubos, &samplers);
    try outputfile.generatePreamble();
    try outputfile.generateSsbos();
    try outputfile.generateUbos();
    try outputfile.generateLoadArguments();
    printLog("{s}", .{outputfile.output.items});
    try outputfile.finalize();
    try outputfile.writeOut(args[2]);
    defer outputfile.deinit();

    printLog("Successfully parsed JSON from: {s}", .{jsonFilePath});
}