// 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/DebugDraw.hpp>
#include <Renderer/EmbeddedShader.hpp>
#include <Renderer/Renderer.hpp>

#include <ResourceManager.hpp>
#include <WorldManager.hpp>

#include <IAEngine/imgui/backends/imgui_impl_sdl3.h>
#include <IAEngine/imgui/backends/imgui_impl_sdlgpu3.h>

#include <algorithm>

namespace ia::iae
{
    STATIC Mat4 IdentityMatrix(1.0f);

    EXTERN SDL_Window *g_windowHandle;

    INT32 Renderer::s_screenWidth{};
    INT32 Renderer::s_screenHeight{};
    Renderer::State Renderer::s_state{};
    SDL_GPUDevice *Renderer::s_gpuDevice{};

    SDL_GPUTexture *Renderer::s_renderTargetSceneColor{};
    SDL_GPUTexture *Renderer::s_renderTargetDebugDrawColor{};

    Pipeline *Renderer::s_geometryPipeline{};
    Pipeline *Renderer::s_postprocessPipeline{};

    Renderer::Geometry *Renderer::s_quadGeometry{};

    Vector<Renderer::DrawEntry> Renderer::s_drawEntries;

    SDL_GPURenderPass *Renderer::s_activeRenderPass{};
    SDL_GPUCommandBuffer *Renderer::s_activeCommandBuffer{};
    SDL_GPUColorTargetInfo Renderer::s_colorTargetInfo{};
    class CameraComponent *Renderer::s_activeCamera{};

    SDL_Rect Renderer::s_scissor{0, 0, 0, 0};
    SDL_GPUViewport Renderer::s_activeViewport{};
    Vec2 Renderer::s_sceneScaleFactor{1.0f, 1.0f};

    VOID Renderer::Initialize()
    {
        if (!(s_gpuDevice = SDL_CreateGPUDevice(SDL_GPU_SHADERFORMAT_SPIRV, Engine::IsDebugMode(), nullptr)))
            THROW_UNKNOWN("Failed to create the SDL GPU Device: ", SDL_GetError());

        if (!SDL_ClaimWindowForGPUDevice(s_gpuDevice, g_windowHandle))
            THROW_UNKNOWN("Failed to initialize SDL GPU for the window: ", SDL_GetError());

        SDL_SetGPUSwapchainParameters(s_gpuDevice, g_windowHandle, SDL_GPU_SWAPCHAINCOMPOSITION_SDR,
                                      SDL_GPU_PRESENTMODE_VSYNC);

        GPUResourceManager::Initialize();

        SDL_GetWindowSizeInPixels(g_windowHandle, &s_screenWidth, &s_screenHeight);
        OnScreenResize(s_screenWidth, s_screenHeight);

        // Initialize Pipelines
        s_geometryPipeline = new Pipeline(
            Pipeline::StageDesc{
                .SourceData = SHADER_SOURCE_GEOMETRY_VERT,
                .SourceLength = sizeof(SHADER_SOURCE_GEOMETRY_VERT),
                .SamplerCount = 0,
                .UniformBufferCount = 3,
            },
            Pipeline::StageDesc{
                .SourceData = SHADER_SOURCE_GEOMETRY_FRAG,
                .SourceLength = sizeof(SHADER_SOURCE_GEOMETRY_FRAG),
                .SamplerCount = 1,
                .UniformBufferCount = 1,
            },
            true, true);
        s_postprocessPipeline = new Pipeline(
            Pipeline::StageDesc{
                .SourceData = SHADER_SOURCE_POSTPROCESS_VERT,
                .SourceLength = sizeof(SHADER_SOURCE_POSTPROCESS_VERT),
                .SamplerCount = 0,
                .UniformBufferCount = 0,
            },
            Pipeline::StageDesc{
                .SourceData = SHADER_SOURCE_POSTPROCESS_FRAG,
                .SourceLength = sizeof(SHADER_SOURCE_POSTPROCESS_FRAG),
                .SamplerCount = 2,
                .UniformBufferCount = 0,
            },
            false, false);

        DebugDraw::Initialize();

        s_colorTargetInfo.clear_color = SDL_FColor{0.0f, 0.0f, 0.0f, 1.0f};
        s_colorTargetInfo.load_op = SDL_GPU_LOADOP_CLEAR;
        s_colorTargetInfo.store_op = SDL_GPU_STOREOP_STORE;

        s_quadGeometry = CreateGeometry(
            {
                {glm::vec3{0, 1, 0}, glm::vec2{0, 1}, glm::vec4{1.0f, 1.0f, 1.0f, 1.0f}},
                {glm::vec3{1, 1, 0}, glm::vec2{1, 1}, glm::vec4{1.0f, 1.0f, 1.0f, 1.0f}},
                {glm::vec3{1, 0, 0}, glm::vec2{1, 0}, glm::vec4{1.0f, 1.0f, 1.0f, 1.0f}},
                {glm::vec3{0, 0, 0}, glm::vec2{0, 0}, glm::vec4{1.0f, 1.0f, 1.0f, 1.0f}},
            },
            {0, 1, 2, 2, 3, 0});
    }

