#include "renderer.h"
#include <dcimgui_impl_sdlgpu3.h>

static bool enable_debug_mode = false;

typedef struct R_Buffer_Impl {
    SDL_GPUBuffer *buffer;
    void *cpu_uniform_data;

    Uint32 num_bytes;
    R_Buffer_Usage usage;
} R_Buffer_Impl;

typedef struct R_Texture_Impl {
    SDL_GPUTexture *texture;

    Uint32 width;
    Uint32 height;

    bool multisampled;
    R_Texture_Format format;
} R_Texture_Impl;

R_Texture R_surface;

static SDL_Window    *window;
static SDL_GPUDevice *device;

static R_Texture_Impl surface;

static SDL_GPUTransferBuffer *transfer_buffer;

static SDL_GPUCommandBuffer *copy_command_buffer;
static SDL_GPUCopyPass      *copy_pass;

static SDL_GPUCommandBuffer *render_command_buffer;
static SDL_GPURenderPass    *render_pass;

static SDL_GPUSampler *bilinear_sampler;

static SDL_GPUGraphicsPipeline *shaders[R_SHADER_COUNT];

static SDL_GPUTextureFormat texture_formats[] = {
    [R_TEXTURE_FORMAT_R8_UNORM]         = SDL_GPU_TEXTUREFORMAT_R8_UNORM,
    [R_TEXTURE_FORMAT_R16_UINT]         = SDL_GPU_TEXTUREFORMAT_R16_UINT,
    [R_TEXTURE_FORMAT_RGBA8_UNORM]      = SDL_GPU_TEXTUREFORMAT_R8G8B8A8_UNORM,
    [R_TEXTURE_FORMAT_RGBA8_UNORM_SRGB] = SDL_GPU_TEXTUREFORMAT_R8G8B8A8_UNORM_SRGB,
    [R_TEXTURE_FORMAT_BGRA8_UNORM]      = SDL_GPU_TEXTUREFORMAT_B8G8R8A8_UNORM,
    [R_TEXTURE_FORMAT_BGRA8_UNORM_SRGB] = SDL_GPU_TEXTUREFORMAT_B8G8R8A8_UNORM_SRGB,
};
static_assert(SDL_arraysize(texture_formats) == R_TEXTURE_FORMAT_COUNT);

static Uint32 texel_sizes[] = {
    [R_TEXTURE_FORMAT_R8_UNORM]         = 1,
    [R_TEXTURE_FORMAT_R16_UINT]         = 2,
    [R_TEXTURE_FORMAT_RGBA8_UNORM]      = 4,
    [R_TEXTURE_FORMAT_RGBA8_UNORM_SRGB] = 4,
    [R_TEXTURE_FORMAT_BGRA8_UNORM]      = 4,
    [R_TEXTURE_FORMAT_BGRA8_UNORM_SRGB] = 4,
};
static_assert(SDL_arraysize(texel_sizes) == R_TEXTURE_FORMAT_COUNT);

static SDL_GPUTextureUsageFlags texture_usages[] = {
    [R_TEXTURE_USAGE_SAMPLED] = SDL_GPU_TEXTUREUSAGE_SAMPLER,
    [R_TEXTURE_USAGE_STORAGE] = SDL_GPU_TEXTUREUSAGE_GRAPHICS_STORAGE_READ,
    [R_TEXTURE_USAGE_TARGET]  = SDL_GPU_TEXTUREUSAGE_COLOR_TARGET,
};
static_assert(SDL_arraysize(texture_usages) == R_TEXTURE_USAGE_COUNT);

static SDL_GPUBufferUsageFlags buffer_usages[] = {
    [R_BUFFER_USAGE_UNIFORM] = 0,
    [R_BUFFER_USAGE_STORAGE] = SDL_GPU_BUFFERUSAGE_GRAPHICS_STORAGE_READ,
    [R_BUFFER_USAGE_VERTEX]  = SDL_GPU_BUFFERUSAGE_VERTEX,
    [R_BUFFER_USAGE_INDEX]   = SDL_GPU_BUFFERUSAGE_INDEX,
};
static_assert(SDL_arraysize(buffer_usages) == R_BUFFER_USAGE_COUNT);

static SDL_GPULoadOp load_ops[] = {
    [R_LOAD_OP_LOAD]  = SDL_GPU_LOADOP_LOAD,
    [R_LOAD_OP_CLEAR] = SDL_GPU_LOADOP_CLEAR,
};
static_assert(SDL_arraysize(load_ops) == R_LOAD_OP_COUNT);

static SDL_GPUStoreOp store_ops[] = {
    [R_STORE_OP_STORE]   = SDL_GPU_STOREOP_STORE,
    [R_STORE_OP_DISCARD] = SDL_GPU_STOREOP_DONT_CARE,
};
static_assert(SDL_arraysize(store_ops) == R_STORE_OP_COUNT);

