#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 <Orange/Render/SmallPipeline.h>

#include <vs_Fullscreen.h>
#include <fs_ColorTester.h>
#include <fs_ColorTester.h>

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

using namespace orange;

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;

    const bool srgbViews = false;
    bool swapchainWasHDR = true;//false;
    auto r_swapchain = Swapchain::Create(*r_renderContext, *r_surface, swapchainWasHDR, srgbViews);
    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_fsColorTester = r_renderContext->CreateShader(fs_ColorTester);
    if (!r_fsColorTester) 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);

    constexpr uint32_t MaxBindlessResources = 32768;
    VkDescriptorPoolSize bindlessPoolSizes[] =
    {
        { VK_DESCRIPTOR_TYPE_SAMPLER,                1u},
        { VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, MaxBindlessResources },
    };

    VkDescriptorPoolCreateInfo poolInfo =
    {
        .sType         = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
        .flags         = VK_DESCRIPTOR_POOL_CREATE_UPDATE_AFTER_BIND_BIT,
        .maxSets       = 1u,
        .poolSizeCount = uint32_t(Size(bindlessPoolSizes)),
        .pPoolSizes    = bindlessPoolSizes,
    };
    
    VkDescriptorPool descriptorPool = VK_NULL_HANDLE;
    vkCreateDescriptorPool(r_renderContext->Device(), &poolInfo, nullptr, &descriptorPool);

    VkDescriptorSetLayoutBinding layoutBindings[] =
    {
        {
            .binding            = 0,
            .descriptorType     = VK_DESCRIPTOR_TYPE_SAMPLER,
            .descriptorCount    = 1,
            .stageFlags         = VK_SHADER_STAGE_FRAGMENT_BIT,
            .pImmutableSamplers = &sampler,
        },
        {
            .binding            = 1,
            .descriptorType     = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE,
            .descriptorCount    = MaxBindlessResources,
            .stageFlags         = VK_SHADER_STAGE_FRAGMENT_BIT,
        },
    };

    constexpr VkDescriptorBindingFlags bindingFlags[] =
    {
        0,
        VK_DESCRIPTOR_BINDING_PARTIALLY_BOUND_BIT_EXT | VK_DESCRIPTOR_BINDING_VARIABLE_DESCRIPTOR_COUNT_BIT_EXT | VK_DESCRIPTOR_BINDING_UPDATE_AFTER_BIND_BIT_EXT,
    };

    Assert(Size(bindingFlags) == Size(layoutBindings));

    VkDescriptorSetLayoutBindingFlagsCreateInfo layoutBindingFlagsInfo =
    {
        .sType         = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_BINDING_FLAGS_CREATE_INFO_EXT,
        .bindingCount  = Size(bindingFlags),
        .pBindingFlags = bindingFlags,
    };

    VkDescriptorSetLayoutCreateInfo layoutInfo =
    {
        .sType        = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
        .pNext        = &layoutBindingFlagsInfo,
        .flags        = VK_DESCRIPTOR_SET_LAYOUT_CREATE_UPDATE_AFTER_BIND_POOL_BIT_EXT,
        .bindingCount = Size(layoutBindings),
        .pBindings    = layoutBindings,
    };

    VkDescriptorSetLayout descriptorSetLayout = VK_NULL_HANDLE;
    vkCreateDescriptorSetLayout(r_renderContext->Device(), &layoutInfo, nullptr, &descriptorSetLayout);

    VkDescriptorSetVariableDescriptorCountAllocateInfoEXT descriptorCountAllocInfo =
    {
        .sType              = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_VARIABLE_DESCRIPTOR_COUNT_ALLOCATE_INFO_EXT,
        .descriptorSetCount = 1,
        .pDescriptorCounts  = &MaxBindlessResources,
    };

    VkDescriptorSetAllocateInfo descriptorSetAllocInfo =
    {
        .sType              = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
        .pNext              = &descriptorCountAllocInfo,
        .descriptorPool     = descriptorPool,
        .descriptorSetCount = 1,
        .pSetLayouts        = &descriptorSetLayout,
    };

    VkDescriptorSet descriptorSet = VK_NULL_HANDLE;
    vkAllocateDescriptorSets(r_renderContext->Device(), &descriptorSetAllocInfo, &descriptorSet);

    struct PushConstants
    {
        uint32_t frameIdx;
        uint32_t imageIdx;
        float    time;
    };

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

    VkPipelineLayoutCreateInfo pipelineLayoutInfo =
    {
        .sType                  = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
        .setLayoutCount         = 1,
        .pSetLayouts            = &descriptorSetLayout,
        .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_fsColorTester,
        .pipelineLayout    = pipelineLayout,
        .colorFormat       = r_swapchain->Format(),
        .depthFormat       = depthFormat,
        .enableDepthTest   = false,
        .enableDepthWrites = false,
        .vertexAttributes  = Span<const VkVertexInputAttributeDescription>(nullptr),
    });


    uint32_t totalImageCount = 0;
    struct ImageUpload
    {
        BufferSlice from;
        VkImage     to;
        VkExtent3D  extent;
    };
    Vector<ImageUpload> imageUploads;
    Vector<bool> imageIsHDR;
    if (!std::filesystem::exists("images")) {
        log::err("No image directory!");
        return 1;
    }

    {
        Vector<std::string> filenames;

        for (const auto &entry : std::filesystem::directory_iterator("images"))
            filenames.push_back(std::string(entry.path()));

        std::sort(filenames.begin(), filenames.end(), [](const auto& a, const auto& b) { return a < b; });

        for (auto& entry : filenames)
        {
            std::cout << entry << std::endl;

            bool imageHDR = true;//std::string(entry.path()).find(".pq.") != std::string::npos;

            auto r_pngData = fs::OpenFileIntoBuffer(entry.c_str());
            if (!r_pngData)
                continue;

            spng_ctx *ctx = spng_ctx_new(0);
            spng_set_png_buffer(ctx, r_pngData->data, r_pngData->size);

            struct spng_ihdr ihdr;
            spng_get_ihdr(ctx, &ihdr);

            auto pngFmt = ihdr.bit_depth == 16 ? SPNG_FMT_RGBA16 : SPNG_FMT_RGBA8;
            auto vkFormat = ihdr.bit_depth == 16 ? VK_FORMAT_R16G16B16A16_UNORM : VK_FORMAT_R8G8B8A8_UNORM;

            size_t pngOutSize = 0zu;
            spng_decoded_image_size(ctx, pngFmt, &pngOutSize);

            std::cout << pngOutSize / 1024 / 1024 << "mb" << std::endl;

            auto pngLinearSlice = *pooler.AllocSlice(pngOutSize, 16);
            spng_decode_image(ctx, pngLinearSlice.ptr, pngOutSize, pngFmt, 0);
            spng_ctx_free(ctx);

            auto r_texture = r_renderContext->CreateImage(pooler, ihdr.width, ihdr.height, vkFormat, VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT);
            if (!r_texture)
                return 1;
            auto r_textureView = r_renderContext->CreateImageView(*r_texture, vkFormat, VK_IMAGE_ASPECT_COLOR_BIT);
            if (!r_textureView)
                return 1;

            VkDescriptorImageInfo descriptorImageInfo =
            {
                .imageView   = *r_textureView,
                .imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
            };

            VkWriteDescriptorSet writeImageDescriptorSet[] =
            {
                {
                    .sType           = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
                    .dstSet          = descriptorSet,
                    .dstBinding      = 1,
                    .dstArrayElement = totalImageCount++,
                    .descriptorCount = 1,
                    .descriptorType  = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE,
                    .pImageInfo      = &descriptorImageInfo,
                },
            };
            vkUpdateDescriptorSets(r_renderContext->Device(), Size(writeImageDescriptorSet), writeImageDescriptorSet, 0, nullptr);

            imageUploads.push_back(ImageUpload
            {
                .from   = pngLinearSlice,
                .to     = *r_texture,
                .extent = VkExtent3D{ ihdr.width, ihdr.height, 1u },
            });

            imageIsHDR.push_back(imageHDR);
        }
    }

    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 currentImageIdx = 0;
    auto t1 = Time::now();
    auto startTime = Time::now();
    uint32_t frameCounter = 0;
    while (r_window->Update(handler))
    {
        frameCounter++;

        const auto t2 = Time::now();
        auto delta = t2 - t1;
        const float time = Seconds(startTime - t2).count();
        t1 = t2;

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

        currentImageIdx = currentImageIdx % totalImageCount;

        bool currentImageIsHDR = imageIsHDR[currentImageIdx];

#if 0
        if (swapchainWasHDR != currentImageIsHDR)
        {
            r_swapchain.Destroy();

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

            swapchainWasHDR = currentImageIsHDR;
        }
#endif

        VkCommandBuffer cmdBuf = r_swapchain->CommandBuffer();
        r_renderContext->BeginCommandBuffer(cmdBuf);
        {
            for (auto& upload : imageUploads)
            {
                VkImageMemoryBarrier transferBarrier =
                {
                    .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
                    .srcAccessMask = 0,
                    .dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT,
                    .oldLayout     = VK_IMAGE_LAYOUT_UNDEFINED,
                    .newLayout     = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
                    .image         = upload.to,
                    .subresourceRange = FirstColorMipSubresourceRange,
                };
                vkCmdPipelineBarrier(cmdBuf, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 0, nullptr, 0, nullptr, 1, &transferBarrier);

                VkBufferImageCopy region =
                {
                    .bufferOffset = upload.from.offset,
                    .imageSubresource =
                    {
                        .aspectMask     = VK_IMAGE_ASPECT_COLOR_BIT,
                        .mipLevel       = 0u,
                        .baseArrayLayer = 0u,
                        .layerCount     = 1u,
                    },
                    .imageExtent = upload.extent,
                };
                vkCmdCopyBufferToImage(cmdBuf, upload.from.buffer, upload.to, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &region);

                VkImageMemoryBarrier optimalBarrier =
                {
                    .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
                    .srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT,
                    .dstAccessMask = VK_ACCESS_MEMORY_READ_BIT,
                    .oldLayout     = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
                    .newLayout     = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
                    .image         = upload.to,
                    .subresourceRange = FirstColorMipSubresourceRange,
                };
                vkCmdPipelineBarrier(cmdBuf, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, 0, 0, nullptr, 0, nullptr, 1, &optimalBarrier);
            }
            imageUploads.clear();

            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 = frameCounter,
                .imageIdx = currentImageIdx,
                .time = time,
            };
            vkCmdPushConstants(cmdBuf, pipelineLayout, VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT, 0, sizeof(PushConstants), &pushConstants);

            vkCmdBindDescriptorSets(cmdBuf, VK_PIPELINE_BIND_POINT_GRAPHICS, pipelineLayout, 0, 1, &descriptorSet, 0, nullptr);

            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;
}