#include <Juliet.h>
#ifdef JULIET_ENABLE_IMGUI

#include <Graphics/ImGuiRenderer.h>

#include <Core/Logging/LogManager.h>
#include <Core/Logging/LogTypes.h>
#include <Core/HAL/Filesystem/Filesystem.h>
#include <Core/Memory/MemoryArena.h>
#include <Graphics/GraphicsPipeline.h>

#include <backends/imgui_impl_win32.h>
#include <imgui.h>

namespace Juliet
{
    namespace
    {
        constexpr uint32 kMaxFramesInFlight = 2;

        struct FrameResources
        {
            GraphicsBuffer*         VertexBuffer     = nullptr;
            GraphicsBuffer*         IndexBuffer      = nullptr;
            GraphicsTransferBuffer* VertexUpload     = nullptr;
            GraphicsTransferBuffer* IndexUpload      = nullptr;
            uint32                  VertexBufferSize = 0; // In bytes
            uint32                  IndexBufferSize  = 0; // In bytes
        };

        struct ImGuiRendererState
        {
            GraphicsDevice*   Device;
            GraphicsPipeline* Pipeline;

            // Resources
            Texture* FontTexture;
            Shader*  VertexShader;
            Shader*  FragmentShader;

            // Frame Data
            FrameResources Frames[kMaxFramesInFlight];
            uint32         FrameIndex;

            bool Initialized;
        };

        ImGuiRendererState g_ImGuiState = {};

        // Helper to resize buffers for a specific frame
        void EnsureBufferSize(FrameResources& frame, uint32 requiredVertexBytes, uint32 requiredIndexBytes)
        {
            bool needVertexRealloc = requiredVertexBytes > frame.VertexBufferSize;
            bool needIndexRealloc  = requiredIndexBytes > frame.IndexBufferSize;

            // If any buffer needs reallocation, wait for GPU to finish using everybody to be safe.
            // Ideally we only wait for this specific frame's fence, but we don't have that granularity easily exposed
            // here. Resizing is rare, so global wait is acceptable.
            if (needVertexRealloc || needIndexRealloc)
            {
                WaitUntilGPUIsIdle(g_ImGuiState.Device);
            }

            if (needVertexRealloc)
            {
                if (frame.VertexBuffer)
                {
                    DestroyGraphicsBuffer(g_ImGuiState.Device, frame.VertexBuffer);
                }
                if (frame.VertexUpload)
                {
                    DestroyGraphicsTransferBuffer(g_ImGuiState.Device, frame.VertexUpload);
                }

                frame.VertexBufferSize = requiredVertexBytes + 5000 * sizeof(ImDrawVert); // Growth factor

                BufferCreateInfo vci = {};
                vci.Size             = frame.VertexBufferSize;
                vci.Usage            = BufferUsage::StructuredBuffer;
                frame.VertexBuffer   = CreateGraphicsBuffer(g_ImGuiState.Device, vci);

                TransferBufferCreateInfo tvci = {};
                tvci.Size                     = frame.VertexBufferSize;
                tvci.Usage                    = TransferBufferUsage::Upload;
                frame.VertexUpload            = CreateGraphicsTransferBuffer(g_ImGuiState.Device, tvci);
            }

            if (needIndexRealloc)
            {
                if (frame.IndexBuffer)
                {
                    DestroyGraphicsBuffer(g_ImGuiState.Device, frame.IndexBuffer);
                }
                if (frame.IndexUpload)
                {
                    DestroyGraphicsTransferBuffer(g_ImGuiState.Device, frame.IndexUpload);
                }

                frame.IndexBufferSize = requiredIndexBytes + 10000 * sizeof(ImDrawIdx);

                BufferCreateInfo ici = {};
                ici.Size             = frame.IndexBufferSize;
                ici.Usage            = BufferUsage::IndexBuffer;
                frame.IndexBuffer    = CreateGraphicsBuffer(g_ImGuiState.Device, ici);

                TransferBufferCreateInfo tici = {};
                tici.Size                     = frame.IndexBufferSize;
                tici.Usage                    = TransferBufferUsage::Upload;
                frame.IndexUpload             = CreateGraphicsTransferBuffer(g_ImGuiState.Device, tici);
            }
        }

    } // namespace

