// IAEngine: 2D Game Engine by IA
// Copyright (C) 2025 IASoft (PVT) LTD (oss@iasoft.dev)
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

#include <IAEngine/Engine.hpp>
#include <Renderer/GPUResourceManager.hpp>
#include <Renderer/Renderer.hpp>

namespace ia::iae
{
    SDL_GPUSampler *g_linearClampSampler{};
    SDL_GPUSampler *g_linearRepeatSampler{};

    VOID GPUResourceManager::Initialize()
    {
        { // Create Samplers
            SDL_GPUSamplerCreateInfo createInfo{.min_filter = SDL_GPU_FILTER_NEAREST,
                                                .mag_filter = SDL_GPU_FILTER_NEAREST,
                                                .mipmap_mode = SDL_GPU_SAMPLERMIPMAPMODE_LINEAR,
                                                .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,
                                                .enable_anisotropy = false};

            g_linearClampSampler = SDL_CreateGPUSampler(Renderer::GetDevice(), &createInfo);

            createInfo.min_filter = SDL_GPU_FILTER_NEAREST;
            createInfo.mag_filter = SDL_GPU_FILTER_NEAREST;
            createInfo.mipmap_mode = SDL_GPU_SAMPLERMIPMAPMODE_LINEAR;
            createInfo.address_mode_u = SDL_GPU_SAMPLERADDRESSMODE_REPEAT,
            createInfo.address_mode_v = SDL_GPU_SAMPLERADDRESSMODE_REPEAT,
            createInfo.address_mode_w = SDL_GPU_SAMPLERADDRESSMODE_REPEAT,

            g_linearRepeatSampler = SDL_CreateGPUSampler(Renderer::GetDevice(), &createInfo);
        }
    }

    VOID GPUResourceManager::Terminate()
    {
        SDL_ReleaseGPUSampler(Renderer::GetDevice(), g_linearClampSampler);
        SDL_ReleaseGPUSampler(Renderer::GetDevice(), g_linearRepeatSampler);
    }

    SDL_GPUSampler *GPUResourceManager::GetSampler_LinearClamp()
    {
        return g_linearClampSampler;
    }

    SDL_GPUSampler *GPUResourceManager::GetSampler_LinearRepeat()
    {
        return g_linearRepeatSampler;
    }