    VOID Renderer::Terminate()
    {
        SDL_WaitForGPUIdle(s_gpuDevice);

        DestroyGeometry(s_quadGeometry);

        DebugDraw::Terminate();

        delete s_geometryPipeline;
        delete s_postprocessPipeline;

        GPUResourceManager::DestroyTexture(s_renderTargetSceneColor);
        GPUResourceManager::DestroyTexture(s_renderTargetDebugDrawColor);

        GPUResourceManager::Terminate();

        SDL_ReleaseWindowFromGPUDevice(s_gpuDevice, g_windowHandle);
        SDL_DestroyGPUDevice(s_gpuDevice);
    }

    VOID Renderer::BeginFrame()
    {
        s_drawEntries.clear();

        if (!(s_activeCommandBuffer = SDL_AcquireGPUCommandBuffer(s_gpuDevice)))
            THROW_UNKNOWN("Failed to acquire SDL GPU command buffer: ", SDL_GetError());

        const auto clearColor = WorldManager::GetActiveScene()->BackgroundColor().GetAsFloatVec();
        s_colorTargetInfo.clear_color = SDL_FColor{clearColor.x, clearColor.y, clearColor.z, 1.0f};
        s_colorTargetInfo.texture = s_renderTargetSceneColor;
        s_activeRenderPass = SDL_BeginGPURenderPass(s_activeCommandBuffer, &s_colorTargetInfo, 1, nullptr);

        SDL_BindGPUGraphicsPipeline(s_activeRenderPass, s_geometryPipeline->GetHandle());
        SDL_PushGPUVertexUniformData(s_activeCommandBuffer, 0,
                                     s_activeCamera ? s_activeCamera->GetProjectionMatrix() : &IdentityMatrix,
                                     sizeof(Mat4));
    }

