#ifndef CONTEXT_H
#define CONTEXT_H

static SDL_GPUGraphicsPipeline* RenderPipeline;
static SDL_GPUSampler* Sampler;
static SDL_GPUTexture* Texture;
static SDL_GPUTransferBuffer* SpriteDataTransferBuffer;
static SDL_GPUBuffer* SpriteDataBuffer;
static const Uint32 SPRITE_COUNT = 10576;
static const char* BasePath = NULL;

typedef struct SpriteInstance
{
    float x, y, z;
    float rotation;
    float w, h, padding_a, padding_b;
    float tex_u, tex_v, tex_w, tex_h;
    float r, g, b, a;
} SpriteInstance;

typedef struct Matrix4x4
{
    float m11, m12, m13, m14;
    float m21, m22, m23, m24;
    float m31, m32, m33, m34;
    float m41, m42, m43, m44;
} Matrix4x4;

Matrix4x4 Matrix4x4_CreateOrthographicOffCenter(
    float left,
    float right,
    float bottom,
    float top,
    float zNearPlane,
    float zFarPlane
) {
    return (Matrix4x4) {
        2.0f / (right - left), 0, 0, 0,
        0, 2.0f / (top - bottom), 0, 0,
        0, 0, 1.0f / (zNearPlane - zFarPlane), 0,
        (left + right) / (left - right), (top + bottom) / (bottom - top), zNearPlane / (zNearPlane - zFarPlane), 1
    };
}

SDL_GPUShader* LoadShader(
    SDL_GPUDevice* device,
    const char* shaderFilename,
    Uint32 samplerCount,
    Uint32 uniformBufferCount,
    Uint32 storageBufferCount,
    Uint32 storageTextureCount
) {
    // Auto-detect the shader stage from the file name for convenience
    SDL_GPUShaderStage stage;
    if (SDL_strstr(shaderFilename, ".vert"))
    {
        stage = SDL_GPU_SHADERSTAGE_VERTEX;
    }
    else if (SDL_strstr(shaderFilename, ".frag"))
    {
        stage = SDL_GPU_SHADERSTAGE_FRAGMENT;
    }
    else
    {
        SDL_Log("Invalid shader stage!");
        return NULL;
    }

    char fullPath[256];
    SDL_GPUShaderFormat backendFormats = SDL_GetGPUShaderFormats(device);
    SDL_GPUShaderFormat format = SDL_GPU_SHADERFORMAT_INVALID;
    const char *entrypoint;

    if (backendFormats & SDL_GPU_SHADERFORMAT_SPIRV) {
        SDL_snprintf(fullPath, sizeof(fullPath), "%s/%s.spv", BasePath, shaderFilename);
        format = SDL_GPU_SHADERFORMAT_SPIRV;
        entrypoint = "main";
    } else if (backendFormats & SDL_GPU_SHADERFORMAT_MSL) {
        SDL_snprintf(fullPath, sizeof(fullPath), "%s/%s.msl", BasePath, shaderFilename);
        format = SDL_GPU_SHADERFORMAT_MSL;
        entrypoint = "main0";
    } else if (backendFormats & SDL_GPU_SHADERFORMAT_DXIL) {
        SDL_snprintf(fullPath, sizeof(fullPath), "%s/%s.dxil", BasePath, shaderFilename);
        format = SDL_GPU_SHADERFORMAT_DXIL;
        entrypoint = "main";
    } else {
        SDL_Log("%s", "Unrecognized backend shader format!");
        return NULL;
    }

    size_t codeSize;
    void* code = SDL_LoadFile(fullPath, &codeSize);
    if (code == NULL)
    {
        SDL_Log("Failed to load shader from disk! %s", fullPath);
        return NULL;
    }

    SDL_GPUShaderCreateInfo shaderInfo = {
        .code =  (Uint8*)code,
        .code_size = codeSize,
        .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("Failed to create shader!");
        SDL_free(code);
        return NULL;
    }

    SDL_free(code);
    return shader;
}

SDL_Surface* LoadImage(const char* imageFilename, int desiredChannels)
{
    char fullPath[256];
    SDL_Surface *result;
    SDL_PixelFormat format;

    SDL_snprintf(fullPath, sizeof(fullPath), "%s/%s", BasePath, imageFilename);

    result = IMG_Load(fullPath);
    if (result == NULL)
    {
        SDL_Log("Failed to load BMP: %s", SDL_GetError());
        return NULL;
    }

    if (desiredChannels == 4)
    {
        format = SDL_PIXELFORMAT_ABGR8888;
    }
    else
    {
        SDL_assert(!"Unexpected desiredChannels");
        SDL_DestroySurface(result);
        return NULL;
    }
    if (result->format != format)
    {
        SDL_Surface *next = SDL_ConvertSurface(result, format);
        SDL_DestroySurface(result);
        result = next;
    }

    return result;
}