    SDL_GPUTexture *GPUResourceManager::CreateTexture(IN SDL_GPUTextureUsageFlags usage, IN INT32 width,
                                                      IN INT32 height, IN INT32 stride, IN PCUINT8 rgbaData,
                                                      IN SDL_GPUTextureFormat format, IN BOOL generateMipmaps)
    {
        const auto mipLevels =
            1; // generateMipmaps ? ia_max((UINT32)(floor(log2(ia_max(width, height))) + 1), (UINT32)1) : (UINT32)1;

        STATIC Vector<UINT8> TMP_COLOR_BUFFER;

        SDL_GPUTextureCreateInfo createInfo{.type = SDL_GPU_TEXTURETYPE_2D,
                                            .format = format,
                                            .usage = usage,
                                            .width = (UINT32) width,
                                            .height = (UINT32) height,
                                            .layer_count_or_depth = 1,
                                            .num_levels = (UINT32) mipLevels,
                                            .sample_count = SDL_GPU_SAMPLECOUNT_1};
        const auto result = SDL_CreateGPUTexture(Renderer::GetDevice(), &createInfo);
        if (!result)
        {
            THROW_UNKNOWN("Failed to create a SDL GPU Texture: ", SDL_GetError());
            return nullptr;
        }
        if (rgbaData)
        {
            TMP_COLOR_BUFFER.resize(width * height * 4);

            if (stride == width)
            {
                for (SIZE_T i = 0; i < TMP_COLOR_BUFFER.size() >> 2; i++)
                {
                    const auto a = static_cast<FLOAT32>(rgbaData[i * 4 + 3]) / 255.0f;
                    TMP_COLOR_BUFFER[i * 4 + 0] = static_cast<UINT8>(rgbaData[i * 4 + 0] * a);
                    TMP_COLOR_BUFFER[i * 4 + 1] = static_cast<UINT8>(rgbaData[i * 4 + 1] * a);
                    TMP_COLOR_BUFFER[i * 4 + 2] = static_cast<UINT8>(rgbaData[i * 4 + 2] * a);
                    TMP_COLOR_BUFFER[i * 4 + 3] = rgbaData[i * 4 + 3];
                }
            }
            else
            {
                for(INT32 y = 0; y < height; y++)
                {
                    for(INT32 x = 0; x < width; x++)
                    {
                        const auto p = &rgbaData[(x + y * stride) * 4];
                        const auto a = static_cast<FLOAT32>(p[3]) / 255.0f;
                        TMP_COLOR_BUFFER[(x + y * width) * 4 + 0] = static_cast<UINT8>(p[0] * a);
                        TMP_COLOR_BUFFER[(x + y * width) * 4 + 1] = static_cast<UINT8>(p[1] * a);
                        TMP_COLOR_BUFFER[(x + y * width) * 4 + 2] = static_cast<UINT8>(p[2] * a);
                        TMP_COLOR_BUFFER[(x + y * width) * 4 + 3] = p[3];
                    }
                }
            }

            SDL_GPUTransferBufferCreateInfo stagingBufferCreateInfo{.usage = SDL_GPU_TRANSFERBUFFERUSAGE_UPLOAD,
                                                                    .size = (UINT32) width * (UINT32) height * 4};
            const auto stagingBuffer = SDL_CreateGPUTransferBuffer(Renderer::GetDevice(), &stagingBufferCreateInfo);
            const auto mappedPtr = SDL_MapGPUTransferBuffer(Renderer::GetDevice(), stagingBuffer, false);
            SDL_memcpy(mappedPtr, TMP_COLOR_BUFFER.data(), width * height * 4);
            SDL_UnmapGPUTransferBuffer(Renderer::GetDevice(), stagingBuffer);

            auto cmdBuffer = SDL_AcquireGPUCommandBuffer(Renderer::GetDevice());
            const auto copyPass = SDL_BeginGPUCopyPass(cmdBuffer);

            SDL_GPUTextureTransferInfo transferInfo{.transfer_buffer = stagingBuffer, .offset = 0};
            SDL_GPUTextureRegion region{.texture = result, .w = (UINT32) width, .h = (UINT32) height, .d = 1};
            SDL_UploadToGPUTexture(copyPass, &transferInfo, &region, false);

            SDL_EndGPUCopyPass(copyPass);
            SDL_SubmitGPUCommandBuffer(cmdBuffer);
            SDL_WaitForGPUIdle(Renderer::GetDevice());
            SDL_ReleaseGPUTransferBuffer(Renderer::GetDevice(), stagingBuffer);

            if (mipLevels > 1)
            {
                cmdBuffer = SDL_AcquireGPUCommandBuffer(Renderer::GetDevice());
                SDL_GenerateMipmapsForGPUTexture(cmdBuffer, result);
                SDL_SubmitGPUCommandBuffer(cmdBuffer);
                SDL_WaitForGPUIdle(Renderer::GetDevice());
            }
        }
        return result;
    }

    SDL_GPUBuffer *GPUResourceManager::CreateDeviceLocalBuffer(IN SDL_GPUBufferUsageFlags usage, IN PCVOID data,
                                                               IN UINT32 dataSize)
    {
        SDL_GPUBufferCreateInfo createInfo{.usage = usage, .size = dataSize};
        const auto result = SDL_CreateGPUBuffer(Renderer::GetDevice(), &createInfo);
        if (!result)
        {
            THROW_UNKNOWN("Failed to create a SDL GPU Buffer: ", SDL_GetError());
            return nullptr;
        }
        if (dataSize && dataSize)
        {
            SDL_GPUTransferBufferCreateInfo stagingBufferCreateInfo{.usage = SDL_GPU_TRANSFERBUFFERUSAGE_UPLOAD,
                                                                    .size = dataSize};
            const auto stagingBuffer = SDL_CreateGPUTransferBuffer(Renderer::GetDevice(), &stagingBufferCreateInfo);
            const auto mappedPtr = SDL_MapGPUTransferBuffer(Renderer::GetDevice(), stagingBuffer, false);
            SDL_memcpy(mappedPtr, data, dataSize);
            SDL_UnmapGPUTransferBuffer(Renderer::GetDevice(), stagingBuffer);

            const auto cmdBuffer = SDL_AcquireGPUCommandBuffer(Renderer::GetDevice());
            const auto copyPass = SDL_BeginGPUCopyPass(cmdBuffer);

            SDL_GPUTransferBufferLocation src{.transfer_buffer = stagingBuffer, .offset = 0};
            SDL_GPUBufferRegion dst{.buffer = result, .offset = 0, .size = dataSize};
            SDL_UploadToGPUBuffer(copyPass, &src, &dst, false);

            SDL_EndGPUCopyPass(copyPass);
            SDL_SubmitGPUCommandBuffer(cmdBuffer);
            SDL_WaitForGPUIdle(Renderer::GetDevice());
            SDL_ReleaseGPUTransferBuffer(Renderer::GetDevice(), stagingBuffer);
        }
        return result;
    }

