#include <Graphics/MeshRenderer.h>

#include <Core/Logging/LogManager.h>
#include <Core/Logging/LogTypes.h>
#include <Core/Math/Matrix.h>
#include <Core/HAL/Filesystem/Filesystem.h>
#include <Graphics/GraphicsDevice.h>
#include <Graphics/Mesh.h>
#include <Graphics/VertexData.h>

namespace Juliet
{
    namespace
    {
        MeshRenderer g_MeshRenderer;
    } // namespace

    bool InitializeMeshRenderer(NonNullPtr<Arena> arena, NonNullPtr<GraphicsDevice> device, NonNullPtr<Window> window)
    {
        bool result = true;

        g_MeshRenderer.Meshes.Create(arena JULIET_DEBUG_PARAM("Meshes"));
        g_MeshRenderer.Vertices.Create(arena JULIET_DEBUG_PARAM("Vertices"));
        g_MeshRenderer.Indices.Create(arena JULIET_DEBUG_PARAM("Indices"));

        GraphicsDevice* graphicsDevice = g_MeshRenderer.Device = device.Get();

        // Create graphics pipeline
        String           entryPoint = WrapString("main");
        ShaderCreateInfo shaderCI   = {};
        shaderCI.EntryPoint         = entryPoint;

        String shaderPath    = GetAssetPath(WrapString("Triangle.vert.dxil"));
        shaderCI.Stage       = ShaderStage::Vertex;
        Shader* vertexShader = CreateShader(graphicsDevice, shaderPath, shaderCI);

        shaderPath             = GetAssetPath(WrapString("SolidColor.frag.dxil"));
        shaderCI.Stage         = ShaderStage::Fragment;
        Shader* fragmentShader = CreateShader(graphicsDevice, shaderPath, shaderCI);

        ColorTargetDescription colorTargetDescription = {};
        colorTargetDescription.Format                 = GetSwapChainTextureFormat(graphicsDevice, window);

        GraphicsPipelineCreateInfo pipelineCI         = {};
        pipelineCI.VertexShader                       = vertexShader;
        pipelineCI.FragmentShader                     = fragmentShader;
        pipelineCI.PrimitiveType                      = PrimitiveType::TriangleList;
        pipelineCI.TargetInfo                         = { .ColorTargetDescriptions = &colorTargetDescription,
                                                          .NumColorTargets         = 1,
                                                          .DepthStencilFormat      = TextureFormat::D32_FLOAT,
                                                          .HasDepthStencilTarget   = true };
        pipelineCI.RasterizerState.FillMode           = FillMode::Solid;
        pipelineCI.DepthStencilState.EnableDepthTest  = true;
        pipelineCI.DepthStencilState.EnableDepthWrite = true;
        pipelineCI.DepthStencilState.CompareOperation = CompareOperation::Less;

        g_MeshRenderer.Pipeline = CreateGraphicsPipeline(graphicsDevice, pipelineCI);
        if (g_MeshRenderer.Pipeline == nullptr)
        {
            LogError(LogCategory::Graphics, "Failed to create graphics pipeline!");
            result = false;
        }

        // Create the vertex and index buffers
        BufferCreateInfo vertexBufferCI = {};
        vertexBufferCI.Size             = kGeometryPage;
        vertexBufferCI.Usage            = BufferUsage::StructuredBuffer;
        g_MeshRenderer.VertexBuffer     = CreateGraphicsBuffer(graphicsDevice, vertexBufferCI);

        BufferCreateInfo indexBufferCI = {};
        indexBufferCI.Size             = kIndexPage;
        indexBufferCI.Usage            = BufferUsage::IndexBuffer;
        g_MeshRenderer.IndexBuffer     = CreateGraphicsBuffer(graphicsDevice, indexBufferCI);

        if (vertexShader)
        {
            DestroyShader(graphicsDevice, vertexShader);
        }
        if (fragmentShader)
        {
            DestroyShader(graphicsDevice, fragmentShader);
        }

        return result;
    }

    void ShutdownMeshRenderer()
    {
        if (g_MeshRenderer.LoadCopyBuffer)
        {
            UnmapGraphicsTransferBuffer(g_MeshRenderer.Device, g_MeshRenderer.LoadCopyBuffer);
            DestroyGraphicsTransferBuffer(g_MeshRenderer.Device, g_MeshRenderer.LoadCopyBuffer);
            g_MeshRenderer.LoadCopyBuffer = nullptr;
        }

        DestroyGraphicsBuffer(g_MeshRenderer.Device, g_MeshRenderer.IndexBuffer);
        DestroyGraphicsBuffer(g_MeshRenderer.Device, g_MeshRenderer.VertexBuffer);

        g_MeshRenderer.Indices.Destroy();
        g_MeshRenderer.Vertices.Destroy();
        g_MeshRenderer.Meshes.Destroy();
    }