    VOID Renderer::EndFrame()
    {
#pragma pack(push, 1)

        STATIC struct
        {
            INT32 FlippedH{false};
            INT32 FlippedV{false};
            Vec2 TextureOffset{0.0f, 0.0f};
            Vec4 ColorOverlay{1.0f, 1.0f, 1.0f, 1.0f};
        } s_fragmentUniform{};

#pragma pack(pop)

        std::sort(s_drawEntries.begin(), s_drawEntries.end(), [](IN CONST DrawEntry &a, IN CONST DrawEntry &b) {
            if(a.Layer != b.Layer)
                return a.Layer < b.Layer;
            return a.SortIndex < b.SortIndex; 
        });

        for (const auto &t : s_drawEntries)
        {
            SDL_PushGPUVertexUniformData(
                s_activeCommandBuffer, 1,
                (s_activeCamera && t.DrawState.CameraRelative) ? s_activeCamera->GetViewMatrix() : &IdentityMatrix,
                sizeof(Mat4));

            SDL_PushGPUVertexUniformData(Renderer::s_activeCommandBuffer, 2, &t.DrawState.ModelMatrix, sizeof(Mat4));

            s_fragmentUniform.ColorOverlay = t.DrawState.ColorOverlay.GetAsFloatVec();
            s_fragmentUniform.FlippedH = t.DrawState.FlippedH;
            s_fragmentUniform.FlippedV = t.DrawState.FlippedV;
            s_fragmentUniform.TextureOffset = t.DrawState.TextureOffset;
            SDL_GPUTextureSamplerBinding textureBinding{.texture = t.DrawState.ActiveTexture,
                                                        .sampler = ((t.DrawState.TextureOffset.x <= FLOAT32_EPSILON) &&
                                                                    (t.DrawState.TextureOffset.y <= FLOAT32_EPSILON))
                                                                       ? GPUResourceManager::GetSampler_LinearClamp()
                                                                       : GPUResourceManager::GetSampler_LinearRepeat()};
            SDL_BindGPUFragmentSamplers(s_activeRenderPass, 0, &textureBinding, 1);
            SDL_PushGPUFragmentUniformData(s_activeCommandBuffer, 0, &s_fragmentUniform, sizeof(s_fragmentUniform));

            // SDL_SetGPUScissor(s_activeRenderPass, &s_scissor);
            // SDL_SetGPUViewport(s_activeRenderPass, &s_activeViewport);

            SDL_GPUBufferBinding bufferBindings[] = {{.buffer = t.GeometryHandle->VertexBuffer, .offset = 0},
                                                     {.buffer = t.GeometryHandle->IndexBuffer, .offset = 0}};
            SDL_BindGPUVertexBuffers(s_activeRenderPass, 0, bufferBindings, 1);
            SDL_BindGPUIndexBuffer(s_activeRenderPass, &bufferBindings[1], SDL_GPU_INDEXELEMENTSIZE_32BIT);
            SDL_DrawGPUIndexedPrimitives(s_activeRenderPass, t.GeometryHandle->IndexCount, 1, 0, 0, 0);
        }

        SDL_EndGPURenderPass(s_activeRenderPass);

        DebugDraw::Render();
        const auto imDrawData = ImGui::GetDrawData();
        ImGui_ImplSDLGPU3_PrepareDrawData(imDrawData, s_activeCommandBuffer);
        s_colorTargetInfo.load_op = SDL_GPU_LOADOP_CLEAR;
        s_colorTargetInfo.clear_color = SDL_FColor{0.0f, 0.0f, 0.0f, 0.0f};
        s_colorTargetInfo.texture = s_renderTargetDebugDrawColor;
        s_activeRenderPass = SDL_BeginGPURenderPass(s_activeCommandBuffer, &s_colorTargetInfo, 1, nullptr);
        ImGui_ImplSDLGPU3_RenderDrawData(imDrawData, s_activeCommandBuffer, s_activeRenderPass);
        SDL_EndGPURenderPass(s_activeRenderPass);

        SDL_GPUTexture *swapChainTexture{};
        if (!SDL_WaitAndAcquireGPUSwapchainTexture(s_activeCommandBuffer, g_windowHandle, &swapChainTexture,
                                                   (PUINT32) &s_screenWidth, (PUINT32) &s_screenHeight))
            THROW_UNKNOWN("Failed to acquire SDL GPU Swapchain texture: ", SDL_GetError());

        if (!swapChainTexture)
            return;

        s_colorTargetInfo.clear_color = SDL_FColor{1.0f, 1.0f, 1.0f, 1.0f};
        s_colorTargetInfo.load_op = SDL_GPU_LOADOP_CLEAR;
        s_colorTargetInfo.texture = swapChainTexture;

        s_activeRenderPass = SDL_BeginGPURenderPass(s_activeCommandBuffer, &s_colorTargetInfo, 1, nullptr);
        SDL_BindGPUGraphicsPipeline(s_activeRenderPass, s_postprocessPipeline->GetHandle());
        SDL_GPUTextureSamplerBinding textureBindings[2] = {
            {.texture = s_renderTargetSceneColor, .sampler = GPUResourceManager::GetSampler_LinearClamp()},
            {.texture = s_renderTargetDebugDrawColor, .sampler = GPUResourceManager::GetSampler_LinearClamp()},
        };
        SDL_BindGPUFragmentSamplers(s_activeRenderPass, 0, textureBindings, 2);
        SDL_DrawGPUPrimitives(s_activeRenderPass, 6, 1, 0, 0);
        SDL_EndGPURenderPass(s_activeRenderPass);

        SDL_SubmitGPUCommandBuffer(s_activeCommandBuffer);
    }

