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

#include "aster/aster.h"

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

#define STB_IMAGE_IMPLEMENTATION
#include "aster/systems/commit_manager.h"
#include "aster/systems/device.h"
#include "aster/util/files.h"
#include "stb_image.h"

#include <EASTL/array.h>

constexpr auto VERTEX_SHADER_FILE = "shader/box.vs.hlsl.spv";
constexpr auto FRAGMENT_SHADER_FILE = "shader/box.ps.hlsl.spv";
constexpr auto SHADER_FILE = "triangle";

struct ImageFile
{
    void *m_Data = nullptr;
    u32 m_Width = 0;
    u32 m_Height = 0;
    u32 m_NumChannels = 0;

    bool Load(cstr fileName);
    [[nodiscard]] usize GetSize() const;
    ~ImageFile();
};

bool
ImageFile::Load(cstr fileName)
{
    int width, height, nrChannels;
    m_Data = stbi_load(fileName, &width, &height, &nrChannels, 4);
    ERROR_IF(!m_Data, "Could not load {}", fileName);

    if (!m_Data)
    {
        return false;
    }

    m_Width = width;
    m_Height = height;
    m_NumChannels = 4;

    return true;
}

usize
ImageFile::GetSize() const
{
    return static_cast<usize>(m_Width) * m_Height * m_NumChannels;
}

ImageFile::~ImageFile()
{
    stbi_image_free(m_Data);
    m_Data = nullptr;
}

struct Vertex
{
    vec3 m_Position;
    f32 m_PositionW = 1.0;
    vec2 m_TexCoord0;
    vec2 m_Padding0_ = {0.0f, 0.0f};
};

