#include "aster/constants.h"
#include "aster/context.h"
#include "aster/device.h"
#include "aster/glfw_context.h"
#include "aster/physical_device.h"
#include "aster/window.h"

#include "aster/global.h"
#include "aster/swapchain.h"

#include <EASTL/array.h>

constexpr u32 MAX_FRAMES_IN_FLIGHT = 3;

bool IsSuitableDevice(const PhysicalDevice *physicalDevice);

PhysicalDevice FindSuitableDevice(const PhysicalDevices &physicalDevices);

QueueAllocation FindAppropriateQueueAllocation(const PhysicalDevice *physicalDevice);

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

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

    u32 BeginFrame(const Swapchain *swapchain, u32 frameIndex) const;
    [[nodiscard]] vk::CommandBuffer AllocateCommandBuffer() const;
};

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

    GlfwContext glfwContext = {};
    Context context = {"Aster", VERSION};
    Window window = {"Aster1", &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 = vk::True}};
    QueueAllocation queueAllocation = FindAppropriateQueueAllocation(&deviceToUse);
    Device device = {&context, &deviceToUse, &enabledDeviceFeatures, {queueAllocation}, "Primary Device"};

    vk::Queue commandQueue = device.GetQueue(queueAllocation.m_Family, 1);

    Swapchain swapchain = {&window, &device, "Primary Chain"};

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

    u32 frameIndex = 0;
    while (window.Poll())
    {
        Frame *currentFrame = &frames[frameIndex];
        u32 imageIndex = currentFrame->BeginFrame(&swapchain, frameIndex);

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

        vk::ImageMemoryBarrier topOfThePipeBarrier = {
            .oldLayout = vk::ImageLayout::eUndefined,
            .newLayout = vk::ImageLayout::eColorAttachmentOptimal,
            .srcQueueFamilyIndex = queueAllocation.m_Family,
            .dstQueueFamilyIndex = queueAllocation.m_Family,
            .image = currentImage,
            .subresourceRange =
                {
                    .aspectMask = vk::ImageAspectFlagBits::eColor,
                    .baseMipLevel = 0,
                    .levelCount = 1,
                    .baseArrayLayer = 0,
                    .layerCount = 1,
                },
        };
        vk::ImageMemoryBarrier renderToPresentBarrier = {
            .oldLayout = vk::ImageLayout::eColorAttachmentOptimal,
            .newLayout = vk::ImageLayout::ePresentSrcKHR,
            .srcQueueFamilyIndex = queueAllocation.m_Family,
            .dstQueueFamilyIndex = queueAllocation.m_Family,
            .image = currentImage,
            .subresourceRange =
                {
                    .aspectMask = vk::ImageAspectFlagBits::eColor,
                    .baseMipLevel = 0,
                    .levelCount = 1,
                    .baseArrayLayer = 0,
                    .layerCount = 1,
                },
        };

        vk::CommandBufferBeginInfo beginInfo = {.flags = vk::CommandBufferUsageFlagBits::eOneTimeSubmit};
        vk::Result 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::eLoad,
            .storeOp = vk::AttachmentStoreOp::eStore,
            .clearValue = vk::ClearColorValue{0.0f, 0.0f, 0.0f, 1.0f},
        };

        vk::RenderingInfo renderingInfo = {
            .renderArea = {.offset = {0, 0}, .extent = swapchain.m_Extent},
            .layerCount = 1,
            .colorAttachmentCount = 1,
            .pColorAttachments = &attachmentInfo,
        };

        cmd.beginRendering(&renderingInfo);

        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);
                break;
            default:
                ERROR("Command queue present failed. Cause: {}", result)
                THEN_ABORT(result);
            }
        }
        frameIndex = (frameIndex + 1) % MAX_FRAMES_IN_FLIGHT;
    }

    return 0;
}

constexpr QueueSupportFlags REQUIRED_QUEUE_SUPPORT = QueueSupportFlags{} | QueueSupportFlagBits::eGraphics |
                                                     QueueSupportFlagBits::eCompute | QueueSupportFlagBits::ePresent |
                                                     QueueSupportFlagBits::eTransfer;

PhysicalDevice
FindSuitableDevice(const PhysicalDevices &physicalDevices)
{
    for (auto &physicalDevice : physicalDevices)
    {
        if (IsSuitableDevice(&physicalDevice))
        {
            return physicalDevice;
        }
    }

    ERROR("No suitable GPU available on the system.")
    THEN_ABORT(vk::Result::eErrorUnknown);
}

QueueAllocation
FindAppropriateQueueAllocation(const PhysicalDevice *physicalDevice)
{
    for (auto &queueFamilyInfo : physicalDevice->m_QueueFamilies)
    {
        if ((queueFamilyInfo.m_Support & REQUIRED_QUEUE_SUPPORT) == REQUIRED_QUEUE_SUPPORT)
        {
            return {
                .m_Family = queueFamilyInfo.m_Index,
                .m_Count = queueFamilyInfo.m_Count,
            };
        }
    }
    ERROR("No suitable queue family on the GPU.")
    THEN_ABORT(vk::Result::eErrorUnknown);
}

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

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

u32
Frame::BeginFrame(const Swapchain *swapchain, u32 frameIndex) const
{
    const auto device = m_Device->m_Device;

    vk::Result result = device.waitForFences(1, &m_FrameAvailableFence, VK_TRUE, MaxValue<u64>);
    ERROR_IF(Failed(result), "Waiting for fence {} failed. Cause: {}", frameIndex, result)
    THEN_ABORT(result);

    u32 imageIndex;
    result = device.acquireNextImageKHR(swapchain->m_Swapchain, MaxValue<u64>, m_ImageAcquireSem, nullptr, &imageIndex);
    ERROR_IF(Failed(result), "Waiting for swapchain image {} failed. Cause: {}", frameIndex, result)
    THEN_ABORT(result);
    // TODO: Deliberate failure on recreation. Will recreate later.

    result = device.resetFences(1, &m_FrameAvailableFence);
    ERROR_IF(Failed(result), "Fence {} reset failed. Cause: {}", frameIndex, result)
    THEN_ABORT(result);

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

    return imageIndex;
}
vk::CommandBuffer
Frame::AllocateCommandBuffer() const
{
    const vk::CommandBufferAllocateInfo allocateInfo = {
        .commandPool = m_Pool, .level = vk::CommandBufferLevel::ePrimary, .commandBufferCount = 1};

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

    return commandBuffer;
}

Frame::~Frame()
{
    auto result = m_Device->m_Device.waitIdle();
    ERROR_IF(Failed(result), "Wait idle failed. Cause: {}", result);
    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");
}

bool
IsSuitableDevice(const PhysicalDevice *physicalDevice)
{
    const bool hasAllRequiredQueues =
        std::ranges::any_of(physicalDevice->m_QueueFamilies, [](const auto &queueFamilyProp) {
            return (queueFamilyProp.m_Support & REQUIRED_QUEUE_SUPPORT) == REQUIRED_QUEUE_SUPPORT;
        });

    const bool isNotCpu = physicalDevice->m_DeviceProperties.deviceType != vk::PhysicalDeviceType::eCpu;

    const bool hasPresentMode = !physicalDevice->m_PresentModes.empty();

    const bool hasSurfaceFormat = !physicalDevice->m_SurfaceFormats.empty();

    return hasSurfaceFormat && hasPresentMode && isNotCpu && hasAllRequiredQueues;
}