R_Texture renderer_texture_create(R_Texture_Format format, R_Texture_Usage usage, bool multisampled, Uint32 width, Uint32 height, void *data, const char *debug_name) {
    SDL_PropertiesID properties = SDL_CreateProperties();
    SDL_SetStringProperty(properties, SDL_PROP_GPU_TEXTURE_CREATE_NAME_STRING, debug_name);
    
    SDL_GPUTexture *texture = SDL_CreateGPUTexture(device, &(SDL_GPUTextureCreateInfo){
        .type                 = SDL_GPU_TEXTURETYPE_2D,
        .format               = texture_formats[format],
        .usage                = texture_usages[usage],
        .width                = width,
        .height               = height,
        .layer_count_or_depth = 1,
        .num_levels           = 1,
        .sample_count         = multisampled ? SDL_GPU_SAMPLECOUNT_8 : SDL_GPU_SAMPLECOUNT_1,
        .props                = properties,
    });
    SDL_DestroyProperties(properties);
    if (!texture) return NULL;

    R_Texture_Impl *result = SDL_calloc(1, sizeof(R_Texture_Impl));
    if (!result) {
        SDL_ReleaseGPUTexture(device, texture);
        return NULL;
    }

    *result = (R_Texture_Impl){
        .texture = texture,

        .width  = width,
        .height = height,

        .multisampled = multisampled,
        .format       = format,
    };

    if (data) {
        renderer_texture_update(result, 0, 0, width, height, data, width * texel_sizes[format]);
    }

    return result;
}

void renderer_texture_update(R_Texture texture, Uint32 offset_x, Uint32 offset_y, Uint32 width, Uint32 height, void *data, Uint32 bytes_per_row) {
    void *gpu_mem = SDL_MapGPUTransferBuffer(device, transfer_buffer, true);
    SDL_memcpy(gpu_mem, data, bytes_per_row * height);
    SDL_UnmapGPUTransferBuffer(device, transfer_buffer);

    SDL_UploadToGPUTexture(copy_pass, 
        &(SDL_GPUTextureTransferInfo){
            .transfer_buffer = transfer_buffer,
            .offset          = 0,
            .pixels_per_row  = width,
            .rows_per_layer  = height,
        },
        &(SDL_GPUTextureRegion){
            .texture   = texture->texture,
            .mip_level = 0,
            .layer     = 0,

            .x = offset_x,
            .y = offset_y,
            .z = 0,

            .w = width,
            .h = height,
            .d = 1,
        },
        false
    );
}

void renderer_texture_destroy(R_Texture texture) {
    SDL_ReleaseGPUTexture(device, texture->texture);
    SDL_free(texture);
}

Uint32 renderer_texture_get_width(R_Texture texture) {
    return texture->width;
}

Uint32 renderer_texture_get_height(R_Texture texture) {
    return texture->height;
}

bool renderer_texture_get_multisampled(R_Texture texture) {
    return texture->multisampled;
}

R_Texture_Format renderer_texture_get_format(R_Texture texture) {
    return texture->format;
}

R_Buffer renderer_buffer_create(R_Buffer_Usage usage, Uint32 num_bytes, void *data, const char *debug_name) {
    R_Buffer_Impl *result = SDL_calloc(1, sizeof(R_Buffer_Impl));
    if (!result) return NULL;

    if (usage == R_BUFFER_USAGE_UNIFORM) {
        *result = (R_Buffer_Impl){
            .cpu_uniform_data = SDL_calloc(1, num_bytes),
            
            .num_bytes = num_bytes,
            .usage     = usage,
        };

        if (!result->cpu_uniform_data) {
            SDL_free(result);
            return NULL;
        }

        if (data) {
            renderer_buffer_update(result, 0, num_bytes, data);
        }
    } else {
        SDL_PropertiesID properties = SDL_CreateProperties();
        SDL_SetStringProperty(properties, SDL_PROP_GPU_BUFFER_CREATE_NAME_STRING, debug_name);
        
        SDL_GPUBuffer *buffer = SDL_CreateGPUBuffer(device, &(SDL_GPUBufferCreateInfo){
            .usage = buffer_usages[usage],
            .size  = num_bytes,
            .props = properties,
        });
        SDL_DestroyProperties(properties);
        if (!buffer)  {
            SDL_free(result);
            return NULL;
        }

        *result = (R_Buffer_Impl){
            .buffer = buffer,
            
            .num_bytes = num_bytes,
            .usage     = usage,
        };

        if (data) {
            renderer_buffer_update(result, 0, num_bytes, data);
        }
    }

    return result;
}

