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

#include "aster/aster.h"

#include "aster/core/buffer.h"
#include "aster/core/constants.h"
#include "aster/core/instance.h"
#include "aster/core/physical_device.h"
#include "aster/core/pipeline.h"
#include "aster/core/swapchain.h"
#include "aster/core/window.h"

#include "aster/systems/device.h"
#include "helpers.h"

#include <EASTL/array.h>

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

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

    Window window = {"Triangle (Aster)", {640, 480}};
    systems::Device device{{
        .m_AppName = "Triangle",
        .m_Window = window,
        .m_Name = "Primary",
        .m_Features = {.m_Vulkan12Features = {.bufferDeviceAddress = true},
                       .m_Vulkan13Features = {.synchronization2 = true, .dynamicRendering = true}},
    }};

    Pipeline pipeline = CreatePipeline(&device.m_Device, &device.m_Swapchain);

    // 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}},
    };
    auto vbo = device.CreateVertexBuffer(vertices.size() * sizeof vertices[0], "VBO");
    vbo->Write(0, vertices.size() * sizeof vertices[0], vertices.data());

    Size2D swapchainSize = device.GetSwapchainSize();
    // Persistent variables
    vk::Viewport viewport = {
        .x = 0,
        .y = Cast<f32>(swapchainSize.m_Height),
        .width = Cast<f32>(swapchainSize.m_Width),
        .height = -Cast<f32>(swapchainSize.m_Height),
        .minDepth = 0.0,
        .maxDepth = 1.0,
    };

    vk::Rect2D scissor = {
        .offset = {0, 0},
        .extent = Cast<vk::Extent2D>(swapchainSize),
    };

    vk::ImageSubresourceRange subresourceRange = {
        .aspectMask = vk::ImageAspectFlagBits::eColor,
        .baseMipLevel = 0,
        .levelCount = 1,
        .baseArrayLayer = 0,
        .layerCount = 1,
    };
    vk::ImageMemoryBarrier2 topOfThePipeBarrier = {
        .srcStageMask = vk::PipelineStageFlagBits2::eColorAttachmentOutput,
        .srcAccessMask = vk::AccessFlagBits2::eNone,
        .dstStageMask = vk::PipelineStageFlagBits2::eColorAttachmentOutput,
        .dstAccessMask = vk::AccessFlagBits2::eColorAttachmentWrite,
        .oldLayout = vk::ImageLayout::eUndefined,
        .newLayout = vk::ImageLayout::eColorAttachmentOptimal,
        .srcQueueFamilyIndex = vk::QueueFamilyIgnored,
        .dstQueueFamilyIndex = vk::QueueFamilyIgnored,
        .subresourceRange = subresourceRange,
    };
    vk::DependencyInfo topOfThePipeDependency = {
        .imageMemoryBarrierCount = 1,
        .pImageMemoryBarriers = &topOfThePipeBarrier,
    };
    vk::ImageMemoryBarrier2 renderToPresentBarrier = {
        .srcStageMask = vk::PipelineStageFlagBits2::eColorAttachmentOutput,
        .srcAccessMask = vk::AccessFlagBits2::eColorAttachmentWrite,
        .dstStageMask = vk::PipelineStageFlagBits2::eBottomOfPipe,
        .dstAccessMask = vk::AccessFlagBits2::eNone,
        .oldLayout = vk::ImageLayout::eColorAttachmentOptimal,
        .newLayout = vk::ImageLayout::ePresentSrcKHR,
        .srcQueueFamilyIndex = vk::QueueFamilyIgnored,
        .dstQueueFamilyIndex = vk::QueueFamilyIgnored,
        .subresourceRange = subresourceRange,
    };
    vk::DependencyInfo renderToPresentDependency = {
        .imageMemoryBarrierCount = 1,
        .pImageMemoryBarriers = &renderToPresentBarrier,
    };

    INFO("Starting loop");
    while (window.Poll())
    {
        systems::Frame &currentFrame = device.GetNextFrame();

        auto context = currentFrame.CreateGraphicsContext();

        topOfThePipeBarrier.image = currentFrame.m_SwapchainImage;
        renderToPresentBarrier.image = currentFrame.m_SwapchainImage;

        context.Begin();

        context.Dependency(topOfThePipeDependency);

        // Render
        vk::RenderingAttachmentInfo attachmentInfo = {
            .imageView = currentFrame.m_SwapchainImageView,
            .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, 0.0f},
        };

        vk::RenderingInfo renderingInfo = {
            .renderArea = scissor,
            .layerCount = 1,
            .colorAttachmentCount = 1,
            .pColorAttachments = &attachmentInfo,
        };

        context.BeginRendering(renderingInfo);

        context.SetViewport(viewport);
        context.BindPipeline(pipeline.m_Pipeline);
        context.BindVertexBuffer(vbo);
        context.Draw(3);

        context.EndRendering();

        context.Dependency(renderToPresentDependency);

        context.End();

        device.Submit(currentFrame, context);

        device.Present(currentFrame);
    }

    device.WaitIdle();

    return 0;
}

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