    bool ImGuiRenderer_Initialize(GraphicsDevice* device)
    {
        printf("ImGuiRenderer_Initialize: device=%p, g_ImGuiState=%p, Initialized=%d\n", (void*)device,
               (void*)&g_ImGuiState, g_ImGuiState.Initialized);
        if (g_ImGuiState.Initialized)
        {
            return true;
        }

        g_ImGuiState.Device = device;

        // Load Shaders
        String           entryPoint = WrapString("main");
        ShaderCreateInfo shaderCI   = {};
        shaderCI.EntryPoint         = entryPoint;

        String vertPath           = GetAssetPath(WrapString("ImGui.vert.dxil"));
        shaderCI.Stage            = ShaderStage::Vertex;
        g_ImGuiState.VertexShader = CreateShader(device, vertPath, shaderCI);

        String fragPath             = GetAssetPath(WrapString("ImGui.frag.dxil"));
        shaderCI.Stage              = ShaderStage::Fragment;
        g_ImGuiState.FragmentShader = CreateShader(device, fragPath, shaderCI);

        if (!g_ImGuiState.VertexShader || !g_ImGuiState.FragmentShader)
        {
            LogError(LogCategory::Graphics, "Failed to load ImGui shaders");
            return false;
        }

        // Build Font Texture
        ImGuiIO&       io = ImGui::GetIO();
        unsigned char* pixels;
        int            width, height;
        io.Fonts->GetTexDataAsRGBA32(&pixels, &width, &height);

        TextureCreateInfo texCI = {};
        texCI.Type              = TextureType::Texture_2D;
        texCI.Width             = static_cast<uint32>(width);
        texCI.Height            = static_cast<uint32>(height);

        texCI.Format        = TextureFormat::R8G8B8A8_UNORM;
        texCI.Flags         = TextureUsageFlag::Sampler;
        texCI.LayerCount    = 1;
        texCI.MipLevelCount = 1;
        texCI.SampleCount   = TextureSampleCount::One;

        g_ImGuiState.FontTexture = CreateTexture(device, texCI);
        io.Fonts->SetTexID(reinterpret_cast<ImTextureID>(g_ImGuiState.FontTexture));

        // Upload data
        uint32 rowPitch        = static_cast<uint32>(width) * 4u;
        uint32 alignedRowPitch = (rowPitch + 255u) & ~255u;
        uint32 textureSize     = alignedRowPitch * static_cast<uint32>(height);

        TransferBufferCreateInfo tbci = {};
        tbci.Size                     = textureSize;
        tbci.Usage                    = TransferBufferUsage::Upload;
        GraphicsTransferBuffer* tb    = CreateGraphicsTransferBuffer(device, tbci);
        if (!tb)
        {
            Log(LogLevel::Error, LogCategory::Graphics, "ImGuiRenderer: Failed to create font upload buffer");
            return false;
        }

        auto* dst = static_cast<uint8*>(MapGraphicsTransferBuffer(device, tb));

        for (uint32 y = 0; y < static_cast<uint32>(height); ++y)
        {
            MemCopy(dst + y * alignedRowPitch, pixels + y * rowPitch, rowPitch);
        }

        UnmapGraphicsTransferBuffer(device, tb);

        CommandList* cmd = AcquireCommandList(device);
        CopyBufferToTexture(cmd, g_ImGuiState.FontTexture, tb);

        SubmitCommandLists(cmd);

        WaitUntilGPUIsIdle(device);

        DestroyGraphicsTransferBuffer(device, tb);

        // Init Frame Data
        g_ImGuiState.FrameIndex = 0;
        for (uint32 i = 0; i < kMaxFramesInFlight; ++i)
        {
            g_ImGuiState.Frames[i] = {};
        }

        g_ImGuiState.Initialized = true;
        return true;
    }

    void ImGuiRenderer_Shutdown(GraphicsDevice* device)
    {
        if (!g_ImGuiState.Initialized)
        {
            return;
        }

        WaitUntilGPUIsIdle(device);

        if (g_ImGuiState.Pipeline)
        {
            DestroyGraphicsPipeline(device, g_ImGuiState.Pipeline);
        }
        if (g_ImGuiState.VertexShader)
        {
            DestroyShader(device, g_ImGuiState.VertexShader);
        }
        if (g_ImGuiState.FragmentShader)
        {
            DestroyShader(device, g_ImGuiState.FragmentShader);
        }
        if (g_ImGuiState.FontTexture)
        {
            DestroyTexture(device, g_ImGuiState.FontTexture);
        }

        for (uint32 i = 0; i < kMaxFramesInFlight; ++i)
        {
            FrameResources& frame = g_ImGuiState.Frames[i];
            if (frame.VertexBuffer)
            {
                DestroyGraphicsBuffer(device, frame.VertexBuffer);
            }
            if (frame.IndexBuffer)
            {
                DestroyGraphicsBuffer(device, frame.IndexBuffer);
            }
            if (frame.VertexUpload)
            {
                DestroyGraphicsTransferBuffer(device, frame.VertexUpload);
            }
            if (frame.IndexUpload)
            {
                DestroyGraphicsTransferBuffer(device, frame.IndexUpload);
            }
            frame = {};
        }

        g_ImGuiState = {};
    }