void renderer_buffer_update(R_Buffer buffer, Uint32 offset, Uint32 num_bytes, void *data) {
    SDL_assert(offset + num_bytes <= buffer->num_bytes);

    if (buffer->usage == R_BUFFER_USAGE_UNIFORM) {
        memcpy(buffer->cpu_uniform_data, data + offset, num_bytes);
    } else {
        void *gpu_mem = SDL_MapGPUTransferBuffer(device, transfer_buffer, true);
        SDL_memcpy(gpu_mem, data, num_bytes);
        SDL_UnmapGPUTransferBuffer(device, transfer_buffer);

        SDL_UploadToGPUBuffer(copy_pass,
            &(SDL_GPUTransferBufferLocation) {
                .transfer_buffer = transfer_buffer,
                .offset          = 0,
            }, 
            &(SDL_GPUBufferRegion){
                .buffer = buffer->buffer,
                .offset = offset,
                .size   = num_bytes,
            },
            false
        );
    }
}

void renderer_buffer_destroy(R_Buffer buffer) {
    if (buffer->buffer) SDL_ReleaseGPUBuffer(device, buffer->buffer);
    if (buffer->cpu_uniform_data) SDL_free(buffer->cpu_uniform_data);
    SDL_free(buffer);
}

Uint32 renderer_buffer_get_size(R_Buffer buffer) {
    return buffer->num_bytes;
}

R_Buffer_Usage renderer_buffer_get_usage(R_Buffer buffer){
    return buffer->usage;
}

void renderer_frame_begin() {
    SDL_assert(R_surface == NULL);

    Uint32 surface_width = 0, surface_height = 0;
    SDL_WaitAndAcquireGPUSwapchainTexture(render_command_buffer, window, &surface.texture, &surface.width, &surface.height);

    R_surface = &surface;
}

void renderer_frame_end() {
    SDL_EndGPUCopyPass(copy_pass);
    SDL_SubmitGPUCommandBuffer(copy_command_buffer);

    if (render_pass) {
        SDL_EndGPURenderPass(render_pass);
        render_pass = NULL;
    }

    SDL_GPUFence *frame_fence = SDL_SubmitGPUCommandBufferAndAcquireFence(render_command_buffer);
    SDL_WaitForGPUFences(device, true, &frame_fence, 1);

    copy_command_buffer   = SDL_AcquireGPUCommandBuffer(device);
    render_command_buffer = SDL_AcquireGPUCommandBuffer(device);

    R_surface = NULL;

    copy_pass = SDL_BeginGPUCopyPass(copy_command_buffer);
}

void renderer_frame_set_target(R_Texture target, R_Texture resolve_target, R_Load_Op load_op, R_Store_Op store_op, R_Color clear_color) {
    if (render_pass) {
        SDL_EndGPURenderPass(render_pass);
        render_pass = NULL;
    }

    render_pass = SDL_BeginGPURenderPass(render_command_buffer, 
        &(SDL_GPUColorTargetInfo){
            .texture              = target->texture,
            .mip_level            = 0,
            .layer_or_depth_plane = 0,

            .clear_color = { clear_color.r, clear_color.g, clear_color.b, clear_color.a },
            .load_op     = load_ops[load_op],
            .store_op    = resolve_target ? (store_op == R_STORE_OP_STORE ? SDL_GPU_STOREOP_RESOLVE_AND_STORE : SDL_GPU_STOREOP_RESOLVE) : store_ops[store_op],

            .resolve_texture   = resolve_target ? resolve_target->texture : NULL,
            .resolve_mip_level = 0,
            .resolve_layer     = 0,

            .cycle                 = false,
            .cycle_resolve_texture = false,
        },
        1,
        NULL
    );
}

void renderer_frame_set_shader(R_Shader shader) {
    SDL_BindGPUGraphicsPipeline(render_pass, shaders[shader]);
}

