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

#include "buffer.h"
#include "constants.h"
#include "context.h"
#include "device.h"
#include "physical_device.h"
#include "window.h"

#include "global.h"
#include "pipeline.h"
#include "swapchain.h"

#include "helpers.h"

#include <EASTL/array.h>

constexpr u32 MAX_FRAMES_IN_FLIGHT = 3;
constexpr auto VERTEX_SHADER_FILE = "shader/triangle.vert.glsl.spv";
constexpr auto FRAGMENT_SHADER_FILE = "shader/triangle.frag.glsl.spv";

vk::ShaderModule CreateShader(const Device *device, cstr shaderFile);
Pipeline CreatePipeline(const Device *device, const Swapchain *swapchain);

struct Vertex
{
    vec3 m_Position;
    vec3 m_Color;

    constexpr static vk::VertexInputBindingDescription
    GetBinding(const u32 binding)
    {
        return {.binding = binding, .stride = sizeof(Vertex), .inputRate = vk::VertexInputRate::eVertex};
    }

    constexpr static eastl::array<vk::VertexInputAttributeDescription, 2>
    GetAttributes(const u32 binding)
    {
        return {
            vk::VertexInputAttributeDescription{
                .location = 0,
                .binding = binding,
                .format = vk::Format::eR32G32B32Sfloat,
                .offset = offsetof(Vertex, m_Position),
            },
            vk::VertexInputAttributeDescription{
                .location = 1,
                .binding = binding,
                .format = vk::Format::eR32G32B32Sfloat,
                .offset = offsetof(Vertex, m_Color),
            },
        };
    }
};

struct Frame
{
    const Device *m_Device;
    vk::CommandPool m_Pool;
    vk::CommandBuffer m_CommandBuffer;
    vk::Fence m_FrameAvailableFence;
    vk::Semaphore m_ImageAcquireSem;
    vk::Semaphore m_RenderFinishSem;

