#include "Orange/Render/Input.h"
#include <Orange/Core/Array.h>
#include <Orange/Core/FileSystem.h>
#include <Orange/Core/Result.h>
#include <Orange/Core/Span.h>
#include <Orange/Core/Time.h>
#include <Orange/Core/Variant.h>
#include <Orange/Core/Vector.h>
#include <Orange/Graphics/Camera.h>
#include <Orange/Graphics/Mesh.h>
#include <Orange/Math/Matrix.h>
#include <Orange/Math/Transformation.h>
#include <Orange/Math/Vector.h>
#include <Orange/Render/RenderContext.h>
#include <Orange/Render/Swapchain.h>
#include <Orange/Render/VulkanHelpers.h>
#include <Orange/Render/Window.h>
#include <Orange/Formats/KeyValues.h>

#include <vs_Fullscreen.h>
#include <fs_HDRTest.h>

#include <iostream>
#include <spng.h>
#include <iomanip>
#include <vulkan/vulkan_core.h>

using namespace orange;

struct MiniPipelineInfo
{
    VkShaderModule vs;
    VkShaderModule fs;

    VkPipelineLayout pipelineLayout;

    VkFormat colorFormat;
    VkFormat depthFormat;

    bool enableDepthTest;
    bool enableDepthWrites;

    Span<const VkVertexInputAttributeDescription> vertexAttributes;
};

static VulkanResult<VkPipeline> CreateGraphicsPipeline(VkDevice device, const MiniPipelineInfo& info)
{
    VkPipelineShaderStageCreateInfo stages[] =
    {
        {
            .sType  = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
            .stage  = VK_SHADER_STAGE_VERTEX_BIT,
            .module = info.vs,
            .pName  = "main",
        },
        {
            .sType  = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
            .stage  = VK_SHADER_STAGE_FRAGMENT_BIT,
            .module = info.fs,
            .pName  = "main",
        },
    };

    Array<VkDynamicState, 2> dynamicStates =
    {
        VK_DYNAMIC_STATE_VIEWPORT,
        VK_DYNAMIC_STATE_SCISSOR,
    };

    VkPipelineDynamicStateCreateInfo dynamicStateInfo =
    {
        .sType             = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO,
        .dynamicStateCount = uint32_t(dynamicStates.size()),
        .pDynamicStates    = dynamicStates.data(),
    };

    VkVertexInputBindingDescription inputBindingDescription =
    {
        .binding   = 0,
        .stride    = sizeof(StaticVertex),
        .inputRate = VK_VERTEX_INPUT_RATE_VERTEX,
    };

    VkPipelineVertexInputStateCreateInfo vertexInputInfo =
    {
        .sType                           = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO,
        .vertexBindingDescriptionCount   = info.vertexAttributes.size ? 1u : 0u,
        .pVertexBindingDescriptions      = &inputBindingDescription,
        .vertexAttributeDescriptionCount = uint32_t(info.vertexAttributes.size),
        .pVertexAttributeDescriptions    = info.vertexAttributes.data,
    };

    VkPipelineInputAssemblyStateCreateInfo inputAssemblyInfo =
    {
        .sType                  = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO,
        .topology               = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST,
        .primitiveRestartEnable = VK_FALSE,
    };

    VkPipelineViewportStateCreateInfo viewportStateInfo =
    {
        .sType         = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO,
        .viewportCount = 1,
        .scissorCount  = 1,
    };

    VkPipelineRasterizationStateCreateInfo rasterizationInfo =
    {
        .sType     = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO,
        .cullMode  = VK_CULL_MODE_BACK_BIT,
        .frontFace = VK_FRONT_FACE_COUNTER_CLOCKWISE,
        .lineWidth = 1.0f,
    };

    VkPipelineMultisampleStateCreateInfo multisampleInfo =
    {
        .sType                = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO,
        .rasterizationSamples = VK_SAMPLE_COUNT_1_BIT,
    };

    VkPipelineDepthStencilStateCreateInfo depthStencilInfo =
    {
        .sType            = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO,
        .depthTestEnable  = info.enableDepthTest,
        .depthWriteEnable = info.enableDepthWrites,
        .depthCompareOp   = VK_COMPARE_OP_LESS,
        .minDepthBounds   = 0.0f,
        .maxDepthBounds   = 1.0f,
    };

    VkPipelineColorBlendAttachmentState attachmentBlendState =
    {
        .colorWriteMask =  VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT |
                           VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT,
    };

    VkPipelineColorBlendStateCreateInfo colorBlendStateInfo =
    {
        .sType           = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO,
        .attachmentCount = 1,
        .pAttachments    = &attachmentBlendState,
    };

    VkPipelineRenderingCreateInfo renderingInfo =
    {
        .sType                   = VK_STRUCTURE_TYPE_PIPELINE_RENDERING_CREATE_INFO,
        .viewMask                = 0u,
        .colorAttachmentCount    = 1u,
        .pColorAttachmentFormats = &info.colorFormat,
        .depthAttachmentFormat   = info.depthFormat,
    };

    VkGraphicsPipelineCreateInfo pipelineInfo =
    {
        .sType               = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO,
        .pNext               = &renderingInfo,
        .stageCount          = Size(stages),
        .pStages             = stages,
        .pVertexInputState   = &vertexInputInfo,
        .pInputAssemblyState = &inputAssemblyInfo,
        .pViewportState      = &viewportStateInfo,
        .pRasterizationState = &rasterizationInfo,
        .pMultisampleState   = &multisampleInfo,
        .pDepthStencilState  = &depthStencilInfo,
        .pColorBlendState    = &colorBlendStateInfo,
        .pDynamicState       = &dynamicStateInfo,
        .layout              = info.pipelineLayout,
    };

    VkPipeline pipeline = VK_NULL_HANDLE;
    VkResult res = VK_SUCCESS;
    if ((res = vkCreateGraphicsPipelines(device, VK_NULL_HANDLE, 1, &pipelineInfo, nullptr, &pipeline)) != VK_SUCCESS)
        return VulkanResult<VkPipeline>::Error(res);

    return VulkanResult<VkPipeline>::Success(pipeline);
}

