//------------------------------------------------------------------------------
//  offscreen.zig
//
//  Render to an offscreen rendertarget texture, and use this texture
//  for rendering to the display.
//------------------------------------------------------------------------------
const std = @import("std");
const stbi = @import("stbi");
const sokol = @import("sokol");
const slog = sokol.log;
const sg = sokol.gfx;
const sapp = sokol.app;
const sglue = sokol.glue;
const simgui = sokol.imgui;
const sappimgui = sokol.appimgui;
const ig = @import("cimgui_docking");
const shd = @import("shaders/main.glsl.zig");

const state = struct {
    var alloc: std.mem.Allocator = undefined;
    var io: std.Io = undefined;

    const Width = 800;
    const Height = 800;
    const SampleCount = 1;
    var show: i32 = 0;

    const raymarch = struct {
        var pass_action: sg.PassAction = .{};
        var pip: sg.Pipeline = .{};
        var bind: sg.Bindings = .{};
        var fs: shd.FsParams = .{
            .resolution = .{ Width, Height },
            .rayCount = 24,
            .maxSteps = 256,
            .size = 200,
        };

        var pixels: [state.Width * state.Height * 4]u8 = @splat(255);
        var pixels_updated: bool = true;
        var pixels_img: sg.Image = undefined;
    };
    const nearest = struct {
        var pass_action: sg.PassAction = .{};
        var pass: sg.Pass = .{};
        var pip1: sg.Pipeline = .{};
        var pip2: sg.Pipeline = .{};
        var bind: sg.Bindings = .{};

        var render_tex1: sg.View = .{};
        var render_tex2: sg.View = .{};
        var input_tex1: sg.View = .{};
        var input_tex2: sg.View = .{};
        var depth_img: sg.Image = .{};
    };
};

export fn init() void {
    sg.setup(.{
        .environment = sglue.environment(),
        .logger = .{ .func = slog.func },
    });

    simgui.setup(.{
        .logger = .{ .func = slog.func },
    });
    sappimgui.setup();
    ig.igGetIO().*.ConfigFlags |= ig.ImGuiConfigFlags_DockingEnable;

    state.raymarch.pass_action.colors[0] = .{
        .load_action = .LOAD,
        .store_action = .STORE,
    };
    state.raymarch.pass_action.depth = .{
        .load_action = .LOAD,
        .store_action = .STORE,
    };

    state.nearest.pass_action.colors[0] = .{
        .load_action = .CLEAR,
    };

    state.raymarch.pixels_img = sg.makeImage(.{
        .usage = .{ .dynamic_update = true },
        .width = state.Width,
        .height = state.Height,
        .sample_count = state.SampleCount,
        .pixel_format = .RGBA8,
    });
    const nearest_tex1 = sg.makeImage(.{
        .usage = .{ .color_attachment = true },
        .width = state.Width,
        .height = state.Height,
        .sample_count = state.SampleCount,
        .pixel_format = .RGBA8,
    });
    const nearest_tex2 = sg.makeImage(.{
        .usage = .{ .color_attachment = true },
        .width = state.Width,
        .height = state.Height,
        .sample_count = state.SampleCount,
        .pixel_format = .RGBA8,
    });
    const nearest_depth_img = sg.makeImage(.{
        .usage = .{ .depth_stencil_attachment = true },
        .width = state.Width,
        .height = state.Height,
        .sample_count = state.SampleCount,
        .pixel_format = .DEPTH,
    });
    state.nearest.render_tex1 = sg.makeView(.{
        .color_attachment = .{ .image = nearest_tex1 },
    });
    state.nearest.render_tex2 = sg.makeView(.{
        .color_attachment = .{ .image = nearest_tex2 },
    });
    state.nearest.input_tex1 = sg.makeView(.{
        .texture = .{ .image = nearest_tex1 },
    });
    state.nearest.input_tex2 = sg.makeView(.{
        .texture = .{ .image = nearest_tex2 },
    });

    state.raymarch.bind.views[shd.VIEW_off_tex] = sg.makeView(.{
        .texture = .{ .image = state.raymarch.pixels_img },
    });
    state.nearest.pass.attachments.depth_stencil = sg.makeView(.{
        .depth_stencil_attachment = .{ .image = nearest_depth_img },
    });

    const vbuf = sg.makeBuffer(.{ .data = sg.asRange(&[6]f32{
        -1.0, -1.0,
        -1.0, 3.0,
        3.0,  -1.0,
    }) });
    state.raymarch.bind.vertex_buffers[0] = vbuf;
    state.nearest.bind.vertex_buffers[0] = vbuf;

    // shader and pipeline object for the default render pass
    state.raymarch.pip = sg.makePipeline(.{
        .shader = sg.makeShader(shd.raymarchShaderDesc(sg.queryBackend())),
        .layout = init: {
            var l = sg.VertexLayoutState{};
            l.attrs[shd.ATTR_raymarch_vpos] = .{ .format = .FLOAT2 };
            break :init l;
        },
        .cull_mode = .BACK,
        .depth = .{
            .pixel_format = .DEPTH,
            .compare = .LESS_EQUAL,
            .write_enabled = true,
        },
    });

    state.nearest.pip1 = sg.makePipeline(.{
        .shader = sg.makeShader(shd.nearestShaderDesc(sg.queryBackend())),
        .layout = init: {
            var l = sg.VertexLayoutState{};
            l.attrs[shd.ATTR_nearest_vpos] = .{ .format = .FLOAT2 };
            break :init l;
        },
        .cull_mode = .BACK,
        .sample_count = state.SampleCount,
        .depth = .{
            .pixel_format = .DEPTH,
            .compare = .LESS_EQUAL,
            .write_enabled = true,
        },
        .colors = init: {
            var c: [sg.max_color_attachments]sg.ColorTargetState = @splat(.{});
            c[0].pixel_format = .RGBA8;
            break :init c;
        },
    });
    state.nearest.pip2 = sg.makePipeline(.{
        .shader = sg.makeShader(shd.nearestShaderDesc(sg.queryBackend())),
        .layout = init: {
            var l = sg.VertexLayoutState{};
            l.attrs[shd.ATTR_nearest_vpos] = .{ .format = .FLOAT2 };
            break :init l;
        },
        .cull_mode = .BACK,
    });

    // a sampler object for sampling the render target texture
    state.raymarch.bind.samplers[shd.SMP_off_smp] = sg.makeSampler(.{
        .min_filter = .LINEAR,
        .mag_filter = .LINEAR,
        .wrap_u = .REPEAT,
        .wrap_v = .REPEAT,
    });
    state.nearest.bind.samplers[shd.SMP_scr_smp] = sg.makeSampler(.{
        .min_filter = .LINEAR,
        .mag_filter = .LINEAR,
        .wrap_u = .REPEAT,
        .wrap_v = .REPEAT,
    });
}