void renderer_frame_draw(R_Buffer vertex_buffer, R_Buffer index_buffer, R_Buffer instance_buffer, Uint32 num_indices, Uint32 num_instances, const R_Draw_Resources *resources) {
    SDL_assert(render_pass);


    for (int i = 0; i < resources->num_vertex_sampled_textures; i++) SDL_BindGPUVertexSamplers(render_pass, i, &(SDL_GPUTextureSamplerBinding){ .texture = resources->vertex_sampled_textures[i]->texture, .sampler = bilinear_sampler }, 1);
    for (int i = 0; i < resources->num_vertex_storage_textures; i++) SDL_BindGPUVertexStorageTextures(render_pass, i, &resources->vertex_storage_textures[i]->texture, 1);
    for (int i = 0; i < resources->num_vertex_storage_buffers;  i++) SDL_BindGPUVertexStorageBuffers(render_pass, i, &resources->vertex_storage_buffers[i]->buffer, 1);
    for (int i = 0; i < resources->num_vertex_uniform_buffers;  i++) SDL_PushGPUVertexUniformData(render_command_buffer, i, resources->vertex_uniform_buffers[i]->cpu_uniform_data, resources->vertex_uniform_buffers[i]->num_bytes);

    for (int i = 0; i < resources->num_fragment_sampled_textures; i++) SDL_BindGPUFragmentSamplers(render_pass, i, &(SDL_GPUTextureSamplerBinding){ .texture = resources->fragment_sampled_textures[i]->texture, .sampler = bilinear_sampler }, 1);
    for (int i = 0; i < resources->num_fragment_storage_textures; i++) SDL_BindGPUFragmentStorageTextures(render_pass, i, &resources->fragment_storage_textures[i]->texture, 1);
    for (int i = 0; i < resources->num_fragment_storage_buffers;  i++) SDL_BindGPUFragmentStorageBuffers(render_pass, i, &resources->fragment_storage_buffers[i]->buffer, 1);
    for (int i = 0; i < resources->num_fragment_uniform_buffers;  i++) SDL_PushGPUFragmentUniformData(render_command_buffer, i, resources->fragment_uniform_buffers[i]->cpu_uniform_data, resources->fragment_uniform_buffers[i]->num_bytes);

    if (vertex_buffer)   SDL_BindGPUVertexBuffers(render_pass, 0, &(SDL_GPUBufferBinding){ .buffer = vertex_buffer->buffer,   .offset = 0 }, 1);
    if (instance_buffer) SDL_BindGPUVertexBuffers(render_pass, 1, &(SDL_GPUBufferBinding){ .buffer = instance_buffer->buffer, .offset = 0 }, 1);

    if (index_buffer) {
        SDL_BindGPUIndexBuffer(render_pass, &(SDL_GPUBufferBinding){ .buffer = index_buffer->buffer, .offset = 0 }, SDL_GPU_INDEXELEMENTSIZE_16BIT);
        SDL_DrawGPUIndexedPrimitives(render_pass, num_indices, num_instances, 0, 0, 0);
    } else {
        SDL_DrawGPUPrimitives(render_pass, num_indices, num_instances, 0, 0);
    }
}

static R_Texture_Format R_Texture_Format_from_SDL_GPUTextureFormat(SDL_GPUTextureFormat format) {
    switch (format) {
        case SDL_GPU_TEXTUREFORMAT_R8_UNORM:            return R_TEXTURE_FORMAT_R8_UNORM;
        case SDL_GPU_TEXTUREFORMAT_R16_UINT:            return R_TEXTURE_FORMAT_R16_UINT;
        case SDL_GPU_TEXTUREFORMAT_R8G8B8A8_UNORM:      return R_TEXTURE_FORMAT_RGBA8_UNORM;
        case SDL_GPU_TEXTUREFORMAT_R8G8B8A8_UNORM_SRGB: return R_TEXTURE_FORMAT_RGBA8_UNORM_SRGB;
        case SDL_GPU_TEXTUREFORMAT_B8G8R8A8_UNORM:      return R_TEXTURE_FORMAT_BGRA8_UNORM;
        case SDL_GPU_TEXTUREFORMAT_B8G8R8A8_UNORM_SRGB: return R_TEXTURE_FORMAT_BGRA8_UNORM_SRGB;

        default: {
            SDL_assert_always(false);
            return R_TEXTURE_FORMAT_BGRA8_UNORM_SRGB;
        };
    };
}

