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

SDL_GPUShader *loadShader(SDL_GPUDevice *device,
    const char *filename,
    Uint32 samplerCount,
    Uint32 uniformBufferCount,
    Uint32 storageBufferCount,
    Uint32 storageTextureCount) {
        SDL_GPUShaderStage stage;
        if(SDL_strstr(filename, ".vert"))
            stage = SDL_GPU_SHADERSTAGE_VERTEX;
        else if(SDL_strstr(filename, ".frag"))
            stage = SDL_GPU_SHADERSTAGE_FRAGMENT;
        else {
            SDL_Log("Invalid shader stage!");
            return NULL;
        }

        SDL_GPUShaderFormat backendFmt = SDL_GetGPUShaderFormats(device);

        SDL_GPUShaderFormat format = SDL_GPU_SHADERFORMAT_INVALID;
        const char *entryPoint;

        if(backendFmt & SDL_GPU_SHADERFORMAT_SPIRV) {
            format = SDL_GPU_SHADERFORMAT_SPIRV;
            entryPoint = "main";
        } else {
            SDL_Log("Unsupported shader format!");
            return NULL;
        }

        size_t fSize = 0;
        void *byteCode = SDL_LoadFile(filename, &fSize);

        if(byteCode == NULL) {
            SDL_Log("Could not load shader bytecode from file: %s", filename);
            return NULL;
        }

        SDL_GPUShaderCreateInfo shaderInfo = {
            .code = (Uint8*)byteCode,
            .code_size = fSize,
            .entrypoint = entryPoint,
            .format = format,
            .stage = stage,
            .num_samplers = samplerCount,
            .num_uniform_buffers = uniformBufferCount,
            .num_storage_buffers = storageBufferCount,
            .num_storage_textures = storageTextureCount
        };
        SDL_GPUShader *shader = SDL_CreateGPUShader(device, &shaderInfo);
        if(shader == NULL) {
            SDL_Log("Could not create shader!");
            SDL_free(byteCode);
            return NULL;
        }

        SDL_free(byteCode);
        return shader;
}

int main(int argc, char **argv) {
    SDL_Window *window = NULL;
    SDL_GPUDevice *GPUDevice = NULL;
    SDL_GPUTexture *swapchainTexture = NULL;
    SDL_GPUCommandBuffer *commandBuf = NULL;
    SDL_GPURenderPass *renderPass = NULL;

    SDL_GPUShader *vertexShader = NULL;
    SDL_GPUShader *fragShader = NULL;

    SDL_GPUGraphicsPipeline *pipeline = NULL;

    bool done = false;

    if(!SDL_Init(SDL_INIT_VIDEO)) {
        SDL_Log("Unable to initialize SDL3: %s", SDL_GetError());
        return 1;
    }

    // Create the GPU device
    GPUDevice = SDL_CreateGPUDevice(SDL_GPU_SHADERFORMAT_SPIRV, true, "vulkan");

    const char* device_driver = SDL_GetGPUDeviceDriver(GPUDevice); 
    SDL_Log("Created GPU device with driver: %s\n", device_driver);

    if(GPUDevice == NULL) {
        SDL_Log("Create GPU failed");
        return 1;
    }

    // Create the window
    window = SDL_CreateWindow("SDL3 GPU", 640, 480, SDL_WINDOW_RESIZABLE);
    SDL_SetWindowPosition(window, SDL_WINDOWPOS_CENTERED, SDL_WINDOWPOS_CENTERED);
    SDL_SetWindowMinimumSize(window, 640, 480);
    SDL_SetWindowMaximumSize(window, 1280, 720);

    if (window == NULL) {
        SDL_Log("Could not create window: %s\n", SDL_GetError());
        return 1;
    }

    // Tie the window to our GPU
    if(!SDL_ClaimWindowForGPUDevice(GPUDevice, window)) {
        SDL_Log("Claiming SDL3 window failed");
        return 1;
    }

    // Shader time
    vertexShader = loadShader(GPUDevice, "triangle.vert.spv", 0, 0, 0, 0);
    if(vertexShader == NULL)
    {
        SDL_Log("Could not create vertex shader!");
        return 1;
    }

    fragShader = loadShader(GPUDevice, "color.frag.spv", 0, 0, 0, 0);
    if(fragShader == NULL)
    {
        SDL_Log("Could not create fragment shader!");
        return 1;
    }

    // Creating the pipelines
    SDL_GPUGraphicsPipelineCreateInfo pipelineCreateInfo = {
        .target_info = {
            .num_color_targets = 1,
            .color_target_descriptions = (SDL_GPUColorTargetDescription[]){{
                .format = SDL_GetGPUSwapchainTextureFormat(GPUDevice, window)
            }},
        },
        .primitive_type = SDL_GPU_PRIMITIVETYPE_TRIANGLELIST,
        .vertex_shader = vertexShader,
        .fragment_shader = fragShader,
    };

    pipelineCreateInfo.rasterizer_state.fill_mode = SDL_GPU_FILLMODE_FILL;
    pipeline = SDL_CreateGPUGraphicsPipeline(GPUDevice, &pipelineCreateInfo);
    if(pipeline == NULL) {
        SDL_Log("Could not create pipeline");
        return 1;
    }

    // Free up shaders that are now bound to pipeline
    SDL_ReleaseGPUShader(GPUDevice, vertexShader);
    SDL_ReleaseGPUShader(GPUDevice, fragShader);
    //SDL_SetGPUSwapchainParameters(GPUDevice, window, SDL_GPU_SWAPCHAINCOMPOSITION_SDR, SDL_GPU_PRESENTMODE_VSYNC);

    // Now to loop
    while (!done) {
        SDL_Event event;

        while (SDL_PollEvent(&event)) {
            if (event.type == SDL_EVENT_QUIT) {
                done = true;
            }
        }

        // Get our command buffer
        commandBuf = SDL_AcquireGPUCommandBuffer(GPUDevice);
        if(commandBuf == NULL) {
            SDL_Log("Could not get command buf!");
            return 1;
        }

        // Get our swapchain
        if(!SDL_WaitAndAcquireGPUSwapchainTexture(commandBuf, window, &swapchainTexture, NULL, NULL)) {
            SDL_Log("Could not create swapchain!");
            return 1;
        }


        // Handle the swapchain
        if(swapchainTexture != NULL) {
            // Clear (background) color
            SDL_GPUColorTargetInfo colorInfo = {
                .texture = swapchainTexture,
                .clear_color = (SDL_FColor){0.25f, 0.5f, 0.1f, 1.0f},
                .load_op = SDL_GPU_LOADOP_CLEAR,
                .store_op = SDL_GPU_STOREOP_STORE
            };

            // Now for the render pass
            renderPass = SDL_BeginGPURenderPass(commandBuf, &colorInfo, 1, NULL);
            SDL_BindGPUGraphicsPipeline(renderPass, pipeline);
            SDL_DrawGPUPrimitives(renderPass, 3, 1, 0, 0);
            SDL_EndGPURenderPass(renderPass);
        } else {
            SDL_Log("Swapchain texture just does not exist! %s", SDL_GetError());
            return 1;
        }

        SDL_SubmitGPUCommandBuffer(commandBuf);
    }

    // Cleanup
    SDL_ReleaseGPUGraphicsPipeline(GPUDevice, pipeline);

    SDL_ReleaseWindowFromGPUDevice(GPUDevice, window);
    SDL_DestroyWindow(window);
    SDL_DestroyGPUDevice(GPUDevice);

    SDL_Quit();
    return 0;
}