#include <stdio.h>
#include <stdlib.h>

#include <SDL3/SDL.h>
#include <SDL3/SDL_gpu.h>

#include "shader_bundle.h"
#include "shader_gpu.h"

#ifndef KAT_SOURCE_DIR
#define KAT_SOURCE_DIR "."
#endif

static void fill_checkerboard(Uint8 *pixels, int width, int height) {

    for (int y = 0; y < height; ++y) {
        for (int x = 0; x < width; ++x) {
            int checker = ((x / 8) + (y / 8)) & 1;
            Uint8 *pixel = pixels + ((y * width) + x) * 4;

            pixel[0] = checker ? 255 : 32;
            pixel[1] = checker ? 220 : 64;
            pixel[2] = checker ? 48 : 180;
            pixel[3] = 255;
        }
    }
}

static int draw_checkerboard(SDL_Window *window, SDL_GPUDevice *device,
        const t_shader_bundle *bundle,
        const t_shader_gpu_pipeline *pipeline,
        t_shader_gpu_resources *resources,
        const void *pixels, int pitch) {

    SDL_GPUCommandBuffer *command_buffer;
    SDL_GPUTexture *swapchain = NULL;
    SDL_GPUColorTargetInfo color_target;
    SDL_GPURenderPass *render_pass;
    t_shader_gpu_common_uniforms uniforms;
    t_shader_gpu_uniform_blocks uniform_blocks;
    Uint32 width = 0;
    Uint32 height = 0;

    command_buffer = SDL_AcquireGPUCommandBuffer(device);
    if (!command_buffer) {
        printf("SDL_AcquireGPUCommandBuffer failed: %s\n", SDL_GetError());
        return -1;
    }

    if (shader_gpu_resources_upload_source(command_buffer,
                resources, pixels,
                resources->source_width, resources->source_height, pitch)) {
        printf("Failed to upload checkerboard source texture: %s\n",
                resources->error);
        SDL_CancelGPUCommandBuffer(command_buffer);
        return -1;
    }

    if (!SDL_WaitAndAcquireGPUSwapchainTexture(command_buffer, window,
                &swapchain, &width, &height)) {
        printf("SDL_WaitAndAcquireGPUSwapchainTexture failed: %s\n",
                SDL_GetError());
        SDL_CancelGPUCommandBuffer(command_buffer);
        return -1;
    }

    if (!swapchain) {
        printf("No swapchain texture acquired\n");
        SDL_CancelGPUCommandBuffer(command_buffer);
        return -1;
    }

    shader_gpu_common_uniforms_fill(&uniforms,
            resources->source_width, resources->source_height,
            resources->source_width, resources->source_height,
            (int) width, (int) height, 0, 1);
    if (shader_gpu_pack_uniform_blocks(bundle, 0, &uniforms,
                &uniform_blocks)) {
        printf("shader_gpu_pack_uniform_blocks failed\n");
        SDL_CancelGPUCommandBuffer(command_buffer);
        return -1;
    }

    SDL_zero(color_target);
    color_target.texture = swapchain;
    color_target.clear_color = (SDL_FColor){ 0.02f, 0.02f, 0.02f, 1.0f };
    color_target.load_op = SDL_GPU_LOADOP_CLEAR;
    color_target.store_op = SDL_GPU_STOREOP_STORE;

    render_pass = SDL_BeginGPURenderPass(command_buffer, &color_target, 1,
            NULL);
    if (!render_pass) {
        printf("SDL_BeginGPURenderPass failed: %s\n", SDL_GetError());
        SDL_CancelGPUCommandBuffer(command_buffer);
        return -1;
    }

    if (shader_gpu_resources_draw_first_pass(command_buffer, render_pass,
                pipeline, resources, &uniform_blocks)) {
        printf("shader_gpu_resources_draw_first_pass failed\n");
        SDL_EndGPURenderPass(render_pass);
        SDL_CancelGPUCommandBuffer(command_buffer);
        return -1;
    }

    SDL_EndGPURenderPass(render_pass);

    if (!SDL_SubmitGPUCommandBuffer(command_buffer)) {
        printf("SDL_SubmitGPUCommandBuffer failed: %s\n", SDL_GetError());
        return -1;
    }

    printf("Drew shader checkerboard to %ux%u swapchain\n", width, height);
    SDL_Delay(250);
    return 0;
}