    Frame(const Device *device, u32 queueFamilyIndex, u32 frameCount);
    ~Frame();
};

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

    Context context = {"Triangle", VERSION};
    Window window = {"Triangle (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_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"};
    Pipeline pipeline = CreatePipeline(&device, &swapchain);

    vk::CommandPool copyPool;
    vk::CommandBuffer copyBuffer;
    {
        vk::CommandPoolCreateInfo poolCreateInfo = {
            .flags = vk::CommandPoolCreateFlagBits::eTransient,
            .queueFamilyIndex = queueAllocation.m_Family,
        };
        auto result = device.m_Device.createCommandPool(&poolCreateInfo, nullptr, &copyPool);
        ERROR_IF(Failed(result), "Copy command pool creation failed. Cause: {}", result) THEN_ABORT(result);
        vk::CommandBufferAllocateInfo bufferAllocateInfo = {
            .commandPool = copyPool,
            .level = vk::CommandBufferLevel::ePrimary,
            .commandBufferCount = 1,
        };
        result = device.m_Device.allocateCommandBuffers(&bufferAllocateInfo, &copyBuffer);
        ERROR_IF(Failed(result), "Copy command buffer allocation failed. Cause: {}", result) THEN_ABORT(result);
    }

    // eastl::array<Vertex, 3> vertices{};
    eastl::array vertices = {
        Vertex{.m_Position = {-0.5f, -0.5f, 0.0f}, .m_Color = {1.0f, 0.0f, 0.0f}},
        Vertex{.m_Position = {0.5f, -0.5f, 0.0f}, .m_Color = {0.0f, 1.0f, 0.0f}},
        Vertex{.m_Position = {0.0f, 0.5f, 0.0f}, .m_Color = {0.0f, 0.0f, 1.0f}},
    };
    VertexBuffer vbo;
    vbo.Init(&device, vertices.size() * sizeof vertices[0], "VBO");
    {
        StagingBuffer staging;
        staging.Init(&device, vertices.size() * sizeof vertices[0], "Staging");
        staging.Write(&device, 0, vertices.size() * sizeof vertices[0], vertices.data());

        vk::Fence fence;
        vk::FenceCreateInfo fenceCreateInfo = {};
        auto result = device.m_Device.createFence(&fenceCreateInfo, nullptr, &fence);
        ERROR_IF(Failed(result), "Fence creation failed. Cause: {}", result) THEN_ABORT(result);

        vk::CommandBufferBeginInfo beginInfo = {.flags = vk::CommandBufferUsageFlagBits::eOneTimeSubmit};
        result = copyBuffer.begin(&beginInfo);
        ERROR_IF(Failed(result), "Copy begin failed. Cause: {}", result) THEN_ABORT(result);

        vk::BufferCopy bufferCopy = {.srcOffset = 0, .dstOffset = 0, .size = staging.GetSize()};
        copyBuffer.copyBuffer(staging.m_Buffer, vbo.m_Buffer, 1, &bufferCopy);

        result = copyBuffer.end();
        ERROR_IF(Failed(result), "Copy end failed. Cause: {}", result) THEN_ABORT(result);

        vk::SubmitInfo submitInfo = {
            .commandBufferCount = 1,
            .pCommandBuffers = &copyBuffer,
        };

        result = commandQueue.submit(1, &submitInfo, fence);
        ERROR_IF(Failed(result), "Submit failed. Cause: {}", result) THEN_ABORT(result) ELSE_INFO("Submit copy");

        result = device.m_Device.waitForFences(1, &fence, true, MaxValue<u64>);
        ERROR_IF(Failed(result), "Fence wait failed. Cause: {}", result) THEN_ABORT(result) ELSE_INFO("Fence wait");

        result = device.m_Device.resetCommandPool(copyPool, {});
        ERROR_IF(Failed(result), "Couldn't reset command pool. Cause: {}", result) THEN_ABORT(result);

        device.m_Device.destroy(fence, nullptr);
        staging.Destroy(&device);
    }

    // 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 renderToPresentBarrier = {
        .oldLayout = vk::ImageLayout::eColorAttachmentOptimal,
        .newLayout = vk::ImageLayout::ePresentSrcKHR,
        .srcQueueFamilyIndex = queueAllocation.m_Family,
        .dstQueueFamilyIndex = queueAllocation.m_Family,
        .subresourceRange = subresourceRange,
    };

    // Frames
    eastl::fixed_vector<Frame, MAX_FRAMES_IN_FLIGHT> frames;
    for (int i = 0; i < MAX_FRAMES_IN_FLIGHT; ++i)
    {
        frames.emplace_back(&device, queueAllocation.m_Family, i);
    }

    INFO("Starting loop");
    u32 frameIndex = 0;
    while (window.Poll())
    {
        Frame *currentFrame = &frames[frameIndex];

        auto result = device.m_Device.waitForFences(1, &currentFrame->m_FrameAvailableFence, true, MaxValue<u64>);
        ERROR_IF(Failed(result), "Waiting for fence {} failed. Cause: {}", frameIndex, result)
        THEN_ABORT(result);

        u32 imageIndex;
        result = device.m_Device.acquireNextImageKHR(swapchain.m_Swapchain, MaxValue<u64>,
                                                     currentFrame->m_ImageAcquireSem, nullptr, &imageIndex);
        if (Failed(result))
        {
            switch (result)
            {
            case vk::Result::eErrorOutOfDateKHR:
            case vk::Result::eSuboptimalKHR:
                INFO("Recreating Swapchain. Cause: {}", result);
                swapchain.Create(&window);
                viewport.y = Cast<f32>(swapchain.m_Extent.height);
                viewport.width = Cast<f32>(swapchain.m_Extent.width);
                viewport.height = -Cast<f32>(swapchain.m_Extent.height);
                scissor.extent = swapchain.m_Extent;
                continue; // Image acquire has failed. We move to the next frame.
            default:
                ERROR("Waiting for swapchain image {} failed. Cause: {}", frameIndex, result)
                THEN_ABORT(result);
            }
        }
        // Reset fences here. In case swapchain was out of date, we leave the fences signalled.
        result = device.m_Device.resetFences(1, &currentFrame->m_FrameAvailableFence);
        ERROR_IF(Failed(result), "Fence {} reset failed. Cause: {}", frameIndex, result)
        THEN_ABORT(result);

        result = device.m_Device.resetCommandPool(currentFrame->m_Pool, {});
        ERROR_IF(Failed(result), "Command pool {} reset failed. Cause: {}", frameIndex, result)
        THEN_ABORT(result);

        vk::ImageView currentImageView = swapchain.m_ImageViews[imageIndex];
        vk::Image currentImage = swapchain.m_Images[imageIndex];
        vk::CommandBuffer cmd = currentFrame->m_CommandBuffer;

        topOfThePipeBarrier.image = currentImage;
        renderToPresentBarrier.image = currentImage;

        vk::CommandBufferBeginInfo beginInfo = {.flags = vk::CommandBufferUsageFlagBits::eOneTimeSubmit};
        result = cmd.begin(&beginInfo);
        ERROR_IF(Failed(result), "Command buffer begin failed. Cause: {}", result)
        THEN_ABORT(result);

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

        // Render
        vk::RenderingAttachmentInfo attachmentInfo = {
            .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::RenderingInfo renderingInfo = {
            .renderArea = {.extent = swapchain.m_Extent},
            .layerCount = 1,
            .colorAttachmentCount = 1,
            .pColorAttachments = &attachmentInfo,
        };

        cmd.beginRendering(&renderingInfo);

        cmd.setViewport(0, 1, &viewport);
        cmd.setScissor(0, 1, &scissor);
        cmd.bindPipeline(vk::PipelineBindPoint::eGraphics, pipeline.m_Pipeline);
        usize offsets = 0;
        cmd.bindVertexBuffers(0, 1, &vbo.m_Buffer, &offsets);
        cmd.draw(3, 1, 0, 0);

        cmd.endRendering();

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

        result = cmd.end();
        ERROR_IF(Failed(result), "Command buffer end failed. Cause: {}", result)
        THEN_ABORT(result);

        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,
        };
        result = commandQueue.submit(1, &submitInfo, currentFrame->m_FrameAvailableFence);
        ERROR_IF(Failed(result), "Command queue submit failed. Cause: {}", result)
        THEN_ABORT(result);

        vk::PresentInfoKHR presentInfo = {
            .waitSemaphoreCount = 1,
            .pWaitSemaphores = &currentFrame->m_RenderFinishSem,
            .swapchainCount = 1,
            .pSwapchains = &swapchain.m_Swapchain,
            .pImageIndices = &imageIndex,
            .pResults = nullptr,
        };
        result = commandQueue.presentKHR(&presentInfo);
        if (Failed(result))
        {
            switch (result)
            {
            case vk::Result::eErrorOutOfDateKHR:
            case vk::Result::eSuboptimalKHR:
                INFO("Recreating Swapchain. Cause: {}", result);
                swapchain.Create(&window);
                viewport.y = Cast<f32>(swapchain.m_Extent.height);
                viewport.width = Cast<f32>(swapchain.m_Extent.width);
                viewport.height = -Cast<f32>(swapchain.m_Extent.height);
                scissor.extent = swapchain.m_Extent;
                break; // Present failed. We redo the frame.
            default:
                ERROR("Command queue present failed. Cause: {}", result)
                THEN_ABORT(result);
            }
        }
        frameIndex = (frameIndex + 1) % MAX_FRAMES_IN_FLIGHT;
    }

    auto result = device.m_Device.waitIdle();
    ERROR_IF(Failed(result), "Wait idle failed. Cause: {}", result);

    device.m_Device.destroy(copyPool, nullptr);
    vbo.Destroy(&device);

    return 0;
}

