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

var stringsArena: std.heap.ArenaAllocator = undefined;

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

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

const OutputFormat = struct {
    extension: []const u8,
    args: []const []const u8,
};

const ShaderFile = struct {
    path: []const u8,
    base_name: []const u8,
    source_dir: []const u8,
};

fn findContentDir(allocator: std.mem.Allocator) ![]const u8 {
    // First try local content directory
    var content_dir = std.fs.cwd().openDir("content", .{}) catch |err| switch (err) {
        error.FileNotFound => {
            // Walk up directories looking for content.txt
            var current_dir = try std.fs.cwd().realpathAlloc(allocator, ".");
            defer allocator.free(current_dir);

            var last_path: ?[]const u8 = null;
            while (last_path == null or !std.mem.eql(u8, current_dir, last_path.?)) {
                const content_txt_path = try std.fs.path.join(allocator, &.{ current_dir, "content.txt" });
                defer allocator.free(content_txt_path);

                if (std.fs.cwd().openFile(content_txt_path, .{})) |file| {
                    defer file.close();

                    const content = try file.readToEndAlloc(allocator, 1024);
                    defer allocator.free(content);

                    const rel_path = std.mem.trim(u8, content, " \t\n\r");
                    const abs_path = try std.fs.path.join(allocator, &.{ current_dir, rel_path });
                    return try stringsArena.allocator().dupe(u8, abs_path);
                } else |_| {
                    last_path = current_dir;
                    const parent = std.fs.path.dirname(current_dir) orelse break;
                    current_dir = try allocator.dupe(u8, parent);
                    if (last_path) |lp| allocator.free(lp);
                }
            }

            return error.ContentDirNotFound;
        },
        else => return err,
    };
    defer content_dir.close();

    const cwd_path = try std.fs.cwd().realpathAlloc(allocator, ".");
    defer allocator.free(cwd_path);

    const content_path = try std.fs.path.join(allocator, &.{ cwd_path, "content" });
    return try stringsArena.allocator().dupe(u8, content_path);
}

fn getShadercrossPath(allocator: std.mem.Allocator) ![]const u8 {
    const cwd_path = try std.fs.cwd().realpathAlloc(allocator, ".");
    defer allocator.free(cwd_path);

    const lib_path = try std.fs.path.join(allocator, &.{ cwd_path, "..", "..", "lib", "sdl3", "shadercross", "bin" });
    defer allocator.free(lib_path);

    const platform_dir = switch (builtin.os.tag) {
        .linux => "linux",
        .macos => "osx",
        .windows => "win64",
        else => return error.UnsupportedPlatform,
    };

    const platform_path = try std.fs.path.join(allocator, &.{ lib_path, platform_dir });
    defer allocator.free(platform_path);

    const exe_name = if (builtin.os.tag == .windows) "shadercross.exe" else "shadercross";
    const exe_path = try std.fs.path.join(allocator, &.{ platform_path, exe_name });

    return try stringsArena.allocator().dupe(u8, exe_path);
}

fn discoverShaders(allocator: std.mem.Allocator) !std.ArrayList(ShaderFile) {
    var result = std.ArrayList(ShaderFile).init(allocator);

    const cwd_path = try std.fs.cwd().realpathAlloc(allocator, ".");
    defer allocator.free(cwd_path);

    const engine_path = try std.fs.path.join(allocator, &.{ cwd_path, "..", "..", "engine" });
    defer allocator.free(engine_path);

    var engine_dir = std.fs.cwd().openDir(engine_path, .{ .iterate = true }) catch |err| {
        printErr("Failed to open engine directory: {}", .{err});
        return result;
    };
    defer engine_dir.close();

    var engine_iter = engine_dir.iterate();
    while (try engine_iter.next()) |entry| {
        if (entry.kind != .directory) continue;

        const module_path = try std.fs.path.join(allocator, &.{ engine_path, entry.name });
        defer allocator.free(module_path);

        const shaders_path = try std.fs.path.join(allocator, &.{ module_path, "shaders" });
        defer allocator.free(shaders_path);

        var shaders_dir = std.fs.cwd().openDir(shaders_path, .{ .iterate = true }) catch continue;
        defer shaders_dir.close();

        var shader_iter = shaders_dir.iterate();
        while (try shader_iter.next()) |shader_entry| {
            if (shader_entry.kind != .file) continue;
            if (!std.mem.endsWith(u8, shader_entry.name, ".hlsl")) continue;

            const full_path = try std.fs.path.join(allocator, &.{ shaders_path, shader_entry.name });
            const base_name = shader_entry.name[0 .. shader_entry.name.len - 5]; // Remove .hlsl

            try result.append(.{
                .path = try stringsArena.allocator().dupe(u8, full_path),
                .base_name = try stringsArena.allocator().dupe(u8, base_name),
                .source_dir = try stringsArena.allocator().dupe(u8, shaders_path),
            });
        }
    }

    return result;
}