    void ImGuiRenderer_NewFrame()
    {
        ImGui_ImplWin32_NewFrame(); // Assumes Win32 initialized elsewhere
        ImGui::NewFrame();
    }

    void ImGuiRenderer_Render(CommandList* cmdList, RenderPass* renderPass)
    {
        if (!g_ImGuiState.Initialized)
        {
            return;
        }

        ImGui::Render();
        ImDrawData* drawData = ImGui::GetDrawData();

        if (!drawData || drawData->CmdListsCount == 0)
        {
            return;
        }

        // Lazy create pipeline
        if (!g_ImGuiState.Pipeline)
        {
            // Assume B8G8R8A8_UNORM for SwapChain
            GraphicsPipelineCreateInfo pci = {};
            pci.VertexShader               = g_ImGuiState.VertexShader;
            pci.FragmentShader             = g_ImGuiState.FragmentShader;
            pci.PrimitiveType              = PrimitiveType::TriangleList;
            pci.RasterizerState.FillMode   = FillMode::Solid;
            pci.RasterizerState.CullMode   = CullMode::None;

            ColorTargetDescription colorDesc = {};
            colorDesc.Format                 = TextureFormat::B8G8R8A8_UNORM; // Match SwapChain

            // Blend State
            colorDesc.BlendState.EnableBlend                 = true;
            colorDesc.BlendState.SourceColorBlendFactor      = BlendFactor::Src_Alpha;
            colorDesc.BlendState.DestinationColorBlendFactor = BlendFactor::One_Minus_Src_Alpha;
            colorDesc.BlendState.ColorBlendOperation         = BlendOperation::Add;
            colorDesc.BlendState.SourceAlphaBlendFactor      = BlendFactor::One;
            colorDesc.BlendState.DestinationAlphaBlendFactor = BlendFactor::One_Minus_Src_Alpha;
            colorDesc.BlendState.AlphaBlendOperation         = BlendOperation::Add;
            colorDesc.BlendState.EnableColorWriteMask        = true;
            colorDesc.BlendState.ColorWriteMask =
                ColorComponentFlags::R | ColorComponentFlags::G | ColorComponentFlags::B | ColorComponentFlags::A;

            pci.TargetInfo.ColorTargetDescriptions = &colorDesc;
            pci.TargetInfo.NumColorTargets         = 1;

            g_ImGuiState.Pipeline = CreateGraphicsPipeline(g_ImGuiState.Device, pci);
        }

        // Cycle Frame
        g_ImGuiState.FrameIndex      = (g_ImGuiState.FrameIndex + 1) % kMaxFramesInFlight;
        FrameResources& currentFrame = g_ImGuiState.Frames[g_ImGuiState.FrameIndex];

        // Upload Buffers
        uint32 totalVtx = static_cast<uint32>(drawData->TotalVtxCount);
        uint32 totalIdx = static_cast<uint32>(drawData->TotalIdxCount);

        EnsureBufferSize(currentFrame, totalVtx * sizeof(ImDrawVert), totalIdx * sizeof(ImDrawIdx));

        auto* vtxDst = static_cast<ImDrawVert*>(MapGraphicsTransferBuffer(g_ImGuiState.Device, currentFrame.VertexUpload));
        auto* idxDst = static_cast<ImDrawIdx*>(MapGraphicsTransferBuffer(g_ImGuiState.Device, currentFrame.IndexUpload));

        for (int n = 0; n < drawData->CmdListsCount; n++)
        {
            const ImDrawList* cmd_list = drawData->CmdLists[n];
            MemCopy(vtxDst, cmd_list->VtxBuffer.Data, static_cast<size_t>(cmd_list->VtxBuffer.Size) * sizeof(ImDrawVert));
            MemCopy(idxDst, cmd_list->IdxBuffer.Data, static_cast<size_t>(cmd_list->IdxBuffer.Size) * sizeof(ImDrawIdx));
            vtxDst += cmd_list->VtxBuffer.Size;
            idxDst += cmd_list->IdxBuffer.Size;
        }

        UnmapGraphicsTransferBuffer(g_ImGuiState.Device, currentFrame.VertexUpload);
        UnmapGraphicsTransferBuffer(g_ImGuiState.Device, currentFrame.IndexUpload);

        // Copy both vertex and index buffers to GPU
        CopyBuffer(cmdList, currentFrame.VertexBuffer, currentFrame.VertexUpload, totalVtx * sizeof(ImDrawVert));
        CopyBuffer(cmdList, currentFrame.IndexBuffer, currentFrame.IndexUpload, totalIdx * sizeof(ImDrawIdx));

        // Transition vertex buffer to SRV state (this barrier waits for copy to complete)
        TransitionBufferToReadable(cmdList, currentFrame.VertexBuffer);

        // SetIndexBuffer transitions from COPY_DEST to INDEX_BUFFER (barrier waits for copy to complete)
        SetIndexBuffer(cmdList, currentFrame.IndexBuffer, IndexFormat::UInt16, totalIdx, 0);

        // Render
        BindGraphicsPipeline(renderPass, g_ImGuiState.Pipeline);

        // Setup Viewport / Scale
        float L = drawData->DisplayPos.x;
        float R = drawData->DisplayPos.x + drawData->DisplaySize.x;
        float T = drawData->DisplayPos.y;
        float B = drawData->DisplayPos.y + drawData->DisplaySize.y;

        float scale[2];
        scale[0] = 2.0f / (R - L);
        scale[1] = -2.0f / (B - T);

        float translate[2];
        translate[0] = -1.0f - L * scale[0];
        translate[1] = 1.0f - T * scale[1];

        // Global State (Display Size, etc)

        uint32 globalVtxOffset = 0;
        uint32 globalIdxOffset = 0;

        ImVec2 clip_off   = drawData->DisplayPos;
        ImVec2 clip_scale = drawData->FramebufferScale;

        for (int n = 0; n < drawData->CmdListsCount; n++)
        {
            const ImDrawList* cmd_list = drawData->CmdLists[n];
            for (int cmd_i = 0; cmd_i < cmd_list->CmdBuffer.Size; cmd_i++)
            {
                const ImDrawCmd* pcmd = &cmd_list->CmdBuffer[cmd_i];
                if (pcmd->UserCallback != nullptr)
                {
                    pcmd->UserCallback(cmd_list, pcmd);
                }
                else
                {
                    // Project scissor/clipping rectangles into framebuffer space
                    ImVec2 clip_min((pcmd->ClipRect.x - clip_off.x) * clip_scale.x,
                                    (pcmd->ClipRect.y - clip_off.y) * clip_scale.y);
                    ImVec2 clip_max((pcmd->ClipRect.z - clip_off.x) * clip_scale.x,
                                    (pcmd->ClipRect.w - clip_off.y) * clip_scale.y);

                    // Skip draw commands with zero-sized scissor rectangles
                    if (clip_max.x <= clip_min.x || clip_max.y <= clip_min.y)
                    {
                        continue;
                    }

                    // Set scissor rect (clipping rectangle)
                    Rectangle scissorRect = {};
                    scissorRect.X         = static_cast<int32>(clip_min.x);
                    scissorRect.Y         = static_cast<int32>(clip_min.y);
                    scissorRect.Width     = static_cast<int32>(clip_max.x - clip_min.x);
                    scissorRect.Height    = static_cast<int32>(clip_max.y - clip_min.y);
                    SetScissorRect(renderPass, scissorRect);

                    // Bind Texture
                    uint32 textureIndex = GetDescriptorIndex(g_ImGuiState.Device, reinterpret_cast<Texture*>(pcmd->GetTexID()));

                    // Push Constants
                    PushData pushData = {}; // Zero-initialize all fields

                    pushData.BufferIndex  = GetDescriptorIndex(g_ImGuiState.Device, currentFrame.VertexBuffer);
                    pushData.TextureIndex = textureIndex;
                    pushData.VertexOffset = pcmd->VtxOffset + globalVtxOffset; // Pass vertex offset for bindless
                    pushData.Scale[0]     = scale[0];
                    pushData.Scale[1]     = scale[1];
                    pushData.Translate[0] = translate[0];
                    pushData.Translate[1] = translate[1];

                    SetPushConstants(cmdList, ShaderStage::Vertex, 0, sizeof(pushData) / 4, &pushData);

                    // Draw - pass 0 for vertexOffset since shader handles it via push constants
                    DrawIndexedPrimitives(renderPass, pcmd->ElemCount, 1, pcmd->IdxOffset + globalIdxOffset, 0, 0);
                }
            }
            globalIdxOffset += static_cast<uint32>(cmd_list->IdxBuffer.Size);
            globalVtxOffset += static_cast<uint32>(cmd_list->VtxBuffer.Size);
        }
    }
} // namespace Juliet

#endif // JULIET_ENABLE_IMGUI
