#include <Graphics/DebugDisplay.h>

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

namespace Juliet
{
    namespace
    {
        constexpr uint32 kMaxDebugVertices = 16384; // 16K vertices = ~450KB per buffer

        struct DebugVertex
        {
            Vector3 Position;
            FColor  Color;
        };

        struct DebugDisplayState
        {
            GraphicsDevice* Device;

            // Pipelines
            GraphicsPipeline* DepthTestedPipeline;
            GraphicsPipeline* OverlayPipeline;

            // Vertex data (CPU side - single array with two regions)
            DebugVertex* Vertices; // Single allocation for all vertices
            uint32       DepthTestedVertexCount;
            uint32       OverlayVertexCount;

            // GPU buffers (consolidated - single buffer pair)
            GraphicsBuffer*         VertexBuffer;
            GraphicsTransferBuffer* TransferBuffer;

            bool Initialized;
        };

        DebugDisplayState g_DebugState = {};

        void AddLine(DebugVertex* vertices, uint32& count, const Vector3& start, const Vector3& end, const FColor& color)
        {
            if (count + 2 > kMaxDebugVertices)
            {
                return;
            }

            vertices[count++] = { start, color };
            vertices[count++] = { end, color };
        }

        void AddSphereWireframe(DebugVertex* vertices, uint32& count, const Vector3& center, float radius, const FColor& color)
        {
            constexpr int   segments = 16;
            constexpr float pi       = 3.14159265358979f;

            // Draw 3 circles (XY, XZ, YZ planes)
            for (int i = 0; i < segments; ++i)
            {
                float a1 = (float)i / segments * 2.0f * pi;
                float a2 = (float)(i + 1) / segments * 2.0f * pi;

                // XY circle
                Vector3 p1 = { center.x + cosf(a1) * radius, center.y + sinf(a1) * radius, center.z };
                Vector3 p2 = { center.x + cosf(a2) * radius, center.y + sinf(a2) * radius, center.z };
                AddLine(vertices, count, p1, p2, color);

                // XZ circle
                p1 = { center.x + cosf(a1) * radius, center.y, center.z + sinf(a1) * radius };
                p2 = { center.x + cosf(a2) * radius, center.y, center.z + sinf(a2) * radius };
                AddLine(vertices, count, p1, p2, color);

                // YZ circle
                p1 = { center.x, center.y + cosf(a1) * radius, center.z + sinf(a1) * radius };
                p2 = { center.x, center.y + cosf(a2) * radius, center.z + sinf(a2) * radius };
                AddLine(vertices, count, p1, p2, color);
            }
        }

        GraphicsPipeline* CreateDebugPipeline(GraphicsDevice* device, TextureFormat colorFormat, bool enableDepthTest)
        {
            String           entryPoint = WrapString("main");
            ShaderCreateInfo shaderCI   = {};
            shaderCI.EntryPoint         = entryPoint;

            String vertPath      = GetAssetPath(WrapString("Debug.vert.dxil"));
            shaderCI.Stage       = ShaderStage::Vertex;
            Shader* vertexShader = CreateShader(device, vertPath, shaderCI);

            String fragPath        = GetAssetPath(WrapString("Debug.frag.dxil"));
            shaderCI.Stage         = ShaderStage::Fragment;
            Shader* fragmentShader = CreateShader(device, fragPath, shaderCI);

            if (!vertexShader || !fragmentShader)
            {
                LogError(LogCategory::Graphics, "Failed to create debug shaders");
                if (vertexShader)
                {
                    DestroyShader(device, vertexShader);
                }
                if (fragmentShader)
                {
                    DestroyShader(device, fragmentShader);
                }
                return nullptr;
            }

            ColorTargetDescription colorDesc = {};
            colorDesc.Format                 = colorFormat;

            GraphicsPipelineCreateInfo createInfo   = {};
            createInfo.VertexShader                 = vertexShader;
            createInfo.FragmentShader               = fragmentShader;
            createInfo.PrimitiveType                = PrimitiveType::LineList;
            createInfo.RasterizerState.FillMode     = FillMode::Solid;
            createInfo.RasterizerState.CullMode     = CullMode::None;
            createInfo.RasterizerState.FrontFace    = FrontFace::CounterClockwise;
            createInfo.MultisampleState.SampleCount = TextureSampleCount::One;

            // No vertex input state - using bindless buffer access via SV_VertexID

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

            // Now that we support depth-stencil targets in the backend, we can enable them.
            createInfo.TargetInfo.HasDepthStencilTarget = true;
            createInfo.TargetInfo.DepthStencilFormat    = TextureFormat::D32_FLOAT;

            if (enableDepthTest)
            {
                createInfo.DepthStencilState.EnableDepthTest  = true;
                createInfo.DepthStencilState.EnableDepthWrite = true;
                createInfo.DepthStencilState.CompareOperation = CompareOperation::Less;
            }
            else
            {
                // Overlay pipeline: depth test/write off, so it draws on top
                createInfo.DepthStencilState.EnableDepthTest  = false;
                createInfo.DepthStencilState.EnableDepthWrite = false;
                createInfo.DepthStencilState.CompareOperation = CompareOperation::Always;
            }

            GraphicsPipeline* pipeline = CreateGraphicsPipeline(device, createInfo);

            DestroyShader(device, vertexShader);
            DestroyShader(device, fragmentShader);

            return pipeline;
        }
    } // namespace