bool renderer_init(SDL_Window *window_) {
    window = window_;

    device = SDL_CreateGPUDevice(SDL_GPU_SHADERFORMAT_SPIRV, enable_debug_mode, NULL);
    if (!device) {
        SDL_LogError(SDL_LOG_CATEGORY_GPU, "Failed to create gpu device.");
        return false;
    }

    if (!SDL_ClaimWindowForGPUDevice(device, window)) {
        SDL_LogError(SDL_LOG_CATEGORY_GPU, "Failed to claim window for gpu device.");
        SDL_DestroyGPUDevice(device);
        return false;
    }

    SDL_SetGPUSwapchainParameters(device, window, SDL_GPU_SWAPCHAINCOMPOSITION_SDR_LINEAR, SDL_GPU_PRESENTMODE_VSYNC);

    surface.format = R_Texture_Format_from_SDL_GPUTextureFormat(SDL_GetGPUSwapchainTextureFormat(device, window));

    bilinear_sampler = SDL_CreateGPUSampler(device, &(SDL_GPUSamplerCreateInfo){
        .min_filter  = SDL_GPU_FILTER_LINEAR,
        .mag_filter  = SDL_GPU_FILTER_LINEAR,
        .mipmap_mode = SDL_GPU_SAMPLERMIPMAPMODE_NEAREST,
    
        .address_mode_u = SDL_GPU_SAMPLERADDRESSMODE_CLAMP_TO_EDGE,
        .address_mode_v = SDL_GPU_SAMPLERADDRESSMODE_CLAMP_TO_EDGE,
        .address_mode_w = SDL_GPU_SAMPLERADDRESSMODE_CLAMP_TO_EDGE,
    
        .mip_lod_bias   = 0.0f,
        .max_anisotropy = 1.0f,
        .compare_op     = SDL_GPU_COMPAREOP_INVALID,

        .min_lod = 0.0f,
        .max_lod = 0.0f,

        .enable_anisotropy = false,
        .enable_compare    = false,
    });

    // Shaders
    { // Basic
        const Uint8 shader_source[] = {
            #embed "shaders/spirv/basic.spv"
        };

        SDL_GPUShader *vertex_shader = SDL_CreateGPUShader(device, &(SDL_GPUShaderCreateInfo){
            .code_size  = sizeof(shader_source),
            .code       = shader_source,
            .entrypoint = "main_vertex",

            .format = SDL_GPU_SHADERFORMAT_SPIRV,
            .stage  = SDL_GPU_SHADERSTAGE_VERTEX,

            .num_samplers         = 0,
            .num_storage_textures = 0,
            .num_storage_buffers  = 0,
            .num_uniform_buffers  = 1,
        });

        SDL_GPUShader *fragment_shader = SDL_CreateGPUShader(device, &(SDL_GPUShaderCreateInfo){
            .code_size  = sizeof(shader_source),
            .code       = shader_source,
            .entrypoint = "main_fragment",

            .format = SDL_GPU_SHADERFORMAT_SPIRV,
            .stage  = SDL_GPU_SHADERSTAGE_FRAGMENT,

            .num_samplers         = 1,
            .num_storage_textures = 0,
            .num_storage_buffers  = 0,
            .num_uniform_buffers  = 1,
        });

        SDL_PropertiesID properties = SDL_CreateProperties();
        SDL_SetStringProperty(properties, SDL_PROP_GPU_BUFFER_CREATE_NAME_STRING, "R_SHADER basic");

        shaders[R_SHADER_BASIC] = SDL_CreateGPUGraphicsPipeline(device, &(SDL_GPUGraphicsPipelineCreateInfo){
            .vertex_shader   = vertex_shader,
            .fragment_shader = fragment_shader,

            .vertex_input_state = {
                .vertex_buffer_descriptions = (SDL_GPUVertexBufferDescription[]){
                    {
                        .slot       = 0,
                        .pitch      = 20,
                        .input_rate = SDL_GPU_VERTEXINPUTRATE_VERTEX,

                        .instance_step_rate = 0,
                    },
                    {
                        .slot       = 1,
                        .pitch      = 8,
                        .input_rate = SDL_GPU_VERTEXINPUTRATE_INSTANCE,

                        .instance_step_rate = 0,
                    },
                },
                .num_vertex_buffers = 2,

                .vertex_attributes = (SDL_GPUVertexAttribute[]){
                    {
                        .location    = 0,
                        .buffer_slot = 0,
                        .format      = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT3,
                        .offset      = 0,
                    },
                    {
                        .location    = 1,
                        .buffer_slot = 0,
                        .format      = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT2,
                        .offset      = 12,
                    },
                    {
                        .location    = 2,
                        .buffer_slot = 1,
                        .format      = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT2,
                        .offset      = 0,
                    },
                },
                .num_vertex_attributes = 3,
            },

            .primitive_type = SDL_GPU_PRIMITIVETYPE_TRIANGLELIST,

            .rasterizer_state = {
                .fill_mode  = SDL_GPU_FILLMODE_FILL,
                .cull_mode  = SDL_GPU_CULLMODE_BACK,
                .front_face = SDL_GPU_FRONTFACE_COUNTER_CLOCKWISE,

                .depth_bias_constant_factor = 0.0f,
                .depth_bias_clamp           = 0.0f,
                .depth_bias_slope_factor    = 0.0f,

                .enable_depth_bias = false,
                .enable_depth_clip = false,
            },

            .multisample_state = {
                .sample_count = SDL_GPU_SAMPLECOUNT_8,
                .sample_mask  = 0,
                .enable_mask  = 0,

                .enable_alpha_to_coverage = false,
            },

            .depth_stencil_state = {
                .compare_op = SDL_GPU_COMPAREOP_GREATER_OR_EQUAL,

                .back_stencil_state  = { .fail_op = SDL_GPU_STENCILOP_INVALID, .pass_op = SDL_GPU_STENCILOP_INVALID, .depth_fail_op = SDL_GPU_STENCILOP_INVALID, .compare_op = SDL_GPU_COMPAREOP_INVALID, },
                .front_stencil_state = { .fail_op = SDL_GPU_STENCILOP_INVALID, .pass_op = SDL_GPU_STENCILOP_INVALID, .depth_fail_op = SDL_GPU_STENCILOP_INVALID, .compare_op = SDL_GPU_COMPAREOP_INVALID, },

                .compare_mask = 0,
                .write_mask   = 0,
        
                .enable_depth_test   = false,
                .enable_depth_write  = false,
                .enable_stencil_test = false,
            },

            .target_info = {
                .color_target_descriptions = (SDL_GPUColorTargetDescription[]){
                    {
                        .format      = texture_formats[surface.format],
                        .blend_state = {
                            .src_color_blendfactor = SDL_GPU_BLENDFACTOR_SRC_ALPHA,
                            .dst_color_blendfactor = SDL_GPU_BLENDFACTOR_ONE_MINUS_SRC_ALPHA,
                            .color_blend_op        = SDL_GPU_BLENDOP_ADD,

                            .src_alpha_blendfactor = SDL_GPU_BLENDFACTOR_SRC_ALPHA,
                            .dst_alpha_blendfactor = SDL_GPU_BLENDFACTOR_ONE_MINUS_SRC_ALPHA,
                            .alpha_blend_op        = SDL_GPU_BLENDOP_ADD,

                            .color_write_mask = 0,

                            .enable_blend            = true,
                            .enable_color_write_mask = false,
                        },
                    }
                },

                .num_color_targets = 1,

                .depth_stencil_format     = SDL_GPU_TEXTUREFORMAT_INVALID,
                .has_depth_stencil_target = false,
            },
        });
        SDL_DestroyProperties(properties);
        SDL_ReleaseGPUShader(device, vertex_shader);
        SDL_ReleaseGPUShader(device, fragment_shader);
    }

    { // World
        const Uint8 shader_source[] = {
            #embed "shaders/spirv/world.spv"
        };

        SDL_GPUShader *vertex_shader = SDL_CreateGPUShader(device, &(SDL_GPUShaderCreateInfo){
            .code_size  = sizeof(shader_source),
            .code       = shader_source,
            .entrypoint = "main_vertex",

            .format = SDL_GPU_SHADERFORMAT_SPIRV,
            .stage  = SDL_GPU_SHADERSTAGE_VERTEX,

            .num_samplers         = 0,
            .num_storage_textures = 1,
            .num_storage_buffers  = 0,
            .num_uniform_buffers  = 2,
        });

        SDL_GPUShader *fragment_shader = SDL_CreateGPUShader(device, &(SDL_GPUShaderCreateInfo){
            .code_size  = sizeof(shader_source),
            .code       = shader_source,
            .entrypoint = "main_fragment",

            .format = SDL_GPU_SHADERFORMAT_SPIRV,
            .stage  = SDL_GPU_SHADERSTAGE_FRAGMENT,

            .num_samplers         = 1,
            .num_storage_textures = 0,
            .num_storage_buffers  = 1,
            .num_uniform_buffers  = 1,
        });

        SDL_PropertiesID properties = SDL_CreateProperties();
        SDL_SetStringProperty(properties, SDL_PROP_GPU_BUFFER_CREATE_NAME_STRING, "R_SHADER world");

        shaders[R_SHADER_WORLD] = SDL_CreateGPUGraphicsPipeline(device, &(SDL_GPUGraphicsPipelineCreateInfo){
            .vertex_shader   = vertex_shader,
            .fragment_shader = fragment_shader,

            .vertex_input_state = {
                .vertex_buffer_descriptions = (SDL_GPUVertexBufferDescription[]){
                    {
                        .slot       = 0,
                        .pitch      = 20,
                        .input_rate = SDL_GPU_VERTEXINPUTRATE_VERTEX,

                        .instance_step_rate = 0,
                    },
                },
                .num_vertex_buffers = 1,

                .vertex_attributes = (SDL_GPUVertexAttribute[]){
                    {
                        .location    = 0,
                        .buffer_slot = 0,
                        .format      = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT3,
                        .offset      = 0,
                    },
                },
                .num_vertex_attributes = 1,
            },

            .primitive_type = SDL_GPU_PRIMITIVETYPE_TRIANGLELIST,

            .rasterizer_state = {
                .fill_mode  = SDL_GPU_FILLMODE_FILL,
                .cull_mode  = SDL_GPU_CULLMODE_BACK,
                .front_face = SDL_GPU_FRONTFACE_COUNTER_CLOCKWISE,

                .depth_bias_constant_factor = 0.0f,
                .depth_bias_clamp           = 0.0f,
                .depth_bias_slope_factor    = 0.0f,

                .enable_depth_bias = false,
                .enable_depth_clip = false,
            },

            .multisample_state = {
                .sample_count = SDL_GPU_SAMPLECOUNT_8,
                .sample_mask  = 0,
                .enable_mask  = 0,

                .enable_alpha_to_coverage = false,
            },

            .depth_stencil_state = {
                .compare_op = SDL_GPU_COMPAREOP_GREATER_OR_EQUAL,

                .back_stencil_state  = { .fail_op = SDL_GPU_STENCILOP_INVALID, .pass_op = SDL_GPU_STENCILOP_INVALID, .depth_fail_op = SDL_GPU_STENCILOP_INVALID, .compare_op = SDL_GPU_COMPAREOP_INVALID, },
                .front_stencil_state = { .fail_op = SDL_GPU_STENCILOP_INVALID, .pass_op = SDL_GPU_STENCILOP_INVALID, .depth_fail_op = SDL_GPU_STENCILOP_INVALID, .compare_op = SDL_GPU_COMPAREOP_INVALID, },

                .compare_mask = 0,
                .write_mask   = 0,
        
                .enable_depth_test   = false,
                .enable_depth_write  = false,
                .enable_stencil_test = false,
            },

            .target_info = {
                .color_target_descriptions = (SDL_GPUColorTargetDescription[]){
                    {
                        .format      = texture_formats[surface.format],
                        .blend_state = {
                            .src_color_blendfactor = SDL_GPU_BLENDFACTOR_SRC_ALPHA,
                            .dst_color_blendfactor = SDL_GPU_BLENDFACTOR_ONE_MINUS_SRC_ALPHA,
                            .color_blend_op        = SDL_GPU_BLENDOP_ADD,

                            .src_alpha_blendfactor = SDL_GPU_BLENDFACTOR_SRC_ALPHA,
                            .dst_alpha_blendfactor = SDL_GPU_BLENDFACTOR_ONE_MINUS_SRC_ALPHA,
                            .alpha_blend_op        = SDL_GPU_BLENDOP_ADD,

                            .color_write_mask = 0,

                            .enable_blend            = true,
                            .enable_color_write_mask = false,
                        },
                    }
                },

                .num_color_targets = 1,

                .depth_stencil_format     = SDL_GPU_TEXTUREFORMAT_INVALID,
                .has_depth_stencil_target = false,
            },
        });
        SDL_DestroyProperties(properties);
        SDL_ReleaseGPUShader(device, vertex_shader);
        SDL_ReleaseGPUShader(device, fragment_shader);
    }

    { // Grid
        const Uint8 shader_source[] = {
            #embed "shaders/spirv/grid.spv"
        };

        SDL_GPUShader *vertex_shader = SDL_CreateGPUShader(device, &(SDL_GPUShaderCreateInfo){
            .code_size  = sizeof(shader_source),
            .code       = shader_source,
            .entrypoint = "main_vertex",

            .format = SDL_GPU_SHADERFORMAT_SPIRV,
            .stage  = SDL_GPU_SHADERSTAGE_VERTEX,

            .num_samplers         = 0,
            .num_storage_textures = 0,
            .num_storage_buffers  = 0,
            .num_uniform_buffers  = 2,
        });

        SDL_GPUShader *fragment_shader = SDL_CreateGPUShader(device, &(SDL_GPUShaderCreateInfo){
            .code_size  = sizeof(shader_source),
            .code       = shader_source,
            .entrypoint = "main_fragment",

            .format = SDL_GPU_SHADERFORMAT_SPIRV,
            .stage  = SDL_GPU_SHADERSTAGE_FRAGMENT,

            .num_samplers         = 0,
            .num_storage_textures = 0,
            .num_storage_buffers  = 0,
            .num_uniform_buffers  = 0,
        });

        SDL_PropertiesID properties = SDL_CreateProperties();
        SDL_SetStringProperty(properties, SDL_PROP_GPU_BUFFER_CREATE_NAME_STRING, "R_SHADER grid");

        shaders[R_SHADER_GRID] = SDL_CreateGPUGraphicsPipeline(device, &(SDL_GPUGraphicsPipelineCreateInfo){
            .vertex_shader   = vertex_shader,
            .fragment_shader = fragment_shader,

            .vertex_input_state = {
                .vertex_buffer_descriptions = (SDL_GPUVertexBufferDescription[]){},
                .num_vertex_buffers = 0,

                .vertex_attributes = (SDL_GPUVertexAttribute[]){},
                .num_vertex_attributes = 0,
            },

            .primitive_type = SDL_GPU_PRIMITIVETYPE_LINELIST,

            .rasterizer_state = {
                .fill_mode  = SDL_GPU_FILLMODE_FILL,
                .cull_mode  = SDL_GPU_CULLMODE_BACK,
                .front_face = SDL_GPU_FRONTFACE_COUNTER_CLOCKWISE,

                .depth_bias_constant_factor = 0.0f,
                .depth_bias_clamp           = 0.0f,
                .depth_bias_slope_factor    = 0.0f,

                .enable_depth_bias = false,
                .enable_depth_clip = false,
            },

            .multisample_state = {
                .sample_count = SDL_GPU_SAMPLECOUNT_8,
                .sample_mask  = 0,
                .enable_mask  = 0,

                .enable_alpha_to_coverage = false,
            },

            .depth_stencil_state = {
                .compare_op = SDL_GPU_COMPAREOP_GREATER_OR_EQUAL,

                .back_stencil_state  = { .fail_op = SDL_GPU_STENCILOP_INVALID, .pass_op = SDL_GPU_STENCILOP_INVALID, .depth_fail_op = SDL_GPU_STENCILOP_INVALID, .compare_op = SDL_GPU_COMPAREOP_INVALID, },
                .front_stencil_state = { .fail_op = SDL_GPU_STENCILOP_INVALID, .pass_op = SDL_GPU_STENCILOP_INVALID, .depth_fail_op = SDL_GPU_STENCILOP_INVALID, .compare_op = SDL_GPU_COMPAREOP_INVALID, },

                .compare_mask = 0,
                .write_mask   = 0,
        
                .enable_depth_test   = false,
                .enable_depth_write  = false,
                .enable_stencil_test = false,
            },

            .target_info = {
                .color_target_descriptions = (SDL_GPUColorTargetDescription[]){
                    {
                        .format      = texture_formats[surface.format],
                        .blend_state = {
                            .src_color_blendfactor = SDL_GPU_BLENDFACTOR_SRC_ALPHA,
                            .dst_color_blendfactor = SDL_GPU_BLENDFACTOR_ONE_MINUS_SRC_ALPHA,
                            .color_blend_op        = SDL_GPU_BLENDOP_ADD,

                            .src_alpha_blendfactor = SDL_GPU_BLENDFACTOR_SRC_ALPHA,
                            .dst_alpha_blendfactor = SDL_GPU_BLENDFACTOR_ONE_MINUS_SRC_ALPHA,
                            .alpha_blend_op        = SDL_GPU_BLENDOP_ADD,

                            .color_write_mask = 0,

                            .enable_blend            = true,
                            .enable_color_write_mask = false,
                        },
                    }
                },

                .num_color_targets = 1,

                .depth_stencil_format     = SDL_GPU_TEXTUREFORMAT_INVALID,
                .has_depth_stencil_target = false,
            },
        });
        SDL_DestroyProperties(properties);
        SDL_ReleaseGPUShader(device, vertex_shader);
        SDL_ReleaseGPUShader(device, fragment_shader);
    }

    for (int i = 0; i < R_SHADER_COUNT; i++) {
        SDL_assert(shaders[i]);
    }

    transfer_buffer = SDL_CreateGPUTransferBuffer(device, &(SDL_GPUTransferBufferCreateInfo){
        .usage = SDL_GPU_TRANSFERBUFFERUSAGE_UPLOAD,
        .size  = 64 * 1024,
    });

    copy_command_buffer   = SDL_AcquireGPUCommandBuffer(device);
    render_command_buffer = SDL_AcquireGPUCommandBuffer(device);

    copy_pass = SDL_BeginGPUCopyPass(copy_command_buffer);

    return true;
}