Frame::Frame(const Device *device, const u32 queueFamilyIndex, const u32 frameCount)
{
    m_Device = device;

    const vk::CommandPoolCreateInfo commandPoolCreateInfo = {
        .flags = vk::CommandPoolCreateFlagBits::eTransient,
        .queueFamilyIndex = queueFamilyIndex,
    };
    vk::Result result = device->m_Device.createCommandPool(&commandPoolCreateInfo, nullptr, &m_Pool);
    ERROR_IF(Failed(result), "Could not command pool for frame {}. Cause: {}", frameCount, result)
    THEN_ABORT(result);

    constexpr vk::FenceCreateInfo fenceCreateInfo = {.flags = vk::FenceCreateFlagBits::eSignaled};
    result = device->m_Device.createFence(&fenceCreateInfo, nullptr, &m_FrameAvailableFence);
    ERROR_IF(Failed(result), "Could not create a fence for frame {}. Cause: {}", frameCount, result)
    THEN_ABORT(result);

    constexpr vk::SemaphoreCreateInfo semaphoreCreateInfo = {};
    result = device->m_Device.createSemaphore(&semaphoreCreateInfo, nullptr, &m_ImageAcquireSem);
    ERROR_IF(Failed(result), "Could not create IA semaphore for frame {}. Cause: {}", frameCount, result)
    THEN_ABORT(result);
    result = device->m_Device.createSemaphore(&semaphoreCreateInfo, nullptr, &m_RenderFinishSem);
    ERROR_IF(Failed(result), "Could not create RF semaphore for frame {}. Cause: {}", frameCount, result)
    THEN_ABORT(result);

    const vk::CommandBufferAllocateInfo allocateInfo = {
        .commandPool = m_Pool, .level = vk::CommandBufferLevel::ePrimary, .commandBufferCount = 1};

    result = m_Device->m_Device.allocateCommandBuffers(&allocateInfo, &m_CommandBuffer);
    ERROR_IF(Failed(result), "Command buffer allocation failed. Cause: {}", result)
    THEN_ABORT(result);

    DEBUG("Frame {} created successfully.", frameCount);
}