    VOID GPUResourceManager::DestroyTexture(IN SDL_GPUTexture *handle)
    {
        if (!handle)
            return;
        SDL_ReleaseGPUTexture(Renderer::GetDevice(), handle);
    }

    VOID GPUResourceManager::DestroyBuffer(IN SDL_GPUBuffer *handle)
    {
        SDL_ReleaseGPUBuffer(Renderer::GetDevice(), handle);
    }

    Vector<UINT8> GPUResourceManager::GetTexturePixelData(IN SDL_GPUTexture *texture, IN INT32 width, IN INT32 height)
    {
        SDL_GPUTransferBufferCreateInfo stagingBufferCreateInfo{.usage = SDL_GPU_TRANSFERBUFFERUSAGE_UPLOAD,
                                                                .size = (UINT32) width * (UINT32) height * 4};
        const auto stagingBuffer = SDL_CreateGPUTransferBuffer(Renderer::GetDevice(), &stagingBufferCreateInfo);

        const auto cmdBuffer = SDL_AcquireGPUCommandBuffer(Renderer::GetDevice());
        const auto copyPass = SDL_BeginGPUCopyPass(cmdBuffer);

        SDL_GPUTextureTransferInfo transferInfo{.transfer_buffer = stagingBuffer, .offset = 0};
        SDL_GPUTextureRegion region{.texture = texture, .w = (UINT32) width, .h = (UINT32) height, .d = 1};
        SDL_DownloadFromGPUTexture(copyPass, &region, &transferInfo);

        SDL_EndGPUCopyPass(copyPass);
        SDL_SubmitGPUCommandBuffer(cmdBuffer);
        SDL_WaitForGPUIdle(Renderer::GetDevice());

        Vector<UINT8> result;
        result.resize(width * height * 4);

        const auto mappedPtr = SDL_MapGPUTransferBuffer(Renderer::GetDevice(), stagingBuffer, false);
        SDL_UnmapGPUTransferBuffer(Renderer::GetDevice(), stagingBuffer);
        SDL_memcpy(result.data(), mappedPtr, result.size());
        SDL_ReleaseGPUTransferBuffer(Renderer::GetDevice(), stagingBuffer);

        return result;
    }

    SDL_GPUTexture *GPUResourceManager::CombineTextures(IN SDL_GPUTexture **textures, IN INT32 unitWidth,
                                                        IN INT32 unitHeight, IN INT32 unitCountX, IN INT32 unitCountY)
    {
        SDL_GPUTextureCreateInfo createInfo{.type = SDL_GPU_TEXTURETYPE_2D,
                                            .format = SDL_GPU_TEXTUREFORMAT_R8G8B8A8_UNORM,
                                            .usage = SDL_GPU_TEXTUREUSAGE_SAMPLER,
                                            .width = (UINT32) unitCountX * unitWidth,
                                            .height = (UINT32) unitCountY * unitHeight,
                                            .layer_count_or_depth = 1,
                                            .num_levels = 1};

        SDL_GPUTexture *result{};
        if (!(result = SDL_CreateGPUTexture(Renderer::GetDevice(), &createInfo)))
            THROW_UNKNOWN("Failed to create a SDL GPU Texture: ", SDL_GetError());

        const auto cmdBuffer = SDL_AcquireGPUCommandBuffer(Renderer::GetDevice());
        const auto copyPass = SDL_BeginGPUCopyPass(cmdBuffer);

        for (INT32 y = 0; y < unitCountY; y++)
        {
            for (INT32 x = 0; x < unitCountX; x++)
            {
                SDL_GPUTextureLocation src{.texture = textures[x + (y * unitCountX)], .x = 0, .y = 0};
                SDL_GPUTextureLocation dst{
                    .texture = result, .x = (UINT32) (x * unitWidth), .y = (UINT32) (y * unitHeight)};
                SDL_CopyGPUTextureToTexture(copyPass, &src, &dst, unitWidth, unitHeight, 1, false);
            }
        }

        SDL_EndGPUCopyPass(copyPass);
        SDL_SubmitGPUCommandBuffer(cmdBuffer);
        SDL_WaitForGPUIdle(Renderer::GetDevice());

        return result;
    }
} // namespace ia::iae