struct Camera
{
    mat4 m_Model;
    mat4 m_View;
    mat4 m_Perspective;
};

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

    Window window = {"Box (Aster)", {640, 480}};

    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,
                .timelineSemaphore = true,
                .bufferDeviceAddress = true,
                .bufferDeviceAddressCaptureReplay = true,
            },
        .m_Vulkan13Features = {.synchronization2 = true, .dynamicRendering = true},
    };

    systems::Device device{{
        .m_Window = window,
        .m_Features = enabledDeviceFeatures,
        .m_AppName = "Box",
        .m_AppVersion = VERSION,
        .m_ShaderSearchPaths = {"shader/"},
    }};

    // TODO: Device in CommitManager.
    systems::CommitManager commitManager{&device.m_Device, 12, 12, 12, device.CreateSampler({})};

    Pipeline pipeline;
    auto pipelineResult =
        device.CreatePipeline(pipeline, {.m_Shaders = {
                                             {.m_ShaderFile = SHADER_FILE, .m_EntryPoints = {"vsmain", "fsmain"}},
                                         }});
    ERROR_IF(pipelineResult, "Could not create pipeline. Cause: {}", pipelineResult.What())
    THEN_ABORT(pipelineResult.uNone);

    auto swapchainSize = device.GetSwapchainSize();
    Camera camera = {
        .m_Model = {1.0f},
        .m_View = lookAt(vec3(0.0f, 2.0f, 2.0f), vec3(0.0f), vec3(0.0f, 1.0f, 0.0f)),
        .m_Perspective = glm::perspective(
            70_deg, static_cast<f32>(swapchainSize.m_Width) / static_cast<f32>(swapchainSize.m_Height), 0.1f, 100.0f),
    };

    eastl::array vertices = {
        Vertex{.m_Position = vec3(0.5f, 0.5f, -0.5f), .m_TexCoord0 = vec2(1.0f, 1.0f)},
        Vertex{.m_Position = vec3(0.5f, -0.5f, -0.5f), .m_TexCoord0 = vec2(1.0f, 0.0f)},
        Vertex{.m_Position = vec3(-0.5f, -0.5f, -0.5f), .m_TexCoord0 = vec2(0.0f, 0.0f)},
        Vertex{.m_Position = vec3(-0.5f, -0.5f, -0.5f), .m_TexCoord0 = vec2(0.0f, 0.0f)},
        Vertex{.m_Position = vec3(-0.5f, 0.5f, -0.5f), .m_TexCoord0 = vec2(0.0f, 1.0f)},
        Vertex{.m_Position = vec3(0.5f, 0.5f, -0.5f), .m_TexCoord0 = vec2(1.0f, 1.0f)},

        Vertex{.m_Position = vec3(-0.5f, -0.5f, 0.5f), .m_TexCoord0 = vec2(0.0f, 0.0f)},
        Vertex{.m_Position = vec3(0.5f, -0.5f, 0.5f), .m_TexCoord0 = vec2(1.0f, 0.0f)},
        Vertex{.m_Position = vec3(0.5f, 0.5f, 0.5f), .m_TexCoord0 = vec2(1.0f, 1.0f)},
        Vertex{.m_Position = vec3(0.5f, 0.5f, 0.5f), .m_TexCoord0 = vec2(1.0f, 1.0f)},
        Vertex{.m_Position = vec3(-0.5f, 0.5f, 0.5f), .m_TexCoord0 = vec2(0.0f, 1.0f)},
        Vertex{.m_Position = vec3(-0.5f, -0.5f, 0.5f), .m_TexCoord0 = vec2(0.0f, 0.0f)},

        Vertex{.m_Position = vec3(-0.5f, 0.5f, 0.5f), .m_TexCoord0 = vec2(1.0f, 1.0f)},
        Vertex{.m_Position = vec3(-0.5f, 0.5f, -0.5f), .m_TexCoord0 = vec2(1.0f, 0.0f)},
        Vertex{.m_Position = vec3(-0.5f, -0.5f, -0.5f), .m_TexCoord0 = vec2(0.0f, 0.0f)},
        Vertex{.m_Position = vec3(-0.5f, -0.5f, -0.5f), .m_TexCoord0 = vec2(0.0f, 0.0f)},
        Vertex{.m_Position = vec3(-0.5f, -0.5f, 0.5f), .m_TexCoord0 = vec2(0.0f, 1.0f)},
        Vertex{.m_Position = vec3(-0.5f, 0.5f, 0.5f), .m_TexCoord0 = vec2(1.0f, 1.0f)},

        Vertex{.m_Position = vec3(0.5f, -0.5f, -0.5f), .m_TexCoord0 = vec2(0.0f, 0.0f)},
        Vertex{.m_Position = vec3(0.5f, 0.5f, -0.5f), .m_TexCoord0 = vec2(1.0f, 0.0f)},
        Vertex{.m_Position = vec3(0.5f, 0.5f, 0.5f), .m_TexCoord0 = vec2(1.0f, 1.0f)},
        Vertex{.m_Position = vec3(0.5f, 0.5f, 0.5f), .m_TexCoord0 = vec2(1.0f, 1.0f)},
        Vertex{.m_Position = vec3(0.5f, -0.5f, 0.5f), .m_TexCoord0 = vec2(0.0f, 1.0f)},
        Vertex{.m_Position = vec3(0.5f, -0.5f, -0.5f), .m_TexCoord0 = vec2(0.0f, 0.0f)},

        Vertex{.m_Position = vec3(-0.5f, -0.5f, -0.5f), .m_TexCoord0 = vec2(0.0f, 0.0f)},
        Vertex{.m_Position = vec3(0.5f, -0.5f, -0.5f), .m_TexCoord0 = vec2(1.0f, 0.0f)},
        Vertex{.m_Position = vec3(0.5f, -0.5f, 0.5f), .m_TexCoord0 = vec2(1.0f, 1.0f)},
        Vertex{.m_Position = vec3(0.5f, -0.5f, 0.5f), .m_TexCoord0 = vec2(1.0f, 1.0f)},
        Vertex{.m_Position = vec3(-0.5f, -0.5f, 0.5f), .m_TexCoord0 = vec2(0.0f, 1.0f)},
        Vertex{.m_Position = vec3(-0.5f, -0.5f, -0.5f), .m_TexCoord0 = vec2(0.0f, 0.0f)},

        Vertex{.m_Position = vec3(0.5f, 0.5f, 0.5f), .m_TexCoord0 = vec2(1.0f, 1.0f)},
        Vertex{.m_Position = vec3(0.5f, 0.5f, -0.5f), .m_TexCoord0 = vec2(1.0f, 0.0f)},
        Vertex{.m_Position = vec3(-0.5f, 0.5f, -0.5f), .m_TexCoord0 = vec2(0.0f, 0.0f)},
        Vertex{.m_Position = vec3(-0.5f, 0.5f, -0.5f), .m_TexCoord0 = vec2(0.0f, 0.0f)},
        Vertex{.m_Position = vec3(-0.5f, 0.5f, 0.5f), .m_TexCoord0 = vec2(0.0f, 1.0f)},
        Vertex{.m_Position = vec3(0.5f, 0.5f, 0.5f), .m_TexCoord0 = vec2(1.0f, 1.0f)},
    };

    ImageFile imageFile;
    bool loaded = imageFile.Load("image/container.jpg");
    assert(loaded);
    INFO("Image {}x{} : {} channels", imageFile.m_Width, imageFile.m_Height, imageFile.m_NumChannels);

    auto vbo = device.CreateStorageBuffer(vertices.size() * sizeof vertices[0], "Vertex Buffer");
    vbo->Write(0, vertices.size() * sizeof vertices[0], vertices.data());

    auto crate = device.CreateTexture2DWithView({
        .m_Format = vk::Format::eR8G8B8A8Srgb,
        .m_Extent = {imageFile.m_Width, imageFile.m_Height},
        .m_Name = "Crate Texture",
    });

    {

        auto imageStaging = device.CreateStagingBuffer(imageFile.GetSize(), "Image Staging");
        imageStaging->Write(0, imageFile.GetSize(), imageFile.m_Data);

        vk::ImageMemoryBarrier2 imageReadyToWrite = {
            .srcStageMask = vk::PipelineStageFlagBits2::eTransfer,
            .srcAccessMask = vk::AccessFlagBits2::eNone,
            .dstStageMask = vk::PipelineStageFlagBits2::eTransfer,
            .dstAccessMask = vk::AccessFlagBits2::eTransferWrite,
            .oldLayout = vk::ImageLayout::eUndefined,
            .newLayout = vk::ImageLayout::eTransferDstOptimal,
            .srcQueueFamilyIndex = vk::QueueFamilyIgnored,
            .dstQueueFamilyIndex = vk::QueueFamilyIgnored,
            .image = crate->GetImage(),
            .subresourceRange =
                {
                    .aspectMask = vk::ImageAspectFlagBits::eColor,
                    .baseMipLevel = 0,
                    .levelCount = 1,
                    .baseArrayLayer = 0,
                    .layerCount = 1,
                },
        };
        vk::DependencyInfo imageReadyToWriteDependency = {
            .imageMemoryBarrierCount = 1,
            .pImageMemoryBarriers = &imageReadyToWrite,
        };

        vk::ImageMemoryBarrier2 imageReadyToRead = {
            .srcStageMask = vk::PipelineStageFlagBits2::eTransfer,
            .srcAccessMask = vk::AccessFlagBits2::eTransferWrite,
            .dstStageMask = vk::PipelineStageFlagBits2::eFragmentShader,
            .dstAccessMask = vk::AccessFlagBits2::eShaderRead,
            .oldLayout = vk::ImageLayout::eTransferDstOptimal,
            .newLayout = vk::ImageLayout::eShaderReadOnlyOptimal,
            .srcQueueFamilyIndex = vk::QueueFamilyIgnored,
            .dstQueueFamilyIndex = vk::QueueFamilyIgnored,
            .image = crate->GetImage(),
            .subresourceRange =
                {
                    .aspectMask = vk::ImageAspectFlagBits::eColor,
                    .baseMipLevel = 0,
                    .levelCount = 1,
                    .baseArrayLayer = 0,
                    .layerCount = 1,
                },
        };
        vk::DependencyInfo imageReadyToReadDependency = {
            .imageMemoryBarrierCount = 1,
            .pImageMemoryBarriers = &imageReadyToRead,
        };

        auto &context = device.CreateTransferContext();
        context.Begin();

        context.Dependency(imageReadyToWriteDependency);

        context.UploadTexture(crate->m_Image, {.m_Data = imageFile.m_Data, .m_NumBytes = imageFile.GetSize()});

        context.Dependency(imageReadyToReadDependency);

        context.End();

        auto recpt = device.Submit(context);
        device.WaitOn(recpt);
    }

    auto ubo = device.CreateStorageBuffer(sizeof camera, "Camera UBO");
    ubo->Write(0, sizeof camera, &camera);

    // Persistent variables

    vk::ImageSubresourceRange subresourceRange = {
        .aspectMask = vk::ImageAspectFlagBits::eColor,
        .baseMipLevel = 0,
        .levelCount = 1,
        .baseArrayLayer = 0,
        .layerCount = 1,
    };
    vk::ImageMemoryBarrier2 topOfThePipeBarrier = {
        // For Color Attachment output ref:
        // https://github.com/KhronosGroup/Vulkan-ValidationLayers/issues/7193#issuecomment-1875960974
        .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,
    };

    eastl::fixed_vector<Ref<ImageView>, MAX_FRAMES_IN_FLIGHT> depthImages;

    auto initDepthImages = [&depthImages, &device](const vk::Extent2D extent) {
        for (u32 i = 0; i < MAX_FRAMES_IN_FLIGHT; ++i)
        {
            depthImages.push_back(device.CreateDepthStencilImageWithView({.m_Extent = extent, .m_Name = "Depth"}));
        }
    };

    initDepthImages(swapchainSize);

    auto recreateDepthBuffers = [&depthImages, &initDepthImages](const vk::Extent2D extent) {
        depthImages.clear();
        initDepthImages(extent);
    };

    struct PCB
    {
        systems::ResId<Buffer> m_Camera;
        systems::ResId<Buffer> m_VertexBuffer;
        systems::ResId<TextureView> m_Texture;
    };

    PCB pcb = {
        .m_Camera = commitManager.CommitBuffer(ubo),
        .m_VertexBuffer = commitManager.CommitBuffer(vbo),
        .m_Texture = commitManager.CommitTexture(crate),
    };

    Time::Init();

    auto prevSwapchainSize = swapchainSize;

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

        camera.m_Model *= rotate(mat4{1.0f}, static_cast<f32>(45.0_deg * Time::m_Delta), vec3(0.0f, 1.0f, 0.0f));
        ubo->Write(0, sizeof camera, &camera);

        auto currentFrame = device.GetNextFrame();

        prevSwapchainSize = swapchainSize;
        swapchainSize = currentFrame.m_SwapchainSize;
        if (swapchainSize != prevSwapchainSize)
        {
            recreateDepthBuffers(swapchainSize);
        }

        vk::Viewport viewport = {
            .x = 0,
            .y = static_cast<f32>(swapchainSize.m_Height),
            .width = static_cast<f32>(swapchainSize.m_Width),
            .height = -static_cast<f32>(swapchainSize.m_Height),
            .minDepth = 0.0,
            .maxDepth = 1.0,
        };

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

        vk::ImageView currentImageView = currentFrame.m_SwapchainImageView;
        vk::Image currentImage = currentFrame.m_SwapchainImage;
        vk::ImageView currentDepthImageView = depthImages[currentFrame.m_FrameIdx]->m_View;

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

        auto context = currentFrame.CreateGraphicsContext();

        context.Begin();

        context.Dependency(topOfThePipeDependency);

        // 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, 0.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 = scissor,
            .layerCount = 1,
            .colorAttachmentCount = static_cast<u32>(attachmentInfos.size()),
            .pColorAttachments = attachmentInfos.data(),
            .pDepthAttachment = &depthAttachment,
        };

        context.BeginRendering(renderingInfo);

        context.SetViewport(viewport);
        context.BindPipeline(pipeline);
        /*cmd.bindDescriptorSets(vk::PipelineBindPoint::eGraphics, pipeline.m_Layout, 0, 1,
                               &commitManager.GetDescriptorSet(), 0, nullptr);*/
        //context.PushConstantBlock(pcb);
        context.Draw(3);

        context.EndRendering();

        context.Dependency(renderToPresentDependency);

        context.End();

        device.Present(currentFrame, context);
    }

    device.WaitIdle();

    return 0;
}