    VOID Renderer::OnScreenResize(IN INT32 newWidth, IN INT32 newHeight)
    {
        s_screenWidth = newWidth;
        s_screenHeight = newHeight;

        if (s_renderTargetSceneColor)
            GPUResourceManager::DestroyTexture(s_renderTargetSceneColor);
        if (s_renderTargetDebugDrawColor)
            GPUResourceManager::DestroyTexture(s_renderTargetDebugDrawColor);

        s_renderTargetSceneColor = GPUResourceManager::CreateTexture(
            SDL_GPU_TEXTUREUSAGE_COLOR_TARGET | SDL_GPU_TEXTUREUSAGE_SAMPLER, s_screenWidth, s_screenHeight, nullptr,
            SDL_GetGPUSwapchainTextureFormat(s_gpuDevice, g_windowHandle));
        s_renderTargetDebugDrawColor = GPUResourceManager::CreateTexture(
            SDL_GPU_TEXTUREUSAGE_COLOR_TARGET | SDL_GPU_TEXTUREUSAGE_SAMPLER, s_screenWidth, s_screenHeight, nullptr,
            SDL_GetGPUSwapchainTextureFormat(s_gpuDevice, g_windowHandle));

        s_scissor = {0, 0, newWidth, newHeight};

        s_activeViewport.x = 0;
        s_activeViewport.y = 0;
        s_activeViewport.w = newWidth;
        s_activeViewport.h = newHeight;
        s_activeViewport.min_depth = 0.0f;
        s_activeViewport.max_depth = 1.0f;

        if (s_activeCamera)
            s_activeCamera->SetViewport(newWidth, newHeight);

        const auto activeScene = WorldManager::GetActiveScene();
        if (activeScene)
        {
            const auto sceneViewport = activeScene->Viewport();
            s_sceneScaleFactor = {(FLOAT32) newWidth / (FLOAT32) sceneViewport.x,
                                  (FLOAT32) newHeight / (FLOAT32) sceneViewport.y};
            IAE_LOG_INFO("Updated Scene Scale Factor: (", s_sceneScaleFactor.x, ", ", s_sceneScaleFactor.y, ")");
        }
    }

    SDL_GPUTextureFormat Renderer::GetRenderTargetFormat()
    {
        return SDL_GetGPUSwapchainTextureFormat(s_gpuDevice, g_windowHandle);
    }

    Renderer::Geometry *Renderer::CreateGeometry(IN CONST Vector<GeometryVertex> &vertices,
                                                 IN CONST Vector<INT32> &indices)
    {
        const auto mesh = new Geometry();
        mesh->VertexBuffer = GPUResourceManager::CreateDeviceLocalBuffer(
            SDL_GPU_BUFFERUSAGE_VERTEX, vertices.data(), static_cast<UINT32>(vertices.size() * sizeof(vertices[0])));
        mesh->IndexBuffer = GPUResourceManager::CreateDeviceLocalBuffer(
            SDL_GPU_BUFFERUSAGE_INDEX, indices.data(), static_cast<UINT32>(indices.size() * sizeof(indices[0])));
        mesh->IndexCount = static_cast<UINT32>(indices.size());
        return mesh;
    }

    VOID Renderer::DestroyGeometry(IN Geometry *handle)
    {
        GPUResourceManager::DestroyBuffer(handle->VertexBuffer);
        GPUResourceManager::DestroyBuffer(handle->IndexBuffer);
        delete handle;
    }

    VOID Renderer::DrawGeometry(IN Geometry *handle, IN UINT8 layer, IN UINT16 sortIndex)
    {
        const auto t = (UINT16) (sortIndex + ((UINT16) s_state.PositionY * s_state.YSortingEnabled));
        s_drawEntries.pushBack(
            DrawEntry{.Layer = layer,
                      .SortIndex = t,
                      .DrawState = s_state,
                      .GeometryHandle = handle});
    }

    VOID Renderer::WaitForGPUIdle()
    {
        SDL_WaitForGPUIdle(s_gpuDevice);
    }
} // namespace ia::iae