void ImGui_ImplRenderer_Init() {
    cImGui_ImplSDLGPU3_Init(&(ImGui_ImplSDLGPU3_InitInfo){
        .Device               = device,
        .ColorTargetFormat    = texture_formats[surface.format],
        .MSAASamples          = SDL_GPU_SAMPLECOUNT_8,
        .SwapchainComposition = SDL_GPU_SWAPCHAINCOMPOSITION_SDR,
        .PresentMode          = SDL_GPU_PRESENTMODE_VSYNC,
    });
}

void ImGui_ImplRenderer_NewFrame() {
    cImGui_ImplSDLGPU3_NewFrame();
}

void ImGui_ImplRenderer_Shutdown() {
    cImGui_ImplSDLGPU3_Shutdown();
}

void ImGui_ImplRenderer_RenderDrawData(void* draw_data) {
    SDL_EndGPUCopyPass(copy_pass);
    cImGui_ImplSDLGPU3_PrepareDrawData(draw_data, copy_command_buffer);
    copy_pass = SDL_BeginGPUCopyPass(copy_command_buffer);

    cImGui_ImplSDLGPU3_RenderDrawData(draw_data, render_command_buffer, render_pass);
}

Uint64 ImGui_ImplRenderer_GetTextureID(R_Texture texture) {
    return (Uint64)texture->texture;
}