    void DebugDisplay_Initialize(GraphicsDevice* device)
    {
        if (g_DebugState.Initialized)
        {
            return;
        }

        g_DebugState.Device = device;

        // Allocate single CPU vertex array (depth-tested in first half, overlay in second half)
        g_DebugState.Vertices = static_cast<DebugVertex*>(Malloc(kMaxDebugVertices * sizeof(DebugVertex)));
        g_DebugState.DepthTestedVertexCount = 0;
        g_DebugState.OverlayVertexCount     = 0;

        // Create single GPU buffer pair (consolidated)
        BufferCreateInfo bufferCI = {};
        bufferCI.Size             = kMaxDebugVertices * sizeof(DebugVertex);
        bufferCI.Usage            = BufferUsage::StructuredBuffer;

        g_DebugState.VertexBuffer = CreateGraphicsBuffer(device, bufferCI);

        TransferBufferCreateInfo transferCI = {};
        transferCI.Size                     = kMaxDebugVertices * sizeof(DebugVertex);
        transferCI.Usage                    = TransferBufferUsage::Upload;

        g_DebugState.TransferBuffer = CreateGraphicsTransferBuffer(device, transferCI);

        g_DebugState.Initialized = true;
    }

    void DebugDisplay_Shutdown(GraphicsDevice* device)
    {
        if (!g_DebugState.Initialized)
        {
            return;
        }

        if (g_DebugState.DepthTestedPipeline)
        {
            DestroyGraphicsPipeline(device, g_DebugState.DepthTestedPipeline);
        }
        if (g_DebugState.OverlayPipeline)
        {
            DestroyGraphicsPipeline(device, g_DebugState.OverlayPipeline);
        }

        if (g_DebugState.VertexBuffer)
        {
            DestroyGraphicsBuffer(device, g_DebugState.VertexBuffer);
        }
        if (g_DebugState.TransferBuffer)
        {
            DestroyGraphicsTransferBuffer(device, g_DebugState.TransferBuffer);
        }

        SafeFree(g_DebugState.Vertices);

        g_DebugState = {};
    }

    void DebugDisplay_DrawLine(const Vector3& start, const Vector3& end, const FColor& color, bool overlay)
    {
        // Depth-tested vertices at beginning, overlay vertices at end (growing backward)
        if (overlay)
        {
            // Overlay vertices grow from end of buffer backward
            uint32 halfMax = kMaxDebugVertices / 2;
            AddLine(g_DebugState.Vertices + halfMax, g_DebugState.OverlayVertexCount, start, end, color);
        }
        else
        {
            // Depth-tested vertices grow from start of buffer forward
            AddLine(g_DebugState.Vertices, g_DebugState.DepthTestedVertexCount, start, end, color);
        }
    }

    void DebugDisplay_DrawSphere(const Vector3& center, float radius, const FColor& color, bool overlay)
    {
        if (overlay)
        {
            uint32 halfMax = kMaxDebugVertices / 2;
            AddSphereWireframe(g_DebugState.Vertices + halfMax, g_DebugState.OverlayVertexCount, center, radius, color);
        }
        else
        {
            AddSphereWireframe(g_DebugState.Vertices, g_DebugState.DepthTestedVertexCount, center, radius, color);
        }
    }

    void DebugDisplay_Prepare(CommandList* cmdList)
    {
        if (!g_DebugState.Initialized)
        {
            return;
        }

        uint32 totalVertices = g_DebugState.DepthTestedVertexCount + g_DebugState.OverlayVertexCount;
        if (totalVertices == 0 || !g_DebugState.VertexBuffer)
        {
            return;
        }

        // Upload all vertex data in one copy (depth-tested at start, overlay at kMaxDebugVertices/2)
        void* ptr = MapGraphicsTransferBuffer(g_DebugState.Device, g_DebugState.TransferBuffer);
        if (ptr)
        {
            // Copy depth-tested vertices (at start)
            if (g_DebugState.DepthTestedVertexCount > 0)
            {
                MemCopy(ptr, g_DebugState.Vertices, g_DebugState.DepthTestedVertexCount * sizeof(DebugVertex));
            }
            // Copy overlay vertices (at kMaxDebugVertices/2 offset)
            if (g_DebugState.OverlayVertexCount > 0)
            {
                uint32 halfMax     = kMaxDebugVertices / 2;
                auto*  overlayDest = static_cast<uint8*>(ptr) + halfMax * sizeof(DebugVertex);
                auto*  overlaySrc  = g_DebugState.Vertices + halfMax;
                MemCopy(overlayDest, overlaySrc, g_DebugState.OverlayVertexCount * sizeof(DebugVertex));
            }
            UnmapGraphicsTransferBuffer(g_DebugState.Device, g_DebugState.TransferBuffer);

            // Single buffer copy (full buffer to ensure both regions are uploaded)
            size_t copySize = kMaxDebugVertices * sizeof(DebugVertex);
            CopyBuffer(cmdList, g_DebugState.VertexBuffer, g_DebugState.TransferBuffer, copySize);
            TransitionBufferToReadable(cmdList, g_DebugState.VertexBuffer);
        }
    }

