// =============================================
//  Aster: model_render.cpp
//  Copyright (c) 2020-2024 Anish Bhobe
// =============================================

#include "buffer.h"
#include "constants.h"
#include "context.h"
#include "device.h"
#include "global.h"
#include "image.h"
#include "physical_device.h"
#include "pipeline.h"
#include "swapchain.h"
#include "window.h"

#include "frame.h"
#include "helpers.h"
#include "light_manager.h"

#include "gpu_resource_manager.h"
#include "gui.h"
#include "model_loader.h"
#include "pipeline_utils.h"

#include <EASTL/array.h>
#include <stb_image.h>
#include <tiny_gltf.h>
#include <filesystem>

constexpr u32 MAX_FRAMES_IN_FLIGHT = 3;
constexpr auto MODEL_FILE = "model/DamagedHelmet.glb";

struct Camera
{
    mat4 m_View;
    mat4 m_Perspective;
    vec4 m_Position;
};

int
main(int, char **)
{
    MIN_LOG_LEVEL(Logger::LogType::eInfo);

    Context context = {"ModelRender [WIP]", VERSION};
    Window window = {"ModelRender [WIP] (Aster)", &context, {640, 480}};

    PhysicalDevices physicalDevices = {&window, &context};
    PhysicalDevice deviceToUse = FindSuitableDevice(physicalDevices);

    INFO("Using {} as the primary device.", deviceToUse.m_DeviceProperties.deviceName.data());

    Features enabledDeviceFeatures = {
        .m_Vulkan10Features = {.samplerAnisotropy = true},
        .m_Vulkan12Features =
            {
                .descriptorIndexing = true,
                .shaderSampledImageArrayNonUniformIndexing = true,
                .shaderStorageBufferArrayNonUniformIndexing = true,
                .shaderStorageImageArrayNonUniformIndexing = true,
                .descriptorBindingUniformBufferUpdateAfterBind = true, // Not related to Bindless
                .descriptorBindingSampledImageUpdateAfterBind = true,
                .descriptorBindingStorageImageUpdateAfterBind = true,
                .descriptorBindingStorageBufferUpdateAfterBind = true,
                .descriptorBindingPartiallyBound = true,
                .runtimeDescriptorArray = true,
            },
        .m_Vulkan13Features = {.dynamicRendering = true},
    };

    QueueAllocation queueAllocation = FindAppropriateQueueAllocation(&deviceToUse);
    Device device = {&context, &deviceToUse, &enabledDeviceFeatures, {queueAllocation}, "Primary Device"};
    vk::Queue commandQueue = device.GetQueue(queueAllocation.m_Family, 0);
    Swapchain swapchain = {&window, &device, "Primary Chain"};
    GpuResourceManager resourceManager = {&device, 1000};

    ModelLoader modelLoader = {&resourceManager, commandQueue, queueAllocation.m_Family, queueAllocation.m_Family};
    LightManager lightManager = LightManager{&resourceManager};

    Model model = modelLoader.LoadModel(MODEL_FILE);

    vk::Format attachmentFormat = vk::Format::eR8G8B8A8Srgb;
    Pipeline pipeline = CreatePipeline(&device, attachmentFormat, &resourceManager);

    Camera camera = {
        .m_View = glm::lookAt(vec3(0.0f, 2.0f, 2.0f), vec3(0.0f), vec3(0.0f, 1.0f, 0.0f)),
        .m_Perspective = glm::perspective(
            70_deg, Cast<f32>(swapchain.m_Extent.width) / Cast<f32>(swapchain.m_Extent.height), 0.1f, 100.0f),
        .m_Position = vec4{0.0f, 2.0f, 2.0f, 1.0f},
    };

    lightManager.AddDirectional(vec3(0.0f, -1.0f, 0.0f), {0.0f, 1.0f, 0.0f});
    lightManager.AddPoint(vec3{2.0f, 1.0f, 0.0f}, {1.0f, 0.0f, 0.0f}, 15.0f);
    lightManager.AddPoint(vec3{-2.0f, 1.0f, 0.0f}, {0.0f, 0.0f, 1.0f}, 15.0f);

    lightManager.Update();

    vk::DescriptorPool descriptorPool;
    vk::DescriptorSet descriptorSet;

    {
        vk::DescriptorSetLayout descriptorSetLayout = pipeline.m_SetLayouts[1];
        eastl::array poolSizes = {
            vk::DescriptorPoolSize{
                .type = vk::DescriptorType::eUniformBuffer,
                .descriptorCount = 1,
            },
        };
        vk::DescriptorPoolCreateInfo descriptorPoolCreateInfo = {
            .maxSets = 1, .poolSizeCount = Cast<u32>(poolSizes.size()), .pPoolSizes = poolSizes.data()};
        AbortIfFailed(device.m_Device.createDescriptorPool(&descriptorPoolCreateInfo, nullptr, &descriptorPool));

        vk::DescriptorSetAllocateInfo descriptorSetAllocateInfo = {
            .descriptorPool = descriptorPool,
            .descriptorSetCount = 1,
            .pSetLayouts = &descriptorSetLayout,
        };
        AbortIfFailed(device.m_Device.allocateDescriptorSets(&descriptorSetAllocateInfo, &descriptorSet));
    }

    UniformBuffer ubo;
    ubo.Init(&device, sizeof camera, "Camera UBO");
    ubo.Write(&device, 0, sizeof camera, &camera);
    vk::DescriptorBufferInfo descriptorBufferInfo = {
        .buffer = ubo.m_Buffer,
        .offset = 0,
        .range = ubo.GetSize(),
    };
    eastl::array writeDescriptors = {
        vk::WriteDescriptorSet{
            .dstSet = descriptorSet,
            .dstBinding = 0,
            .dstArrayElement = 0,
            .descriptorCount = 1,
            .descriptorType = vk::DescriptorType::eUniformBuffer,
            .pBufferInfo = &descriptorBufferInfo,
        },
    };
    device.m_Device.updateDescriptorSets(Cast<u32>(writeDescriptors.size()), writeDescriptors.data(), 0, nullptr);

    resourceManager.Update();

    // Persistent variables
    vk::Viewport viewport = {
        .x = 0,
        .y = Cast<f32>(swapchain.m_Extent.height),
        .width = Cast<f32>(swapchain.m_Extent.width),
        .height = -Cast<f32>(swapchain.m_Extent.height),
        .minDepth = 0.0,
        .maxDepth = 1.0,
    };

    vk::Rect2D scissor = {
        .offset = {0, 0},
        .extent = swapchain.m_Extent,
    };

    vk::ImageSubresourceRange subresourceRange = {
        .aspectMask = vk::ImageAspectFlagBits::eColor,
        .baseMipLevel = 0,
        .levelCount = 1,
        .baseArrayLayer = 0,
        .layerCount = 1,
    };
    vk::ImageMemoryBarrier topOfThePipeBarrier = {
        .oldLayout = vk::ImageLayout::eUndefined,
        .newLayout = vk::ImageLayout::eColorAttachmentOptimal,
        .srcQueueFamilyIndex = queueAllocation.m_Family,
        .dstQueueFamilyIndex = queueAllocation.m_Family,
        .subresourceRange = subresourceRange,
    };
    vk::ImageMemoryBarrier renderToTransferSrcBarrier = {
        .oldLayout = vk::ImageLayout::eColorAttachmentOptimal,
        .newLayout = vk::ImageLayout::eTransferSrcOptimal,
        .srcQueueFamilyIndex = queueAllocation.m_Family,
        .dstQueueFamilyIndex = queueAllocation.m_Family,
        .subresourceRange = subresourceRange,
    };
    vk::ImageMemoryBarrier acquireToTransferDstBarrier = {
        .oldLayout = vk::ImageLayout::eUndefined,
        .newLayout = vk::ImageLayout::eTransferDstOptimal,
        .srcQueueFamilyIndex = queueAllocation.m_Family,
        .dstQueueFamilyIndex = queueAllocation.m_Family,
        .subresourceRange = subresourceRange,
    };
    eastl::array postRenderBarriers = {
        renderToTransferSrcBarrier,
        acquireToTransferDstBarrier,
    };
    vk::ImageMemoryBarrier transferDstToPresentBarrier = {
        .oldLayout = vk::ImageLayout::eTransferDstOptimal,
        .newLayout = vk::ImageLayout::ePresentSrcKHR,
        .srcQueueFamilyIndex = queueAllocation.m_Family,
        .dstQueueFamilyIndex = queueAllocation.m_Family,
        .subresourceRange = subresourceRange,
    };

    FrameManager frameManager = {&device, queueAllocation.m_Family, MAX_FRAMES_IN_FLIGHT};
    eastl::fixed_vector<DepthImage, MAX_FRAMES_IN_FLIGHT> depthImages(frameManager.m_FramesInFlight);
    eastl::fixed_vector<AttachmentImage, MAX_FRAMES_IN_FLIGHT> attachmentImages(frameManager.m_FramesInFlight);
    {
        auto depthIter = depthImages.begin();
        auto attachmentIter = attachmentImages.begin();
        for (u32 index = 0; index < frameManager.m_FramesInFlight; ++index)
        {
            auto name = fmt::format("Depth Frame{}", index);
            depthIter->Init(&device, swapchain.m_Extent, name.c_str());

            name = fmt::format("Attachment0 Frame{}", index);
            attachmentIter->Init(&device, swapchain.m_Extent, attachmentFormat, name.c_str());

            ++depthIter;
            ++attachmentIter;
        }
    }

    gui::Init(&context, &device, &window, attachmentFormat, frameManager.m_FramesInFlight, queueAllocation.m_Family,
              commandQueue);
    bool rotating = false;

    Time::Init();

    INFO("Starting loop");
    while (window.Poll())
    {
        Time::Update();

        gui::StartBuild();

        gui::Begin("Settings");
        gui::Text("Delta: %0.6f ms", 1000.0f * Time::m_Delta);
        gui::Text("FPS: %0.6f", 1.0f / Time::m_Delta);
        gui::Checkbox("Rotate", &rotating);
        if (gui::Button("Exit"))
        {
            window.RequestExit();
        }
        gui::End();

        gui::EndBuild();

        if (rotating)
        {
            model.SetModelTransform(
                rotate(model.GetModelTransform(), Cast<f32>(45.0_deg * Time::m_Delta), vec3(0.0f, 1.0f, 0.0f)));
        }
        model.Update();
        ubo.Write(&device, 0, sizeof camera, &camera);

        Frame *currentFrame = frameManager.GetNextFrame(&swapchain, &window);

        u32 imageIndex = currentFrame->m_ImageIdx;
        vk::Image currentSwapchainImage = swapchain.m_Images[imageIndex];
        vk::CommandBuffer cmd = currentFrame->m_CommandBuffer;
        vk::ImageView currentDepthImageView = depthImages[currentFrame->m_FrameIdx].m_View;

        AttachmentImage *currentAttachment = &attachmentImages[currentFrame->m_ImageIdx];
        vk::Image currentImage = currentAttachment->m_Image;
        vk::ImageView currentImageView = currentAttachment->m_View;

        topOfThePipeBarrier.image = currentImage;
        postRenderBarriers[0].image = currentImage;
        postRenderBarriers[1].image = currentSwapchainImage;
        transferDstToPresentBarrier.image = currentSwapchainImage;

        vk::CommandBufferBeginInfo beginInfo = {.flags = vk::CommandBufferUsageFlagBits::eOneTimeSubmit};
        AbortIfFailed(cmd.begin(&beginInfo));

        cmd.pipelineBarrier(vk::PipelineStageFlagBits::eTopOfPipe, vk::PipelineStageFlagBits::eColorAttachmentOutput,
                            {}, 0, nullptr, 0, nullptr, 1, &topOfThePipeBarrier);

        // Render
        eastl::array attachmentInfos = {
            vk::RenderingAttachmentInfo{
                .imageView = currentImageView,
                .imageLayout = vk::ImageLayout::eColorAttachmentOptimal,
                .resolveMode = vk::ResolveModeFlagBits::eNone,
                .loadOp = vk::AttachmentLoadOp::eClear,
                .storeOp = vk::AttachmentStoreOp::eStore,
                .clearValue = vk::ClearColorValue{0.0f, 0.0f, 0.0f, 1.0f},
            },
        };

        vk::RenderingAttachmentInfo depthAttachment = {
            .imageView = currentDepthImageView,
            .imageLayout = vk::ImageLayout::eDepthAttachmentOptimal,
            .resolveMode = vk::ResolveModeFlagBits::eNone,
            .loadOp = vk::AttachmentLoadOp::eClear,
            .storeOp = vk::AttachmentStoreOp::eDontCare,
            .clearValue = vk::ClearDepthStencilValue{.depth = 1.0f, .stencil = 0},
        };

        vk::RenderingInfo renderingInfo = {
            .renderArea = {.extent = ToExtent2D(currentAttachment->m_Extent)},
            .layerCount = 1,
            .colorAttachmentCount = Cast<u32>(attachmentInfos.size()),
            .pColorAttachments = attachmentInfos.data(),
            .pDepthAttachment = &depthAttachment,
        };

        cmd.beginRendering(&renderingInfo);

        cmd.setViewport(0, 1, &viewport);
        cmd.setScissor(0, 1, &scissor);
        cmd.bindPipeline(vk::PipelineBindPoint::eGraphics, pipeline.m_Pipeline);
        cmd.bindDescriptorSets(vk::PipelineBindPoint::eGraphics, pipeline.m_Layout, 0, 1,
                               &resourceManager.m_DescriptorSet, 0, nullptr);
        cmd.bindDescriptorSets(vk::PipelineBindPoint::eGraphics, pipeline.m_Layout, 1, 1, &descriptorSet, 0, nullptr);

        cmd.bindIndexBuffer(model.m_IndexBuffer.m_Buffer, 0, vk::IndexType::eUint32);

        u32 pcbOffset = 0;
        cmd.pushConstants(pipeline.m_Layout, vk::ShaderStageFlagBits::eAll, pcbOffset, sizeof model.m_Handles,
                          &model.m_Handles);
        pcbOffset += sizeof model.m_Handles;
        cmd.pushConstants(pipeline.m_Layout, vk::ShaderStageFlagBits::eAll, pcbOffset, sizeof lightManager.m_MetaInfo,
                          &lightManager.m_MetaInfo);
        pcbOffset += sizeof lightManager.m_MetaInfo;

        for (auto &prim : model.m_MeshPrimitives)
        {
            u32 innerPcbOffset = pcbOffset;
            cmd.pushConstants(pipeline.m_Layout, vk::ShaderStageFlagBits::eAll, innerPcbOffset,
                              sizeof prim.m_MaterialIdx, &prim.m_MaterialIdx);
            innerPcbOffset += sizeof prim.m_MaterialIdx;
            cmd.pushConstants(pipeline.m_Layout, vk::ShaderStageFlagBits::eAll, innerPcbOffset,
                              sizeof prim.m_TransformIdx, &prim.m_TransformIdx);
            innerPcbOffset += sizeof prim.m_TransformIdx;
            cmd.drawIndexed(prim.m_IndexCount, 1, prim.m_FirstIndex, Cast<i32>(prim.m_VertexOffset), 0);
        }

        cmd.endRendering();

        gui::Draw(cmd, currentAttachment);

        cmd.pipelineBarrier(vk::PipelineStageFlagBits::eColorAttachmentOutput, vk::PipelineStageFlagBits::eAllCommands,
                            {}, 0, nullptr, 0, nullptr, postRenderBarriers.size(), postRenderBarriers.data());

        vk::ImageBlit blitRegion = {
            .srcSubresource =
                {
                    .aspectMask = vk::ImageAspectFlagBits::eColor,
                    .mipLevel = 0,
                    .baseArrayLayer = 0,
                    .layerCount = 1,
                },
            .srcOffsets =
                std::array{
                    vk::Offset3D{0, 0, 0},
                    ToOffset3D(currentAttachment->m_Extent),
                },
            .dstSubresource =
                {
                    .aspectMask = vk::ImageAspectFlagBits::eColor,
                    .mipLevel = 0,
                    .baseArrayLayer = 0,
                    .layerCount = 1,
                },
            .dstOffsets =
                std::array{
                    vk::Offset3D{0, 0, 0},
                    vk::Offset3D{Cast<i32>(swapchain.m_Extent.width), Cast<i32>(swapchain.m_Extent.height), 1},
                },
        };
        cmd.blitImage(currentImage, postRenderBarriers[0].newLayout, currentSwapchainImage,
                      postRenderBarriers[1].newLayout, 1, &blitRegion, vk::Filter::eLinear);

        cmd.pipelineBarrier(vk::PipelineStageFlagBits::eAllCommands, vk::PipelineStageFlagBits::eAllCommands, {}, 0,
                            nullptr, 0, nullptr, 1, &transferDstToPresentBarrier);

        AbortIfFailed(cmd.end());

        vk::PipelineStageFlags waitDstStage = vk::PipelineStageFlagBits::eColorAttachmentOutput;
        vk::SubmitInfo submitInfo = {
            .waitSemaphoreCount = 1,
            .pWaitSemaphores = &currentFrame->m_ImageAcquireSem,
            .pWaitDstStageMask = &waitDstStage,
            .commandBufferCount = 1,
            .pCommandBuffers = &cmd,
            .signalSemaphoreCount = 1,
            .pSignalSemaphores = &currentFrame->m_RenderFinishSem,
        };
        AbortIfFailed(commandQueue.submit(1, &submitInfo, currentFrame->m_FrameAvailableFence));

        currentFrame->Present(commandQueue, &swapchain, &window);
    }

    AbortIfFailed(device.m_Device.waitIdle());

    gui::Destroy(&device);

    for (auto &depthImage : depthImages)
    {
        depthImage.Destroy(&device);
    }
    for (auto &attachmentImage : attachmentImages)
    {
        attachmentImage.Destroy(&device);
    }
    ubo.Destroy(&device);
    device.m_Device.destroy(descriptorPool, nullptr);

    return 0;
}