namespace ia::iae
{
    IVec2 Engine::GetDisplayExtent()
    {
        return IVec2(Renderer::s_screenWidth, Renderer::s_screenHeight);
    }

    VOID Engine::SetActiveCamera(IN CameraComponent *cameraComponent)
    {
        Renderer::s_activeCamera = cameraComponent;
        Renderer::OnScreenResize(Renderer::s_screenWidth, Renderer::s_screenHeight);
    }

    CameraComponent *Engine::GetActiveCamera()
    {
        return Renderer::s_activeCamera;
    }

    Vec2 Engine::GetRendererScalingFactor()
    {
        return Renderer::s_sceneScaleFactor;
    }

    Handle Engine::CreateGeometry(IN CONST Vector<GeometryVertex> &vertices, IN CONST Vector<INT32> &indices)
    {
        return (Handle) Renderer::CreateGeometry(vertices, indices);
    }

    VOID Engine::DestroyGeometry(IN Handle geometry)
    {
        Renderer::DestroyGeometry((Renderer::Geometry *) geometry);
    }

    VOID Engine::DrawGeometry(IN Handle handle, IN UINT8 layer, IN UINT16 sortIndex)
    {
        Renderer::DrawGeometry((Renderer::Geometry *) handle, layer, sortIndex);
    }

    VOID Engine::SetRenderState_Scissor(IN IVec4 rect)
    {
        Renderer::s_scissor = rect.z ? SDL_Rect{0, 0, Renderer::s_screenWidth, Renderer::s_screenHeight}
                                     : SDL_Rect{rect.x, rect.y, rect.z, rect.w};
    }

    VOID Engine::SetRenderState_FlippedH(IN BOOL value)
    {
        Renderer::s_state.FlippedH = value;
    }

    VOID Engine::SetRenderState_FlippedV(IN BOOL value)
    {
        Renderer::s_state.FlippedV = value;
    }

    VOID Engine::SetRenderState_TextureOffset(IN Vec2 off)
    {
        Renderer::s_state.TextureOffset = off;
    }

    VOID Engine::SetRenderState_ColorOverlay(IN Color color)
    {
        Renderer::s_state.ColorOverlay = color;
    }

    VOID Engine::SetRenderState_CameraRelative(IN BOOL value)
    {
        Renderer::s_state.CameraRelative = value;
    }

    VOID Engine::SetRenderState_Texture(IN Handle image)
    {
        Renderer::s_state.ActiveTexture = ResourceManager::GetTextureFromImage(image);
    }

    VOID Engine::SetRenderState_Transform(IN Vec2 position, IN Vec2 scale, IN FLOAT32 rotation)
    {
        position *= Renderer::s_sceneScaleFactor;
        scale *= Renderer::s_sceneScaleFactor;

        Renderer::s_state.PositionY = position.y;

        Renderer::s_state.ModelMatrix = glm::translate(glm::mat4(1.0f), glm::vec3{position.x, position.y, 0});
        Renderer::s_state.ModelMatrix =
            glm::rotate(Renderer::s_state.ModelMatrix, rotation, glm::vec3(0.0f, 0.0f, 1.0f));
        Renderer::s_state.ModelMatrix = glm::scale(Renderer::s_state.ModelMatrix, glm::vec3{scale.x, scale.y, 1.0f});
    }

    VOID Engine::SetRenderState_TransformUI(IN Vec2 position, IN Vec2 scale, IN FLOAT32 rotation)
    {
        Renderer::s_state.ModelMatrix = glm::translate(glm::mat4(1.0f), glm::vec3{position.x, position.y, 0});
        Renderer::s_state.ModelMatrix =
            glm::rotate(Renderer::s_state.ModelMatrix, rotation, glm::vec3(0.0f, 0.0f, 1.0f));
        Renderer::s_state.ModelMatrix = glm::scale(Renderer::s_state.ModelMatrix, glm::vec3{scale.x, scale.y, 1.0f});
    }

    VOID Engine::SetRenderState_YSortingEnabled(IN BOOL value)
    {
        Renderer::s_state.YSortingEnabled = value;
    }

    Handle Engine::GetGeometry_Quad()
    {
        return (Handle) Renderer::s_quadGeometry;
    }
} // namespace ia::iae