    void DebugDisplay_Flush(CommandList* cmdList, RenderPass* renderPass, const Camera& camera)
    {
        if (!g_DebugState.Initialized)
        {
            return;
        }

        // Lazy-create pipelines (need swapchain format)
        if (!g_DebugState.DepthTestedPipeline)
        {
            // Use B8G8R8A8_UNORM which matches the SDR swapchain
            g_DebugState.DepthTestedPipeline = CreateDebugPipeline(g_DebugState.Device, TextureFormat::B8G8R8A8_UNORM, true);
            g_DebugState.OverlayPipeline = CreateDebugPipeline(g_DebugState.Device, TextureFormat::B8G8R8A8_UNORM, false);
        }

        uint32 bufferIndex = GetDescriptorIndex(g_DebugState.Device, g_DebugState.VertexBuffer);

        // Render depth-tested primitives (vertices at offset 0 in buffer)
        if (g_DebugState.DepthTestedVertexCount > 0 && g_DebugState.DepthTestedPipeline && g_DebugState.VertexBuffer)
        {
            // Pack VP matrix + buffer index into push constants
            PushData pushData = {};
            pushData.ViewProjection = Camera_GetViewProjectionMatrix(camera);
            pushData.MeshIndex      = 0;
            pushData.TransformsBufferIndex = 0; // Not used by debug shader but layout must match
            pushData.BufferIndex    = bufferIndex;
            pushData.TextureIndex   = 0;
            pushData.VertexOffset   = 0; // Depth-tested vertices start at 0
            pushData.Padding        = 0;
            pushData.Scale[0] = 1.0f; pushData.Scale[1] = 1.0f;
            pushData.Translate[0] = 0.0f; pushData.Translate[1] = 0.0f;
            
            // Dummy light data as we don't light debug primitives
            pushData.GlobalLightDirection = {0,0,-1};
            pushData.GlobalLightPad = 0;
            pushData.GlobalLightColor = {1,1,1};
            pushData.GlobalAmbientIntensity = 1.0f;
            
            SetPushConstants(cmdList, ShaderStage::Vertex, 0, sizeof(pushData) / sizeof(uint32), &pushData);

            DrawPrimitives(renderPass, g_DebugState.DepthTestedVertexCount, 1, 0, 0);
        }

        // Render overlay primitives (vertices at offset kMaxDebugVertices/2 in buffer)
        if (g_DebugState.OverlayVertexCount > 0 && g_DebugState.OverlayPipeline && g_DebugState.VertexBuffer)
        {
            BindGraphicsPipeline(renderPass, g_DebugState.OverlayPipeline);

            // Pack VP matrix + buffer index into push constants
            PushData pushData = {};
            pushData.ViewProjection = Camera_GetViewProjectionMatrix(camera);
            pushData.MeshIndex      = 0;
            pushData.TransformsBufferIndex = 0;
            pushData.BufferIndex    = bufferIndex;
            pushData.TextureIndex   = 0;
            pushData.VertexOffset   = kMaxDebugVertices / 2; // Overlay vertices start at half
            pushData.Padding        = 0;
            pushData.Scale[0] = 1.0f; pushData.Scale[1] = 1.0f;
            pushData.Translate[0] = 0.0f; pushData.Translate[1] = 0.0f;
            
            // Dummy light data as we don't light debug primitives
            pushData.GlobalLightDirection = {0,0,-1};
            pushData.GlobalLightPad = 0;
            pushData.GlobalLightColor = {1,1,1};
            pushData.GlobalAmbientIntensity = 1.0f;
            
            SetPushConstants(cmdList, ShaderStage::Vertex, 0, sizeof(pushData) / sizeof(uint32), &pushData);

            DrawPrimitives(renderPass, g_DebugState.OverlayVertexCount, 1, 0, 0);
        }

        // Clear for next frame
        g_DebugState.DepthTestedVertexCount = 0;
        g_DebugState.OverlayVertexCount     = 0;
    }
} // namespace Juliet