int main(int argc, char** argv)
{
    (void)argc; (void)argv;

    auto r_window = Window::Create();
    if (!r_window)
        return 1;

    auto r_renderContext = RenderContext::Create("HDR Test");
    if (!r_renderContext)
        return 1;

    auto r_surface = r_window->CreateSurface(r_renderContext->Instance());
    if (!r_surface)
        return 1;

    auto r_swapchain = Swapchain::Create(*r_renderContext, *r_surface, true);
    if (!r_swapchain)
        return 1;

    auto r_buffer = r_renderContext->CreateBuffer(256 * 1024 * 1024);
    if (!r_buffer)
        return 1;
    auto pooler = MemoryPool{ *r_buffer };

    auto r_vsFullscreen = r_renderContext->CreateShader(vs_Fullscreen);
    if (!r_vsFullscreen) return 1;
    auto r_fsHDRTest = r_renderContext->CreateShader(fs_HDRTest);
    if (!r_fsHDRTest) return 1;

    VkSamplerCreateInfo samplerInfo =
    {
        .sType            = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO,
        .magFilter        = VK_FILTER_LINEAR,
        .minFilter        = VK_FILTER_LINEAR,
        .mipmapMode       = VK_SAMPLER_MIPMAP_MODE_LINEAR,
        .addressModeU     = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
        .addressModeV     = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
        .addressModeW     = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
        .mipLodBias       = 0.0f,
        .anisotropyEnable = VK_TRUE,
        .maxAnisotropy    = 16.0f,
        .minLod           = -FLT_MAX,
        .maxLod           = FLT_MAX,
    };
    VkSampler sampler = VK_NULL_HANDLE;
    vkCreateSampler(r_renderContext->Device(), &samplerInfo, nullptr, &sampler);

    struct PushConstants
    {
        uint32_t frameIdx;
        uint32_t currentTestIdx;
        float    time;
        float    targetNits;
    };

    VkPushConstantRange pushConstantRange =
    {
        .stageFlags = VK_SHADER_STAGE_VERTEX_BIT,
        .offset     = 0,
        .size       = sizeof(PushConstants),
    };

    VkPipelineLayoutCreateInfo pipelineLayoutInfo =
    {
        .sType                  = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
        .setLayoutCount         = 0,
        .pSetLayouts            = nullptr,
        .pushConstantRangeCount = 1,
        .pPushConstantRanges    = &pushConstantRange,
    };
    VkPipelineLayout pipelineLayout = VK_NULL_HANDLE;
    vkCreatePipelineLayout(r_renderContext->Device(), &pipelineLayoutInfo, nullptr, &pipelineLayout);

    //
    VkFormat depthFormat = *FindDepthFormat(r_renderContext->PhysicalDevice());
    auto pipeline = *CreateGraphicsPipeline(r_renderContext->Device(), MiniPipelineInfo
    {
        .vs                = *r_vsFullscreen,
        .fs                = *r_fsHDRTest,
        .pipelineLayout    = pipelineLayout,
        .colorFormat       = r_swapchain->Format(),
        .depthFormat       = depthFormat,
        .enableDepthTest   = false,
        .enableDepthWrites = false,
        .vertexAttributes  = Span<const VkVertexInputAttributeDescription>(nullptr),
    });

    struct FrameData
    {
        uint32_t testIdx;
    };

    auto frameDataSlice = *pooler.AllocSlice(MaxFramesInFlight * sizeof(FrameData), 16);
    FrameData* frameData = reinterpret_cast<FrameData*>(frameDataSlice.ptr);

    Input::InputHandler handler;

    r_window->EnableRelativeMouse(false);

    int32_t currentTest = 0;
    auto t1 = Time::now();
    auto startTime = Time::now();
    int32_t stops = 0;
    while (r_window->Update(handler))
    {
        const auto t2 = Time::now();
        auto delta = t2 - t1;
        const float dt = Seconds(delta).count();
        const float time = Seconds(startTime - t2).count();
        t1 = t2;

        if (handler.getFirst(Input::ButtonCodes::Key_Right))
            currentTest++;
        else if (handler.getFirst(Input::ButtonCodes::Key_Left))
            currentTest--;

        if (handler.getFirst(Input::ButtonCodes::Key_Up))
            stops++;
        else if (handler.getFirst(Input::ButtonCodes::Key_Down))
            stops--;

        stops = Clamp(stops, 0, 5);

        float targetNits = 0.0f;
        switch (stops)
        {
            case 0: targetNits = 80.0f; break;
            case 1: targetNits = 100.0f; break;
            case 2: targetNits = 203.0f; break;
            case 3: targetNits = 400.00f; break;
            case 4: targetNits = 1000.00f; break;
            case 5: targetNits = 10000.00f; break;
        }

        VkCommandBuffer cmdBuf = r_swapchain->CommandBuffer();
        r_renderContext->BeginCommandBuffer(cmdBuf);
        {
            VkViewport viewport =
            {
                .x        = 0.0f,
                .y        = 0.0f,
                .width    = float(r_swapchain->Extent().width),
                .height   = float(r_swapchain->Extent().height),
                .minDepth = 0.0f,
                .maxDepth = 1.0f,
            };

            VkRect2D scissor =
            {
                .offset = {0u, 0u},
                .extent = r_swapchain->Extent(),
            };

            vkCmdSetViewport(cmdBuf, 0, 1, &viewport);
            vkCmdSetScissor(cmdBuf, 0, 1, &scissor);

            PushConstants pushConstants =
            {
                .frameIdx = r_swapchain->CurrentFrame(),
                .currentTestIdx = currentTest,
                .time = time,
                .targetNits = targetNits,
            };
            vkCmdPushConstants(cmdBuf, pipelineLayout, VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT, 0, sizeof(PushConstants), &pushConstants);

            const VkImageMemoryBarrier undefinedToColorBarrier =
            {
                .sType            = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
                .dstAccessMask    = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT,
                .oldLayout        = VK_IMAGE_LAYOUT_UNDEFINED,
                .newLayout        = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
                .image            = r_swapchain->Image(),
                .subresourceRange = FirstColorMipSubresourceRange,
            };

            vkCmdPipelineBarrier(
                cmdBuf,
                VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT,
                VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT,
                0,
                0, nullptr,
                0, nullptr,
                1, &undefinedToColorBarrier);

            const VkRenderingAttachmentInfoKHR colorAttachmentInfos[] =
            {
                {
                    .sType       = VK_STRUCTURE_TYPE_RENDERING_ATTACHMENT_INFO_KHR,
                    .imageView   = r_swapchain->ImageView(),
                    .imageLayout = VK_IMAGE_LAYOUT_ATTACHMENT_OPTIMAL_KHR,
                    .loadOp      = VK_ATTACHMENT_LOAD_OP_CLEAR,
                    .storeOp     = VK_ATTACHMENT_STORE_OP_STORE,
                    .clearValue  = { .color = { .float32 = { 1.0f, 1.0f, 1.0f, 1.0f } } },
                },
            };

            const VkRenderingInfo renderInfo =
            {
                .sType                = VK_STRUCTURE_TYPE_RENDERING_INFO,
                .renderArea           = { {}, r_swapchain->Extent() },
                .layerCount           = 1,
                .colorAttachmentCount = Size(colorAttachmentInfos),
                .pColorAttachments    = colorAttachmentInfos,
                .pDepthAttachment     = nullptr,
            };
            vkCmdBeginRendering(cmdBuf, &renderInfo);
            {
                vkCmdBindPipeline(cmdBuf, VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline);
                vkCmdDraw(cmdBuf, 3, 1, 0, 0);
            }
            vkCmdEndRendering(cmdBuf);

            const VkImageMemoryBarrier colorToPresentBarrier =
            {
                .sType            = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
                .srcAccessMask    = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT,
                .oldLayout        = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
                .newLayout        = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR,
                .image            = r_swapchain->Image(),
                .subresourceRange = FirstColorMipSubresourceRange,
            };

            vkCmdPipelineBarrier(
                cmdBuf,
                VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT,
                VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT,
                0,
                0, nullptr,
                0, nullptr,
                1, &colorToPresentBarrier);
        }
        r_renderContext->EndCommandBuffer(cmdBuf);
        r_swapchain->Present();
    }

    return 0;
}