// 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>

namespace ia::iae
{
    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_renderTargetSceneDepth{};
    SDL_GPUTexture *Renderer::s_renderTargetDebugDrawColor{};

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

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

    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_state.ColorTargetInfo.clear_color = SDL_FColor{0.0f, 0.0f, 0.0f, 1.0f};
        s_state.ColorTargetInfo.load_op = SDL_GPU_LOADOP_CLEAR;
        s_state.ColorTargetInfo.store_op = SDL_GPU_STOREOP_STORE;

        s_state.DepthStencilTargetInfo.cycle = true;
        s_state.DepthStencilTargetInfo.clear_depth = 0;
        s_state.DepthStencilTargetInfo.clear_stencil = 0;
        s_state.DepthStencilTargetInfo.load_op = SDL_GPU_LOADOP_CLEAR;
        s_state.DepthStencilTargetInfo.store_op = SDL_GPU_STOREOP_STORE;
        s_state.DepthStencilTargetInfo.stencil_load_op = SDL_GPU_LOADOP_CLEAR;
        s_state.DepthStencilTargetInfo.stencil_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_renderTargetSceneDepth);
        GPUResourceManager::DestroyTexture(s_renderTargetSceneColor);
        GPUResourceManager::DestroyTexture(s_renderTargetDebugDrawColor);

        GPUResourceManager::Terminate();

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

    VOID Renderer::BeginFrame()
    {
        STATIC Mat4 IdentityMatrix(1.0f);

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

        const auto clearColor = WorldManager::GetActiveScene()->BackgroundColor().GetAsFloatVec();
        s_state.ColorTargetInfo.clear_color = SDL_FColor{clearColor.x, clearColor.y, clearColor.z, 1.0f};
        s_state.ColorTargetInfo.texture = s_renderTargetSceneColor;
        s_state.DepthStencilTargetInfo.texture = s_renderTargetSceneDepth;
        s_state.ActiveRenderPass = SDL_BeginGPURenderPass(s_state.ActiveCommandBuffer, &s_state.ColorTargetInfo, 1,
                                                          &s_state.DepthStencilTargetInfo);

        SDL_BindGPUGraphicsPipeline(s_state.ActiveRenderPass, s_geometryPipeline->GetHandle());
        SDL_PushGPUVertexUniformData(
            s_state.ActiveCommandBuffer, 0,
            s_state.ActiveCamera ? s_state.ActiveCamera->GetProjectionMatrix() : &IdentityMatrix, sizeof(Mat4));
        SDL_PushGPUVertexUniformData(
            s_state.ActiveCommandBuffer, 1,
            (s_state.ActiveCamera && !s_state.CameraRelative) ? s_state.ActiveCamera->GetViewMatrix() : &IdentityMatrix,
            sizeof(Mat4));
    }

    VOID Renderer::EndFrame()
    {
        SDL_EndGPURenderPass(s_state.ActiveRenderPass);

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

        SDL_GPUTexture *swapChainTexture{};
        if (!SDL_WaitAndAcquireGPUSwapchainTexture(s_state.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_state.ColorTargetInfo.clear_color = SDL_FColor{1.0f, 1.0f, 1.0f, 1.0f};
        s_state.ColorTargetInfo.load_op = SDL_GPU_LOADOP_CLEAR;
        s_state.ColorTargetInfo.texture = swapChainTexture;

        s_state.ActiveRenderPass =
            SDL_BeginGPURenderPass(s_state.ActiveCommandBuffer, &s_state.ColorTargetInfo, 1, nullptr);
        SDL_BindGPUGraphicsPipeline(s_state.ActiveRenderPass, s_postprocessPipeline->GetHandle());
        SDL_GPUTextureSamplerBinding textureBindings[2] = {
            {.texture = s_renderTargetSceneColor, .sampler = GPUResourceManager::GetSampler_LinearRepeat()},
            {.texture = s_renderTargetDebugDrawColor, .sampler = GPUResourceManager::GetSampler_LinearRepeat()},
        };
        SDL_BindGPUFragmentSamplers(s_state.ActiveRenderPass, 0, textureBindings, 2);
        SDL_DrawGPUPrimitives(s_state.ActiveRenderPass, 6, 1, 0, 0);
        SDL_EndGPURenderPass(s_state.ActiveRenderPass);

        SDL_SubmitGPUCommandBuffer(s_state.ActiveCommandBuffer);
    }

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

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

        s_renderTargetSceneDepth =
            GPUResourceManager::CreateTexture(SDL_GPU_TEXTUREUSAGE_DEPTH_STENCIL_TARGET, s_screenWidth, s_screenHeight,
                                              nullptr, SDL_GPU_TEXTUREFORMAT_D16_UNORM);
        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_state.Scissor = {0, 0, newWidth, newHeight};

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

        if (s_state.ActiveCamera)
            s_state.ActiveCamera->SetViewport(newWidth, newHeight);

        const auto activeScene = WorldManager::GetActiveScene();
        if (activeScene)
        {
            const auto sceneExtent = activeScene->Extent();
            s_state.SceneScaleFactor = {(FLOAT32) newWidth / (FLOAT32) sceneExtent.x,
                                        (FLOAT32) newHeight / (FLOAT32) sceneExtent.y};
            IAE_LOG_INFO("Updated Scene Scale Factor: (", s_state.SceneScaleFactor.x, ", ", s_state.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)
    {
#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)

        s_fragmentUniform.ColorOverlay = s_state.ColorOverlay.GetAsFloatVec();
        s_fragmentUniform.FlippedH = s_state.FlippedH;
        s_fragmentUniform.FlippedV = s_state.FlippedV;
        s_fragmentUniform.TextureOffset = s_state.TextureOffset;
        SDL_GPUTextureSamplerBinding textureBinding{.texture = s_state.ActiveTexture,
                                                    .sampler = GPUResourceManager::GetSampler_LinearRepeat()};
        SDL_BindGPUFragmentSamplers(s_state.ActiveRenderPass, 0, &textureBinding, 1);
        SDL_PushGPUFragmentUniformData(s_state.ActiveCommandBuffer, 0, &s_fragmentUniform, sizeof(s_fragmentUniform));

        SDL_SetGPUScissor(s_state.ActiveRenderPass, &s_state.Scissor);
        SDL_SetGPUViewport(s_state.ActiveRenderPass, &s_state.ActiveViewport);

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

    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_state.ActiveCamera = cameraComponent;
        Renderer::OnScreenResize(Renderer::s_screenWidth, Renderer::s_screenHeight);
    }

    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)
    {
        Renderer::DrawGeometry((Renderer::Geometry *) handle);
    }

    VOID Engine::SetRenderState_Scissor(IN IVec4 rect)
    {
        Renderer::s_state.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, IN UINT8 layer,
                                          IN INT16 sortIndex)
    {
        IA_ASSERT(sortIndex <= 0x1FFF);

        if (Renderer::s_state.YSortingEnabled)
            sortIndex += static_cast<INT16>(position.y);

        position *= Renderer::s_state.SceneScaleFactor;
        scale *= Renderer::s_state.SceneScaleFactor;

        Renderer::s_state.ModelMatrix =
            glm::translate(glm::mat4(1.0f), glm::vec3{position.x, position.y,
                                                      static_cast<FLOAT32>((layer << 13) | (sortIndex & 0x1FFF))});
        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});

        SDL_PushGPUVertexUniformData(Renderer::s_state.ActiveCommandBuffer, 2, &Renderer::s_state.ModelMatrix,
                                     sizeof(Mat4));
    }

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

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