static int run_pipeline_spike(const char *bundle_dir) {

    SDL_Window *window = NULL;
    SDL_Renderer *renderer = NULL;
    SDL_PropertiesID renderer_props;
    SDL_GPUDevice *device;
    SDL_GPUShaderFormat formats;
    SDL_GPUTextureFormat target_format;
    t_shader_bundle bundle;
    t_shader_gpu_pipeline pipeline;
    t_shader_gpu_resources resources;
    Uint8 *pixels = NULL;
    int status = 1;
    const int source_width = 64;
    const int source_height = 64;

    shader_bundle_init(&bundle);
    shader_gpu_pipeline_init(&pipeline);
    shader_gpu_resources_init(&resources);

    window = SDL_CreateWindow("kat5200 SDL GPU shader pipeline spike",
            640, 480, 0);
    if (!window) {
        printf("SDL_CreateWindow failed: %s\n", SDL_GetError());
        goto done;
    }

    renderer = SDL_CreateRenderer(window, "gpu");
    if (!renderer) {
        printf("SDL_CreateRenderer(window, \"gpu\") failed: %s\n",
                SDL_GetError());
        goto done;
    }

    renderer_props = SDL_GetRendererProperties(renderer);
    device = SDL_GetPointerProperty(renderer_props,
            SDL_PROP_RENDERER_GPU_DEVICE_POINTER, NULL);
    if (!device) {
        printf("Renderer did not expose SDL_PROP_RENDERER_GPU_DEVICE_POINTER\n");
        goto done;
    }

    formats = SDL_GetGPUShaderFormats(device);
    printf("GPU driver: %s\n", SDL_GetGPUDeviceDriver(device));
    printf("Shader formats: 0x%08x\n", formats);
    target_format = SDL_GetGPUSwapchainTextureFormat(device, window);

    if (shader_bundle_load(bundle_dir, &bundle)) {
        printf("Failed to load bundle: %s\n", bundle_dir);
        goto done;
    }

    if (shader_gpu_pipeline_create(device, &bundle, target_format, &pipeline)) {
        printf("Failed to create pipeline for bundle %s: %s\n",
                bundle.id ? bundle.id : bundle_dir, pipeline.error);
        goto done;
    }

    printf("Created %d SDL GPU pipeline(s) for bundle %s using %s\n",
            pipeline.num_passes, bundle.id,
            shader_gpu_format_name(pipeline.shader_format));

    if (shader_gpu_resources_create(device, source_width, source_height, 0,
                &resources)) {
        printf("Failed to create shader resources: %s\n", resources.error);
        goto done;
    }

    pixels = malloc((size_t) source_width * source_height * 4);
    if (!pixels) {
        printf("Failed to allocate checkerboard pixels\n");
        goto done;
    }
    fill_checkerboard(pixels, source_width, source_height);

    if (draw_checkerboard(window, device, &bundle, &pipeline, &resources,
                pixels, source_width * 4))
        goto done;

    status = 0;

done:
    free(pixels);
    shader_gpu_resources_free(&resources);
    shader_gpu_pipeline_free(&pipeline);
    shader_bundle_free(&bundle);
    if (renderer)
        SDL_DestroyRenderer(renderer);
    if (window)
        SDL_DestroyWindow(window);

    return status;
}

int main(int argc, char **argv) {

    const char *bundle_dir = KAT_SOURCE_DIR
        "/resources/shaders/bundles/bilinear";
    int status;

    if (argc > 1)
        bundle_dir = argv[1];

    if (!SDL_Init(SDL_INIT_VIDEO)) {
        printf("SDL_Init failed: %s\n", SDL_GetError());
        return 1;
    }

    status = run_pipeline_spike(bundle_dir);

    SDL_Quit();
    return status ? 1 : 0;
}