    void LoadMeshesOnGPU(NonNullPtr<CommandList> cmdList)
    {
        if (g_MeshRenderer.Meshes.IsEmpty())
        {
            return;
        }

        // Loading everything in one go.
        // Destroy the buffer at the end

        TransferBufferCreateInfo uploadBCI = {};
        uploadBCI.Usage                    = TransferBufferUsage::Upload;
        uploadBCI.Size                     = kGeometryPage + kIndexPage;
        g_MeshRenderer.LoadCopyBuffer      = CreateGraphicsTransferBuffer(g_MeshRenderer.Device, uploadBCI);

        void* map = MapGraphicsTransferBuffer(g_MeshRenderer.Device, g_MeshRenderer.LoadCopyBuffer);

        Vertex* vertices = g_MeshRenderer.Vertices.Data;
        if (!vertices)
        {
            UnmapGraphicsTransferBuffer(g_MeshRenderer.Device, g_MeshRenderer.LoadCopyBuffer);
            DestroyGraphicsTransferBuffer(g_MeshRenderer.Device, g_MeshRenderer.LoadCopyBuffer);
            g_MeshRenderer.LoadCopyBuffer = nullptr;
            return;
        }

        Index* indices = g_MeshRenderer.Indices.Data;
        if (!indices)
        {
            UnmapGraphicsTransferBuffer(g_MeshRenderer.Device, g_MeshRenderer.LoadCopyBuffer);
            DestroyGraphicsTransferBuffer(g_MeshRenderer.Device, g_MeshRenderer.LoadCopyBuffer);
            g_MeshRenderer.LoadCopyBuffer = nullptr;
            return;
        }

        // Copy all meshes! This supports only one page for now
        // Copy ALL Vertices in one block
        size_t totalVertexBytes = g_MeshRenderer.Vertices.Count * sizeof(Vertex);
        MemCopy(map, g_MeshRenderer.Vertices.Data, totalVertexBytes);

        // Copy ALL Indices in one block
        size_t indexOfByteOffset = (kGeometryPage + 255) & static_cast<size_t>(~255);
        uint8* ptrOneByte        = static_cast<uint8*>(map);
        size_t totalIndexBytes   = g_MeshRenderer.Indices.Count * sizeof(Index);

        MemCopy(ptrOneByte + indexOfByteOffset, g_MeshRenderer.Indices.Data, totalIndexBytes);

        // index_t index             = 0;
        // for (Mesh& mesh : g_MeshRenderer.Meshes)
        // {
        //     // Vertices first
        //     Vertex* mapVertices = static_cast<Vertex*>(map);
        //     MemCopy(mapVertices + index, vertices + mesh.VertexOffset, mesh.VertexCount * sizeOfVertex);
        //
        //     // Indices next
        //     uint8*  ptrOneByte = static_cast<uint8*>(map);
        //     uint16* dst        = reinterpret_cast<uint16*>(ptrOneByte + indexOfByteOffset);
        //     MemCopy(dst, indices, mesh.IndexCount * sizeOfIndex);
        //
        //     ++index;
        // }

        CopyBuffer(cmdList, g_MeshRenderer.VertexBuffer, g_MeshRenderer.LoadCopyBuffer, totalVertexBytes, 0, 0);
        CopyBuffer(cmdList, g_MeshRenderer.IndexBuffer, g_MeshRenderer.LoadCopyBuffer, totalIndexBytes, 0, indexOfByteOffset);

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

    void RenderMeshes(NonNullPtr<RenderPass> pass, NonNullPtr<CommandList> cmdList, PushData& pushData)
    {
        // First destroy any buffer that needs to be
        if (g_MeshRenderer.LoadCopyBuffer)
        {
            WaitUntilGPUIsIdle(g_MeshRenderer.Device);
            UnmapGraphicsTransferBuffer(g_MeshRenderer.Device, g_MeshRenderer.LoadCopyBuffer);
            DestroyGraphicsTransferBuffer(g_MeshRenderer.Device, g_MeshRenderer.LoadCopyBuffer);
            g_MeshRenderer.LoadCopyBuffer = nullptr;
        }

        BindGraphicsPipeline(pass, g_MeshRenderer.Pipeline);

        uint32 vertexDescriptorIndex = GetDescriptorIndex(g_MeshRenderer.Device, g_MeshRenderer.VertexBuffer);

        pushData.BufferIndex = vertexDescriptorIndex;

        SetIndexBuffer(cmdList, g_MeshRenderer.IndexBuffer, IndexFormat::UInt16, g_MeshRenderer.Indices.Count, 0);

        for (Mesh& mesh : g_MeshRenderer.Meshes)
        {
            pushData.Model = mesh.Transform;
            SetPushConstants(cmdList, ShaderStage::Vertex, 0, sizeof(pushData) / 4, &pushData);
            DrawIndexedPrimitives(pass, static_cast<uint32>(mesh.IndexCount), 1, static_cast<uint32>(mesh.IndexOffset),
                                  static_cast<uint32>(mesh.VertexOffset), 0);
        }
    }

    MeshID AddCube()
    {
        Mesh result = {};

        constexpr Vertex vertexData[] = {
            // Front Face (Z = -0.5f) — Red
            { { -0.5f, 0.5f, -0.5f }, { 0.8f, 0.2f, 0.2f, 1.0f } },
            { { 0.5f, 0.5f, -0.5f }, { 0.8f, 0.2f, 0.2f, 1.0f } },
            { { 0.5f, -0.5f, -0.5f }, { 0.8f, 0.2f, 0.2f, 1.0f } },
            { { -0.5f, -0.5f, -0.5f }, { 0.8f, 0.2f, 0.2f, 1.0f } },

            // Back Face (Z = 0.5f) — Green
            { { 0.5f, 0.5f, 0.5f }, { 0.2f, 0.8f, 0.2f, 1.0f } },
            { { -0.5f, 0.5f, 0.5f }, { 0.2f, 0.8f, 0.2f, 1.0f } },
            { { -0.5f, -0.5f, 0.5f }, { 0.2f, 0.8f, 0.2f, 1.0f } },
            { { 0.5f, -0.5f, 0.5f }, { 0.2f, 0.8f, 0.2f, 1.0f } },

            // Top Face (Y = 0.5f) — Blue
            { { -0.5f, 0.5f, 0.5f }, { 0.2f, 0.2f, 0.8f, 1.0f } },
            { { 0.5f, 0.5f, 0.5f }, { 0.2f, 0.2f, 0.8f, 1.0f } },
            { { 0.5f, 0.5f, -0.5f }, { 0.2f, 0.2f, 0.8f, 1.0f } },
            { { -0.5f, 0.5f, -0.5f }, { 0.2f, 0.2f, 0.8f, 1.0f } },

            // Bottom Face (Y = -0.5f) — Yellow
            { { -0.5f, -0.5f, -0.5f }, { 0.8f, 0.8f, 0.2f, 1.0f } },
            { { 0.5f, -0.5f, -0.5f }, { 0.8f, 0.8f, 0.2f, 1.0f } },
            { { 0.5f, -0.5f, 0.5f }, { 0.8f, 0.8f, 0.2f, 1.0f } },
            { { -0.5f, -0.5f, 0.5f }, { 0.8f, 0.8f, 0.2f, 1.0f } },

            // Right Face (X = 0.5f) — Cyan
            { { 0.5f, 0.5f, -0.5f }, { 0.2f, 0.8f, 0.8f, 1.0f } },
            { { 0.5f, 0.5f, 0.5f }, { 0.2f, 0.8f, 0.8f, 1.0f } },
            { { 0.5f, -0.5f, 0.5f }, { 0.2f, 0.8f, 0.8f, 1.0f } },
            { { 0.5f, -0.5f, -0.5f }, { 0.2f, 0.8f, 0.8f, 1.0f } },

            // Left Face (X = -0.5f) — Magenta
            { { -0.5f, 0.5f, 0.5f }, { 0.8f, 0.2f, 0.8f, 1.0f } },
            { { -0.5f, 0.5f, -0.5f }, { 0.8f, 0.2f, 0.8f, 1.0f } },
            { { -0.5f, -0.5f, -0.5f }, { 0.8f, 0.2f, 0.8f, 1.0f } },
            { { -0.5f, -0.5f, 0.5f }, { 0.8f, 0.2f, 0.8f, 1.0f } }
        };
        constexpr size_t cubeVertexCount = ArraySize(vertexData);
        result.VertexCount               = cubeVertexCount;
        result.VertexOffset              = g_MeshRenderer.Vertices.Count;
        g_MeshRenderer.Vertices.PushBack(vertexData, cubeVertexCount);

        constexpr uint16 indices[]      = { 0,  1,  2,  0,  2,  3,    // Front
                                            4,  5,  6,  4,  6,  7,    // Back
                                            8,  9,  10, 8,  10, 11,   // Top
                                            12, 13, 14, 12, 14, 15,   // Bottom
                                            16, 17, 18, 16, 18, 19,   // Right
                                            20, 21, 22, 20, 22, 23 }; // Left
        constexpr size_t indexCubeCount = ArraySize(indices);
        result.IndexCount               = indexCubeCount;
        result.IndexOffset              = g_MeshRenderer.Indices.Count;
        g_MeshRenderer.Indices.PushBack(indices, indexCubeCount);

        MeshID id = g_MeshRenderer.Meshes.Count;
        g_MeshRenderer.Meshes.PushBack(result);
        return id;
    }

    MeshID AddQuad()
    {
        // Mesh result = {};

        // // Using the exact 6 vertices from the working triangles!
        // constexpr Vertex vertexData[] = {
        //     // Triangle 1 (Clockwise)
        //     { { -0.5f, -0.5f, 0.0f }, { 1.0f, 0.0f, 0.0f, 1.0f } }, // 0: Red
        //     { { 0.0f, 0.5f, 0.0f }, { 0.0f, 1.0f, 0.0f, 1.0f } },   // 1: Green
        //     { { 0.5f, -0.5f, 0.0f }, { 0.0f, 0.0f, 1.0f, 1.0f } },  // 2: Blue
        //
        //     // Triangle 2 (Clockwise)
        //     { { -0.5f, 0.5f, 0.0f }, { 1.0f, 1.0f, 0.0f, 1.0f } }, // 3: Yellow
        //     { { 0.0f, 0.8f, 0.0f }, { 0.0f, 1.0f, 1.0f, 1.0f } },  // 4: Cyan
        //     { { 0.5f, 0.5f, 0.0f }, { 1.0f, 0.0f, 1.0f, 1.0f } }   // 5: Magenta
        // };
        // constexpr size_t triVertexCount = ArraySize(vertexData);
        // result.VertexCount              = triVertexCount;
        // result.Vertices                 = ArenaPushArray<Vertex>(arena.Get(), triVertexCount);
        // MemCopy(result.Vertices, vertexData, sizeof(Vertex) * triVertexCount);
        //
        // // Just the 6 indices for the two triangles
        // constexpr uint16 indices[]     = { 0, 1, 2, 3, 4, 5 };
        // constexpr size_t triIndexCount = ArraySize(indices);
        // result.IndexCount              = triIndexCount;
        // result.Indices = ArenaPushArray<uint16>(arena.Get(), triIndexCount JULIET_DEBUG_PARAM("Indices"));
        // MemCopy(result.Indices, indices, sizeof(uint16) * triIndexCount);
        //
        // g_MeshRenderer.Meshes.PushBack(std::move(result));
        return g_MeshRenderer.Meshes.Count - 1;
    }

    void SetMeshTransform(MeshID id, const Matrix& transform)
    {
        Assert(id < static_cast<MeshID>(g_MeshRenderer.Meshes.Count));
        g_MeshRenderer.Meshes.Data[id].Transform = transform;
    }
#if ALLOW_SHADER_HOT_RELOAD
    void ReloadMeshRendererShaders()
    {
        auto* pipeline = g_MeshRenderer.Pipeline;
        auto* device   = g_MeshRenderer.Device;

        String           entryPoint = WrapString("main");
        ShaderCreateInfo shaderCI   = {};
        shaderCI.EntryPoint         = entryPoint;
        String shaderPath           = GetAssetPath(WrapString("Triangle.vert.dxil"));
        shaderCI.Stage              = ShaderStage::Vertex;
        Shader* vertexShader        = CreateShader(device, shaderPath, shaderCI);

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

        UpdateGraphicsPipelineShaders(device, pipeline, vertexShader, fragmentShader);

        if (vertexShader)
        {
            DestroyShader(device, vertexShader);
        }
        if (fragmentShader)
        {
            DestroyShader(device, fragmentShader);
        }
    }
#endif
} // namespace Juliet