static int render(SDL_GPUDevice *device, TileMap &tilemap, Tileset &tileset, SDL_Window *window) {
    Matrix4x4 cameraMatrix = Matrix4x4_CreateOrthographicOffCenter(
                                                                   0,
                                                                   640,
                                                                   480,
                                                                   0,
                                                                   0,
                                                                   -1
                                                                   );
    
    SDL_GPUCommandBuffer* cmdBuf = SDL_AcquireGPUCommandBuffer(device);
    if (cmdBuf == NULL)
    {
        SDL_Log("AcquireGPUCommandBuffer failed: %s", SDL_GetError());
        return -1;
    }
    
    SDL_GPUTexture* swapchainTexture;
    if (!SDL_WaitAndAcquireGPUSwapchainTexture(cmdBuf, window, &swapchainTexture, NULL, NULL)) {
        SDL_Log("WaitAndAcquireGPUSwapchainTexture failed: %s", SDL_GetError());
        return -1;
    }
    
    if (swapchainTexture != NULL)
    {
        // Build sprite instance transfer
        SpriteInstance* dataPtr = (SpriteInstance*)SDL_MapGPUTransferBuffer(
                                                                            device,
                                                                            SpriteDataTransferBuffer,
                                                                            true
                                                                            );
        
        int numCols = 100;
        int numRows = 100;
        
        for (int j = 0; j < numRows; j++) { // Iterate over rows first
            for (int i = 0; i < numCols; i++) { // Then iterate over columns
                int index = j * numCols + i;  // Correct row-major indexing
                int n = tilemap.layers[0][index];  // Tile index in tileset
                
                // Position in world space
                dataPtr[index].x = i * tileset.tileWidth;
                dataPtr[index].y = j * tileset.tileHeight;
                dataPtr[index].z = 0;
                dataPtr[index].rotation = 0;
                
                // Tile dimensions
                dataPtr[index].w = tileset.tileWidth;
                dataPtr[index].h = tileset.tileHeight;
                
                // Get tile rect from tileset
                float tex_x = tileset.tiles[n-1].rect.x;
                float tex_y = tileset.tiles[n-1].rect.y;
                float tex_w = tileset.tileWidth;
                float tex_h = tileset.tileHeight;
                float tex_w_max = tex_w;
                float tex_h_max = tex_h;
                
                // Convert to normalized UV coordinates
                dataPtr[index].tex_u = (tex_x / tileset.imageWidth);
                dataPtr[index].tex_v = (tex_y / tileset.imageHeight);
                dataPtr[index].tex_w = (tex_w_max / tileset.imageWidth);
                dataPtr[index].tex_h = (tex_h_max / tileset.imageHeight);
                
                // Default colors (white)
                dataPtr[index].r = 1.0f;
                dataPtr[index].g = 1.0f;
                dataPtr[index].b = 1.0f;
                dataPtr[index].a = 1.0f;
            }
        }
        
        
        SDL_UnmapGPUTransferBuffer(device, SpriteDataTransferBuffer);
        
        SDL_GPUTransferBufferLocation buffer_location {
            .transfer_buffer = SpriteDataTransferBuffer,
            .offset = 0
        };
        
        SDL_GPUBufferRegion buffer_region {
            .buffer = SpriteDataBuffer,
            .offset = 0,
            .size = SPRITE_COUNT * sizeof(SpriteInstance)
        };
        
        // Upload instance data
        SDL_GPUCopyPass* copyPass = SDL_BeginGPUCopyPass(cmdBuf);
        SDL_UploadToGPUBuffer(
                              copyPass,
                              &buffer_location,
                              &buffer_region,
                              true
                              );
        SDL_EndGPUCopyPass(copyPass);
        
        SDL_GPUColorTargetInfo target_info{
            .texture = swapchainTexture,
            .cycle = false,
            .load_op = SDL_GPU_LOADOP_CLEAR,
            .store_op = SDL_GPU_STOREOP_STORE,
            .clear_color = { 0, 0, 0, 1 }
        };
        
        // Render sprites
        SDL_GPURenderPass* renderPass = SDL_BeginGPURenderPass(
                                                               cmdBuf,
                                                               &target_info,
                                                               1,
                                                               NULL
                                                               );
        
        SDL_BindGPUGraphicsPipeline(renderPass, RenderPipeline);
        SDL_BindGPUVertexStorageBuffers(
                                        renderPass,
                                        0,
                                        &SpriteDataBuffer,
                                        1
                                        );
        SDL_GPUTextureSamplerBinding sampler_binding{
            .texture = Texture,
            .sampler = Sampler
        };
        SDL_BindGPUFragmentSamplers(
                                    renderPass,
                                    0,
                                    &sampler_binding,
                                    1
                                    );
        SDL_PushGPUVertexUniformData(
                                     cmdBuf,
                                     0,
                                     &cameraMatrix,
                                     sizeof(Matrix4x4)
                                     );
        SDL_DrawGPUPrimitives(
                              renderPass,
                              SPRITE_COUNT * 6,
                              1,
                              0,
                              0
                              );
        
        SDL_EndGPURenderPass(renderPass);
    }
    
    SDL_SubmitGPUCommandBuffer(cmdBuf);
    
    
    // Present the frame
    SDL_GL_SwapWindow(window);
}

#endif