fn runCommand(allocator: std.mem.Allocator, argv: []const []const u8) !void {
    printLog("Running command: {s}", .{try std.mem.join(allocator, " ", argv)});

    var child = std.process.Child.init(argv, allocator);
    child.stdout_behavior = .Inherit;
    child.stderr_behavior = .Inherit;

    const term = try child.spawnAndWait();
    switch (term) {
        .Exited => |code| {
            if (code != 0) {
                return error.CommandFailed;
            }
        },
        else => return error.CommandFailed,
    }
}

fn ensureDir(path: []const u8) !void {
    std.fs.cwd().makePath(path) catch |err| switch (err) {
        error.PathAlreadyExists => {},
        else => return err,
    };
}

fn cookShaders(allocator: std.mem.Allocator, shader_files: []const ShaderFile, content_dir: []const u8, shadercross_path: []const u8) !void {
    const cooked_root = try std.fs.path.join(allocator, &.{ content_dir, "_shaders" });
    defer allocator.free(cooked_root);

    const output_formats = [_]OutputFormat{
        .{ .extension = "dxil", .args = &.{} },
        .{ .extension = "spv", .args = &.{} },
        .{ .extension = "msl", .args = &.{} },
    };

    var spv_files = std.ArrayList(struct { spv_path: []const u8, source_dir: []const u8 }).init(allocator);
    defer spv_files.deinit();

    // Compile shaders to all formats
    for (output_formats) |format| {
        const format_dir = try std.fs.path.join(allocator, &.{ cooked_root, format.extension });
        defer allocator.free(format_dir);

        try ensureDir(format_dir);

        for (shader_files) |shader| {
            const output_name = try std.fmt.allocPrint(allocator, "{s}.{s}", .{ shader.base_name, format.extension });
            defer allocator.free(output_name);

            const output_path = try std.fs.path.join(allocator, &.{ format_dir, output_name });
            defer allocator.free(output_path);

            var cmd_args = std.ArrayList([]const u8).init(allocator);
            defer cmd_args.deinit();

            try cmd_args.append(shadercross_path);
            try cmd_args.append(shader.path);
            for (format.args) |arg| {
                try cmd_args.append(arg);
            }
            try cmd_args.append("-o");
            try cmd_args.append(output_path);

            try runCommand(allocator, cmd_args.items);

            // Store SPV files for reflection processing
            if (std.mem.eql(u8, format.extension, "spv")) {
                try spv_files.append(.{
                    .spv_path = try stringsArena.allocator().dupe(u8, output_path),
                    .source_dir = shader.source_dir,
                });
            }
        }
    }

    // Generate reflection data for SPV files
    for (spv_files.items) |spv_file| {
        // Get file size to ensure it's not empty
        const file = std.fs.cwd().openFile(spv_file.spv_path, .{}) catch continue;
        defer file.close();

        const file_size = try file.getEndPos();
        if (file_size <= 1) {
            printErr("SPV file is empty or too small: {s}", .{spv_file.spv_path});
            continue;
        }

        const base_name = std.fs.path.basename(spv_file.spv_path);
        const name_no_ext = base_name[0 .. base_name.len - 4]; // Remove .spv

        const json_name = try std.fmt.allocPrint(allocator, "{s}.json", .{name_no_ext});
        defer allocator.free(json_name);

        const json_path = try std.fs.path.join(allocator, &.{ spv_file.source_dir, json_name });
        defer allocator.free(json_path);

        const cmd_args = [_][]const u8{
            "spirv-cross",
            spv_file.spv_path,
            "--reflect",
            "--output",
            json_path,
        };

        try runCommand(allocator, &cmd_args);
    }
}

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();

    printLog("Starting shader compilation...", .{});

    // Find content directory
    const content_dir = findContentDir(allocator) catch |err| {
        printErr("Failed to find content directory: {}", .{err});
        return;
    };
    printLog("Content target dir: {s}", .{content_dir});

    // Get shadercross path
    const shadercross_path = getShadercrossPath(allocator) catch |err| {
        printErr("Failed to get shadercross path: {}", .{err});
        return;
    };
    printLog("Shadercross path: {s}", .{shadercross_path});

    // Discover shader files
    var shader_files = discoverShaders(allocator) catch |err| {
        printErr("Failed to discover shaders: {}", .{err});
        return;
    };
    defer shader_files.deinit();

    printLog("Found {} shader files:", .{shader_files.items.len});
    for (shader_files.items) |shader| {
        printLog("  {s}", .{shader.path});
    }

    // Cook shaders
    cookShaders(allocator, shader_files.items, content_dir, shadercross_path) catch |err| {
        printErr("Failed to cook shaders: {}", .{err});
        return;
    };

    printLog("Shader compilation completed successfully!", .{});
}