export fn event(ev: [*c]const sapp.Event) void {
    sappimgui.trackEvent(ev.*);
    _ = simgui.handleEvent(ev.*);

    switch (ev.*.type) {
        .KEY_DOWN => {
            if (ev.*.key_code == .ESCAPE) {
                sapp.requestQuit();
            }
        },
        else => {},
    }
}

export fn frame() void {
    // call simgui.newFrame() before any ImGui calls
    sappimgui.trackFrame();
    simgui.newFrame(.{
        .width = sapp.width(),
        .height = sapp.height(),
        .delta_time = sapp.frameDuration(),
        .dpi_scale = sapp.dpiScale(),
    });

    // const dt: f32 = @floatCast(sapp.frameDuration() * 60);
    // const aspect = sapp.widthf() / sapp.heightf();
    state.raymarch.fs.resolution[0] = sapp.widthf();
    state.raymarch.fs.resolution[1] = sapp.heightf();

    if (ig.igBegin("abc", null, 0)) {
        // general
        _ = ig.igInputInt("show", &state.show);
        // fs
        _ = ig.igInputInt("rayCount", &state.raymarch.fs.rayCount);
        _ = ig.igInputInt("maxSteps", &state.raymarch.fs.maxSteps);
        _ = ig.igDragFloatEx("size", &state.raymarch.fs.size, 1, 0, 1000, "%f", 0);

        if (ig.igButton("load png")) {
            var x: c_int = 0;
            var y: c_int = 0;
            var n: c_int = 0;
            const img = stbi.stbi_load("test.png", &x, &y, &n, 0);
            @memcpy(&state.raymarch.pixels, img);
            stbi.stbi_image_free(img);
            state.raymarch.pixels_updated = true;
        }
    }
    ig.igEnd();

    if (state.raymarch.pixels_updated) {
        state.raymarch.pixels_updated = false;
        var imageData: sg.ImageData = .{};
        imageData.mip_levels[0] = sg.asRange(&state.raymarch.pixels);
        sg.updateImage(state.raymarch.pixels_img, imageData);
    }

    if (state.show == 0) {
        sg.beginPass(.{ .action = state.raymarch.pass_action, .swapchain = sglue.swapchain() });
        sg.applyPipeline(state.raymarch.pip);
        sg.applyBindings(state.raymarch.bind);
        sg.applyUniforms(shd.UB_fs_params, sg.asRange(&state.raymarch.fs));
        sg.draw(0, 3, 1);
        simgui.render();
        sg.endPass();
    } else if (state.show == 1) {
        // render to tex1
        state.nearest.pass.attachments.colors[0] = state.nearest.render_tex1;
        state.nearest.bind.views[shd.VIEW_scr_tex] = state.nearest.input_tex2;
        sg.beginPass(state.nearest.pass);
        sg.applyPipeline(state.nearest.pip1);
        sg.applyBindings(state.nearest.bind);
        sg.draw(0, 3, 1);
        sg.endPass();

        // render from tex1 to tex2
        state.nearest.pass.attachments.colors[0] = state.nearest.render_tex2;
        state.nearest.bind.views[shd.VIEW_scr_tex] = state.nearest.input_tex1;
        sg.beginPass(state.nearest.pass);
        sg.applyPipeline(state.nearest.pip1);
        sg.applyBindings(state.nearest.bind);
        sg.draw(0, 3, 1);
        sg.endPass();
        
        // render from tex2 to screen
        state.nearest.bind.views[shd.VIEW_scr_tex] = state.nearest.input_tex2;
        sg.beginPass(.{.action = state.nearest.pass_action, .swapchain = sglue.swapchain() });
        sg.applyPipeline(state.nearest.pip2);
        sg.applyBindings(state.nearest.bind);
        sg.draw(0, 3, 1);
        simgui.render();
        sg.endPass();
    }

    sg.commit();
}

export fn cleanup() void {
    sappimgui.shutdown();
    simgui.shutdown();
    sg.shutdown();
}

pub fn main(_init: std.process.Init) void {
    state.alloc = _init.arena.allocator();
    state.io = _init.io;

    sapp.run(.{
        .init_cb = init,
        .frame_cb = frame,
        .event_cb = event,
        .cleanup_cb = cleanup,
        .sample_count = 4,
        .width = state.Width,
        .height = state.Height,
        .icon = .{ .sokol_default = true },
        .window_title = "offscreen.zig",
        .logger = .{ .func = slog.func },
    });
}