Pipeline
CreatePipeline(const Device *device, const Swapchain *swapchain)
{
    // Pipeline Setup
    auto vertexShaderModule = CreateShader(device, VERTEX_SHADER_FILE);
    auto fragmentShaderModule = CreateShader(device, FRAGMENT_SHADER_FILE);

    eastl::array<vk::PipelineShaderStageCreateInfo, 2> shaderStages = {{
        {
            .stage = vk::ShaderStageFlagBits::eVertex,
            .module = vertexShaderModule,
            .pName = "main",
        },
        {
            .stage = vk::ShaderStageFlagBits::eFragment,
            .module = fragmentShaderModule,
            .pName = "main",
        },
    }};

    vk::PipelineLayoutCreateInfo pipelineLayoutCreateInfo = {
        .setLayoutCount = 0,
        .pSetLayouts = nullptr,
        .pushConstantRangeCount = 0,
        .pPushConstantRanges = nullptr,
    };
    vk::PipelineLayout pipelineLayout;
    vk::Result result = device->m_Device.createPipelineLayout(&pipelineLayoutCreateInfo, nullptr, &pipelineLayout);
    ERROR_IF(Failed(result), "Could not create a pipeline layout. Cause: {}", result) THEN_ABORT(result);
    device->SetName(pipelineLayout, "Triangle Layout");

    vk::VertexInputBindingDescription inputBindingDescription = Vertex::GetBinding(0);
    auto inputAttributeDescription = Vertex::GetAttributes(0);

    vk::PipelineVertexInputStateCreateInfo vertexInputStateCreateInfo = {
        .vertexBindingDescriptionCount = 1,
        .pVertexBindingDescriptions = &inputBindingDescription,
        .vertexAttributeDescriptionCount = Cast<u32>(inputAttributeDescription.size()),
        .pVertexAttributeDescriptions = inputAttributeDescription.data(),
    };
    vk::PipelineInputAssemblyStateCreateInfo inputAssemblyStateCreateInfo = {
        .topology = vk::PrimitiveTopology::eTriangleList,
        .primitiveRestartEnable = false,
    };

    vk::PipelineViewportStateCreateInfo viewportStateCreateInfo = {
        .viewportCount = 1,
        .scissorCount = 1,
    };

    vk::PipelineRasterizationStateCreateInfo rasterizationStateCreateInfo = {
        .depthClampEnable = false,
        .rasterizerDiscardEnable = false,
        .polygonMode = vk::PolygonMode::eFill,
        .cullMode = vk::CullModeFlagBits::eNone,
        .frontFace = vk::FrontFace::eCounterClockwise,
        .depthBiasEnable = false,
        .lineWidth = 1.0,
    };
    vk::PipelineMultisampleStateCreateInfo multisampleStateCreateInfo = {
        .rasterizationSamples = vk::SampleCountFlagBits::e1,
        .sampleShadingEnable = false,
    };
    vk::PipelineDepthStencilStateCreateInfo depthStencilStateCreateInfo = {
        .depthTestEnable = false,
        .depthWriteEnable = false,
    };
    vk::PipelineColorBlendAttachmentState colorBlendAttachmentState = {
        .blendEnable = false,
        .srcColorBlendFactor = vk::BlendFactor::eSrcColor,
        .dstColorBlendFactor = vk::BlendFactor::eOneMinusSrcColor,
        .colorBlendOp = vk::BlendOp::eAdd,
        .srcAlphaBlendFactor = vk::BlendFactor::eSrcAlpha,
        .dstAlphaBlendFactor = vk::BlendFactor::eOneMinusSrcAlpha,
        .alphaBlendOp = vk::BlendOp::eAdd,
        .colorWriteMask = vk::ColorComponentFlagBits::eR | vk::ColorComponentFlagBits::eG |
                          vk::ColorComponentFlagBits::eB | vk::ColorComponentFlagBits::eA,
    };
    vk::PipelineColorBlendStateCreateInfo colorBlendStateCreateInfo = {
        .logicOpEnable = false,
        .attachmentCount = 1,
        .pAttachments = &colorBlendAttachmentState,
    };

    eastl::array dynamicStates = {
        vk::DynamicState::eScissor,
        vk::DynamicState::eViewport,
    };

    vk::PipelineDynamicStateCreateInfo dynamicStateCreateInfo = {
        .dynamicStateCount = Cast<u32>(dynamicStates.size()),
        .pDynamicStates = dynamicStates.data(),
    };

    vk::PipelineRenderingCreateInfo renderingCreateInfo = {
        .viewMask = 0,
        .colorAttachmentCount = 1,
        .pColorAttachmentFormats = &swapchain->m_Format,
    };

    vk::GraphicsPipelineCreateInfo pipelineCreateInfo = {
        .pNext = &renderingCreateInfo,
        .stageCount = Cast<u32>(shaderStages.size()),
        .pStages = shaderStages.data(),
        .pVertexInputState = &vertexInputStateCreateInfo,
        .pInputAssemblyState = &inputAssemblyStateCreateInfo,
        .pViewportState = &viewportStateCreateInfo,
        .pRasterizationState = &rasterizationStateCreateInfo,
        .pMultisampleState = &multisampleStateCreateInfo,
        .pDepthStencilState = &depthStencilStateCreateInfo,
        .pColorBlendState = &colorBlendStateCreateInfo,
        .pDynamicState = &dynamicStateCreateInfo,
        .layout = pipelineLayout,
    };
    vk::Pipeline pipeline;
    result = device->m_Device.createGraphicsPipelines(nullptr, 1, &pipelineCreateInfo, nullptr, &pipeline);
    ERROR_IF(Failed(result), "Could not create a graphics pipeline. Cause: {}", result)
    THEN_ABORT(result);
    device->SetName(pipeline, "Triangle Pipeline");

    device->m_Device.destroy(vertexShaderModule, nullptr);
    device->m_Device.destroy(fragmentShaderModule, nullptr);

    return {device, pipelineLayout, pipeline, {}};
}

vk::ShaderModule
CreateShader(const Device *device, cstr shaderFile)
{
    eastl::vector<u32> shaderCode = ReadFile(shaderFile);

    const vk::ShaderModuleCreateInfo shaderModuleCreateInfo = {
        .codeSize = shaderCode.size() * sizeof(u32),
        .pCode = shaderCode.data(),
    };
    vk::ShaderModule shaderModule;

    vk::Result result = device->m_Device.createShaderModule(&shaderModuleCreateInfo, nullptr, &shaderModule);
    ERROR_IF(Failed(result), "Shader {} could not be created. Cause: {}", shaderFile, result)
    THEN_ABORT(result);
    return shaderModule;
}

Frame::~Frame()
{
    m_Device->m_Device.destroy(m_RenderFinishSem, nullptr);
    m_Device->m_Device.destroy(m_ImageAcquireSem, nullptr);
    m_Device->m_Device.destroy(m_FrameAvailableFence, nullptr);
    m_Device->m_Device.destroy(m_Pool, nullptr);

    DEBUG("Destoryed Frame");
}