#include <AssetManager/AssetManager.hpp>
#include <UniqueID/UniqueID.hpp>
#include <Utilities/RGBABytes.hpp>
#include <cstdlib>
#include <filesystem>
#include <iostream>
#include <chrono>
#include <GraphicsContext/GraphicsContext.hpp>
#include <Display/Display.hpp>
#include <Utilities/Utilities.hpp>
#include <Log/Log.hpp>
#include <ToolTip/ToolTip.hpp>
#include <Version/Version.hpp>
#include <Vk/ShaderModule/ShaderModule.hpp>
#include <Vk/ImageView/ImageView.hpp>
#include <Vk/Sampler/Sampler.hpp>
#include <Vk/ShaderConst.hpp>
#include <IndexedUnitCube/IndexedUnitCube.hpp>
#include <PerfHistory/PerfHistory.hpp>
#include <Settings/Settings.hpp>
#include <PixelMap/PixelMap.hpp>
#include <Utilities/Perf.hpp>
#include <memory>
#include <sstream>
#include <VoxelTree/VoxelTree.hpp>

#ifdef DEBUG_BUILD
#   define WINDOW_CAPTION_FORMAT "Voxelverse $+ [DEBUG BUILD]"
#else
#   define WINDOW_CAPTION_FORMAT "Voxelverse $+"
#endif

void GraphicsContext::VirtualViewport::apply(Vk::CommandBuffer &cmdbuf)
{
    if (size.x <= 0 || size.y <= 0 || (pos.x + size.x) <= 0 || (pos.y + size.y) <= 0)
    {
        cmdbuf.setViewport(0.0f, 0.0f, 1.0f, 1.0f);
        cmdbuf.setScissor({{3, 3}, {1, 1}});
        return;
    }

    cmdbuf.setViewport(
        (float) (pos.x * pixelSize), (float) (pos.y * pixelSize),
        (float) (size.x * pixelSize), (float) (size.y * pixelSize)
    );

    cmdbuf.setScissor({
        {(int32_t) (clipPos.x * pixelSize), (int32_t) (clipPos.y * pixelSize)},
        {(uint32_t) (clipSize.x * pixelSize), (uint32_t) (clipSize.y * pixelSize)}
    });
}

GraphicsContext* GraphicsContext::instance;
CommandQueue<GraphicsContext> GraphicsContext::m_asyncCommandQueue;
std::atomic<int> GraphicsContext::m_numPendingSubmissions;
std::atomic<uint64_t> GraphicsContext::m_frameCounter;

glm::vec4 GraphicsContext::color;
std::shared_ptr<Vk::DescriptorSetLayout> GraphicsContext::materialLayout;
std::shared_ptr<Vk::DescriptorSetLayout> GraphicsContext::voxelPaletteDescriptorSetLayout;
std::shared_ptr<Vk::DescriptorSetLayout> GraphicsContext::arrayDescriptorSetLayout;
std::shared_ptr<Vk::Pipeline> GraphicsContext::svoUpdatePipeline;

void GraphicsContext::checkFrameTime()
{
#ifndef NO_PERF_TIMERS
    if (!settings.frameDumpThresholdMS.has_value())
    {
        return;
    }

    PerfTimeline &timeline = PerfHistory::frame.getTimeline(0);
    PerfTimeline::Ticks frameDuration = timeline.getCurrentDuration();
    int frameMS = frameDuration / 1ms;

    if (frameMS > *settings.frameDumpThresholdMS)
    {
        Log::debug(formatString("Frame #$+ over threshold, took $+ms", m_frameCounter.load(), frameMS));

        PixelMap chart = makePerfTimelineChart(timeline, *settings.frameDumpThresholdMS * 1ms);
        
        std::string path = getDumpPath(formatString("$+__$+ms.png", m_frameCounter.load(), frameMS));
        chart.writePNG(path);
    }
#endif
}

glm::mat4 GraphicsContext::makeRectMatrix(int x, int y, int width, int height)
{
    glm::ivec2 size = m_viewportStack.top().size;

    float clipX = 2.0f * (float) x / (float) size.x - 1.0f;
    float clipY = 2.0f * (float) y / (float) size.y - 1.0f;
    float clipWidth = 2.0f * (float) width / (float) size.x;
    float clipHeight = 2.0f * (float) height / (float) size.y;

    return glm::mat4(
        clipWidth, 0.0f, 0.0f, 0.0f,
        0.0f, clipHeight, 0.0f, 0.0f,
        0.0f, 0.0f, 1.0f, 0.0f,
        clipX, clipY, 0.0f, 1.0
    );
}

void GraphicsContext::recreateSwapchain()
{
    // Wait for the device to be idle.
    m_device->waitIdle();

    // Create the swapchain.
    int fbWidth, fbHeight;
    SDL_GetWindowSizeInPixels(m_sdlwin, &fbWidth, &fbHeight);

    if (fbWidth == 0 || fbHeight == 0)
    {
        m_swapchainDirty = true;
        return;
    }

    // Clean up existing objects.
    m_presentFramebuffers.clear();
    m_swapchain = nullptr;
    m_guiRenderPass = nullptr;
    m_guiPipeline = nullptr;

    m_swapchain = std::make_shared<Vk::Swapchain>(m_device, m_surface, fbWidth, fbHeight);

    m_guiRenderPass = Vk::RenderPass::Builder()
        .addColorAttachment(
            m_swapchain->getImageFormat(),
            VK_ATTACHMENT_LOAD_OP_LOAD,
            VK_ATTACHMENT_STORE_OP_STORE,
            VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
            VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL
        )
        .build(m_device);
    
    m_clearRenderPass = Vk::RenderPass::Builder()
        .addColorAttachment(
            m_swapchain->getImageFormat(),
            VK_ATTACHMENT_LOAD_OP_CLEAR,
            VK_ATTACHMENT_STORE_OP_STORE,
            VK_IMAGE_LAYOUT_UNDEFINED,
            VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL
        )
        .build(m_device);

    m_guiPipeline = Vk::Pipeline::GraphicsBuilder{}
        .addShaderStage(VK_SHADER_STAGE_VERTEX_BIT, *Vk::ShaderModule::vert_gui)
        .addShaderStage(VK_SHADER_STAGE_FRAGMENT_BIT, *Vk::ShaderModule::frag_gui)
        .addDynamicState(VK_DYNAMIC_STATE_VIEWPORT)
        .addDynamicState(VK_DYNAMIC_STATE_SCISSOR)
        .addPushConstantRange(0, sizeof(GuiRectInfo), VK_SHADER_STAGE_VERTEX_BIT)
        .addDescriptorSetLayout(materialLayout)
        .addAlphaBlendedAttachment()
        .build(m_device, m_guiRenderPass);

    // Create image views and framebuffers for the swapchain.
    for (VkImage image : m_swapchain->images())
    {
        std::vector<std::shared_ptr<Vk::ImageView>> imageViews;
        imageViews.push_back(
            std::make_shared<Vk::ImageView>(
                m_device,
                image,
                m_swapchain->getImageFormat()
            )
        );

        m_presentFramebuffers.push_back(
            std::make_shared<Vk::Framebuffer>(
                m_device,
                m_guiRenderPass,
                imageViews,
                fbWidth, fbHeight
            )
        );
    }
}

GraphicsContext::GraphicsContext() :
    m_nextFrameInFlight(0),
    m_windowWidth(1600), m_windowHeight(800),
    m_fullscreen(false),
    m_swapchainDirty(false),
    m_tooltip(nullptr)
{
    int iconComp, iconWidth, iconHeight;
    std::string logoPath = assets.resolve("logo.png");
    unsigned char *iconData = stbi_load(logoPath.c_str(), &iconWidth, &iconHeight, &iconComp, STBI_rgb_alpha);

    if (iconData == nullptr)
    {
        crash(formatString("Failed to load the window icon `$+`!", logoPath));
    }

    // Set the static instance pointer.
    instance = this;

    // Create the game window.
    std::string windowTitle = formatString(WINDOW_CAPTION_FORMAT, Version::ofThisBuild);
    if (envVars.devMode)
    {
        windowTitle += " [DEV MODE]";
    }

    m_sdlwin = SDL_CreateWindow(
        windowTitle.c_str(),
        m_windowWidth, m_windowHeight,
        SDL_WINDOW_VULKAN | SDL_WINDOW_RESIZABLE
    );

    if (m_sdlwin == nullptr)
    {
        crash(formatString("Failed to create window: $+", SDL_GetError()));
    }

    // Set up the icon. On Windows we don't do anything because the EXE icon is used by default,
    // and it correctly sets up different sizes.
#ifndef _WIN32
    SDL_Surface *iconSurface = SDL_CreateSurfaceFrom(
        64, 64,
        SDL_PIXELFORMAT_RGBA32,
        iconData,
        64 * 4
    );
    SDL_SetWindowIcon(m_sdlwin, iconSurface);
    SDL_DestroySurface(iconSurface);
#endif

    // Create the global Vulkan objects.
    m_instance = std::make_shared<Vk::Instance>(m_sdlwin);
    m_surface = std::make_shared<Vk::Surface>(m_sdlwin, m_instance);
    m_device = std::make_shared<Vk::Device>(m_instance, m_surface);
    Vk::ShaderModule::createAll(m_device);

    // Create the GUI-related Vulkan objects.
    materialLayout = Vk::DescriptorSetLayout::Builder{}
        .addBinding(
            GLSL_COLOR_IMAGE_INDEX,
            VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
            1,
            VK_SHADER_STAGE_FRAGMENT_BIT
        )
        .addBinding(
            GLSL_NORMALMAP_IMAGE_INDEX,
            VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
            1,
            VK_SHADER_STAGE_FRAGMENT_BIT
        )
        .addBinding(
            GLSL_LIGHTMAP_IMAGE_INDEX,
            VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
            1,
            VK_SHADER_STAGE_FRAGMENT_BIT
        )
        .addBinding(
            GLSL_SPECULARMAP_IMAGE_INDEX,
            VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
            1,
            VK_SHADER_STAGE_FRAGMENT_BIT
        )
        .addBinding(
            GLSL_BIASMAP_IMAGE_INDEX,
            VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
            1,
            VK_SHADER_STAGE_FRAGMENT_BIT
        )
        .addBinding(
            GLSL_IMAGE_METADATA_INDEX,
            VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
            1,
            VK_SHADER_STAGE_FRAGMENT_BIT
        )
        .build(m_device);

    voxelPaletteDescriptorSetLayout = Vk::DescriptorSetLayout::Builder{}
        .addBinding(
            0,
            VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            1,
            VK_SHADER_STAGE_FRAGMENT_BIT | VK_SHADER_STAGE_COMPUTE_BIT
        )
        .addBinding(
            1,
            VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            1,
            VK_SHADER_STAGE_FRAGMENT_BIT | VK_SHADER_STAGE_COMPUTE_BIT
        )
        .build(GraphicsContext::instance->getDevice());

    arrayDescriptorSetLayout = Vk::DescriptorSetLayout::Builder{}
        .addBinding(
            0,
            VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            1,
            VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT | VK_SHADER_STAGE_COMPUTE_BIT
        )
        .build(GraphicsContext::instance->getDevice());
    
    svoUpdatePipeline = Vk::Pipeline::ComputeBuilder{}
        .addDescriptorSetLayout(arrayDescriptorSetLayout)
        .addDescriptorSetLayout(arrayDescriptorSetLayout)
        .addPushConstantRange(0, sizeof(VoxelTree::UpdatePushConst), VK_SHADER_STAGE_COMPUTE_BIT)
        .build(getDevice(), *Vk::ShaderModule::comp_svoUpdate);

    // Create the swapchain and related objects.
    recreateSwapchain();

    // Create the graphics command pool.
    m_graphicsCommandPool = std::make_shared<Vk::CommandPool>(
        m_device,
        VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT,
        m_device->getQueueFamilies().graphicsFamily.value()
    );

    // Create the per-frame objects.
    int i;
    for (i=0; i<maxFramesInFlight; i++)
    {
        FrameInFlight &frame = m_framesInFlight[i];

        frame.precmdbuf = std::make_shared<Vk::CommandBuffer>(m_device, m_graphicsCommandPool, VK_COMMAND_BUFFER_LEVEL_PRIMARY);
        frame.cmdbuf = std::make_shared<Vk::CommandBuffer>(m_device, m_graphicsCommandPool, VK_COMMAND_BUFFER_LEVEL_PRIMARY);
        frame.imageAvailableSemaphore = std::make_shared<Vk::Semaphore>(m_device);
        frame.renderFinishedSemaphore = std::make_shared<Vk::Semaphore>(m_device);
        frame.fence = std::make_shared<Vk::Fence>(m_device, VK_FENCE_CREATE_SIGNALED_BIT);
        
        if (m_device->areTimerQueriesEnabled())
        {
            frame.queryPool = std::make_unique<Vk::QueryPool>(m_device, 2);
        }
    }

    // Create the samplers.
    Vk::Sampler::createAll(m_device);

    // Create the one-time command buffer.
    std::shared_ptr<Vk::CommandBuffer> initCommandBuffer = std::make_shared<Vk::CommandBuffer>(
        m_device,
        m_graphicsCommandPool,
        VK_COMMAND_BUFFER_LEVEL_PRIMARY
    );

    initCommandBuffer->reset();
    initCommandBuffer->begin(VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT);

    // Load the game icon into a material.
    Material::internalGameIcon.widthValue = 64;
    Material::internalGameIcon.heightValue = 64;
    Material::internalGameIcon.addComponent(m_device, *initCommandBuffer, iconData);
    Material::internalGameIcon.buildMetadata(m_device, *initCommandBuffer);
    Material::internalGameIcon.buildDescriptorSet(m_device);

    // Load the "white" material.
    static uint8_t whitePixel[4] = {0xFF, 0xFF, 0xFF, 0xFF};
    Material::white.widthValue = 1;
    Material::white.heightValue = 1;
    Material::white.addComponent(m_device, *initCommandBuffer, whitePixel);
    Material::white.buildMetadata(m_device, *initCommandBuffer);
    Material::white.buildDescriptorSet(m_device);

    // Initialize indexed unit cube.
    IndexedUnitCube::init(m_device, *initCommandBuffer);

    // Submit the initialization commands, and wait for them to complete.
    initCommandBuffer->end();
    initCommandBuffer->submit();
    m_device->waitIdle();

    // Clean up the initialization stuff.
    Material::internalGameIcon.finalize();
    Material::white.finalize();
    IndexedUnitCube::finalize();
}

GraphicsContext::~GraphicsContext()
{
    m_device->waitIdle();
}

void GraphicsContext::toggleFullscreen()
{
    m_fullscreen = !m_fullscreen;
    if (m_fullscreen)
    {
        SDL_SetWindowFullscreen(m_sdlwin, true);
        SDL_GetWindowSizeInPixels(m_sdlwin, &m_windowWidth, &m_windowHeight);
    }
    else
    {
        SDL_SetWindowFullscreen(m_sdlwin, false);
        SDL_GetWindowSizeInPixels(m_sdlwin, &m_windowWidth, &m_windowHeight);
    }

    m_swapchainDirty = true;
}

bool GraphicsContext::drawFrame()
{
    // Re-create the swapchain if dirty.
    if (m_swapchainDirty)
    {
        m_swapchainDirty = false;
        recreateSwapchain();
    }

    if (m_swapchainDirty)
    {
        return false;
    }

    std::chrono::time_point<std::chrono::steady_clock> frameStartTime = std::chrono::steady_clock::now();

    // Process an async graphics command.
    int maxToSubmit = 1;
    if (m_numPendingSubmissions > 20)
    {
        maxToSubmit = 2;
    }
    
    if (m_numPendingSubmissions > 100)
    {
        maxToSubmit = 4;
    }

    if (m_numPendingSubmissions > 200)
    {
        maxToSubmit = 8;
    }

    int i;
    for (i=0; i<maxToSubmit; i++)
    {
        if (m_asyncCommandQueue.processNextCommand(*this))
        {
            m_numPendingSubmissions--;
        }
        else
        {
            break;
        }
    }

    // Try to get a frame in flight.
    m_thisFrame = m_nextFrameInFlight;
    FrameInFlight &frame = m_framesInFlight[m_thisFrame];
    if (!frame.fence->wait(0))
    {
        Uint32 deltaTime = SDL_GetTicks() - frame.drawStarted;
        if (deltaTime > 5000)
        {
            crash("GPU freeze detected: A single frame has been rendering for more than 10s!");
        }

        return false;
    }

    if (frame.timerQueryReady)
    {
        drawTracker.submitSample(frame.queryPool->getDeltaTimeMS(0, 1));
    }

    // If a screenshot buffer is set, it is now filled with the RGBA data, so we can read it.
    if (frame.screenshotBuffer != nullptr)
    {
        // TODO: Can we support other image formats for screenshots? Currently we assume B8G8R8A8_SRGB.
        std::error_code ec;
        std::filesystem::create_directories("screenshots", ec);

        UniqueID screenshotID = UniqueID::generate();
        std::string path = formatString("screenshots/screenshot-$+.png", screenshotID);

        std::vector<unsigned char> convertedPixels(4 * frame.screenshotExtent.width * frame.screenshotExtent.height);
        std::memcpy(
            convertedPixels.data(),
            frame.screenshotBuffer->hostdata(),
            convertedPixels.size()
        );

        for (uint32_t y = 0; y < frame.screenshotExtent.height; y++)
        {
            for (uint32_t x = 0; x < frame.screenshotExtent.width; x++)
            {
                unsigned char *pixel = &convertedPixels[4 * (y * frame.screenshotExtent.width + x)];
                unsigned char temp = pixel[0];
                pixel[0] = pixel[2];
                pixel[2] = temp;
                pixel[3] = 0xFF;
            }
        }

        stbi_write_png(
            path.c_str(),
            frame.screenshotExtent.width,
            frame.screenshotExtent.height,
            4,
            convertedPixels.data(),
            4 * frame.screenshotExtent.width
        );

        Log::info(formatString("Screenshot saved to '$+'.", path));
        frame.screenshotBuffer = nullptr;
    }

    // Clear the held resources.
    frame.heldResources.clear();

    // Try to acquire the next swapchain image.
    std::optional<uint32_t> optionalNextImageIndex = m_swapchain->acquireNextImage(
        0,
        frame.imageAvailableSemaphore.get()
    );
    if (!optionalNextImageIndex.has_value())
    {
        recreateSwapchain();
        return false;
    }

    PerfHistory::frame.beginTimeline();

    // Image acquired, advance to the next frame and reset the fence.
    uint32_t imageIndex = optionalNextImageIndex.value();
    m_nextFrameInFlight = (m_nextFrameInFlight + 1) % maxFramesInFlight;
    m_frameCounter++;
    frame.fence->reset();

    // Set up the pre-command buffer.
    m_precmdbuf = frame.precmdbuf;
    m_precmdbuf->reset();
    m_precmdbuf->begin();

    // Set up the command buffer.
    m_cmdbuf = frame.cmdbuf;
    m_framebuffer = m_presentFramebuffers[imageIndex];
    m_cmdbuf->reset();
    m_cmdbuf->waitFor(*frame.imageAvailableSemaphore, VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT);
    m_cmdbuf->signal(*frame.renderFinishedSemaphore);
    m_cmdbuf->begin();

    if (frame.queryPool != nullptr)
    {
        m_cmdbuf->resetQueryPool(*frame.queryPool, 0, 2);
        m_cmdbuf->writeTimestamp(*frame.queryPool, 0);
    }

    // Draw the display.
    std::shared_ptr<Display> display = Display::getCurrent();
    display->draw(*this);

    // Prepare screenshot if necessary.
    if (screenshotRequested)
    {
        screenshotRequested = false;

        frame.screenshotExtent = m_swapchain->getExtent();
        frame.screenshotBuffer = std::make_shared<Vk::Buffer>(
            m_device,
            frame.screenshotExtent.width * frame.screenshotExtent.height * 4,
            VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT
        );

        m_cmdbuf->copyImageToBuffer(
            m_swapchain->images()[imageIndex],
            *frame.screenshotBuffer,
            frame.screenshotExtent.width,
            frame.screenshotExtent.height
        );
    }

    // Ensure the image is presentable.
    m_cmdbuf->transitionImageLayout(
        m_swapchain->images()[imageIndex],
        VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
        VK_IMAGE_LAYOUT_PRESENT_SRC_KHR
    );

    // Write to the end-of-frame timestamp.
    if (frame.queryPool != nullptr)
    {
        m_cmdbuf->writeTimestamp(*frame.queryPool, 1);
        frame.timerQueryReady = true;
    }

    // Submit the pre-command buffer.
    m_precmdbuf->end();
    m_precmdbuf->submit();

    // Submit the command buffer.
    m_cmdbuf->end();
    m_cmdbuf->submit(frame.fence.get());
    frame.drawStarted = SDL_GetTicks();

    // Present the frame.
    VkSemaphore presentSemaphores[] = {frame.renderFinishedSemaphore->inner()};
    VkPresentInfoKHR presentInfo{};
    presentInfo.sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR;

    presentInfo.waitSemaphoreCount = 1;
    presentInfo.pWaitSemaphores = presentSemaphores;

    VkSwapchainKHR swapChains[] = {m_swapchain->inner()};
    presentInfo.swapchainCount = 1;
    presentInfo.pSwapchains = swapChains;
    presentInfo.pImageIndices = &imageIndex;

    vkQueuePresentKHR(m_device->presentQueue.inner(), &presentInfo);

    std::chrono::time_point<std::chrono::steady_clock> frameEndTime = std::chrono::steady_clock::now();
    std::chrono::duration<double, std::milli> frameDuration = frameEndTime - frameStartTime;

    notifyFrameTime(frameDuration.count());

    checkFrameTime();

    return true;
}

void GraphicsContext::clearGui(glm::vec4 color)
{
    VkClearValue clearValue;
    clearValue.color = {{color.x, color.y, color.z, color.w}};

    std::vector<VkClearValue> values;
    values.push_back(clearValue);

    m_cmdbuf->beginRenderPass(*m_clearRenderPass, *m_framebuffer, m_swapchain->getExtent(), values, VK_SUBPASS_CONTENTS_INLINE);
    m_cmdbuf->endRenderPass();
}

void GraphicsContext::beginGui()
{
    // Delete the current tooltip if any.
    delete m_tooltip;
    m_tooltip = nullptr;

    // Set up the initial virtual viewport.
    VirtualViewport initialViewport;
    initialViewport.pos = initialViewport.clipPos = glm::ivec2(0, 0);
    initialViewport.size = initialViewport.clipSize = glm::ivec2(getVirtualWidth(), getVirtualHeight());

    m_viewportStack = std::stack<VirtualViewport>();
    m_viewportStack.emplace(initialViewport);

    // Set up the pipeline.
    std::vector<VkClearValue> clearValues{{{{0.0f, 0.0f, 0.0f, 0.0f}}}};
    m_cmdbuf->beginRenderPass(*m_guiRenderPass, *m_framebuffer, m_swapchain->getExtent(), clearValues, VK_SUBPASS_CONTENTS_INLINE);
    m_cmdbuf->bindPipeline(VK_PIPELINE_BIND_POINT_GRAPHICS, *m_guiPipeline);
    initialViewport.apply(*m_cmdbuf);
}

void GraphicsContext::drawGuiRect(const GuiRectInfo &info)
{
    m_cmdbuf->pushConstants(*m_guiPipeline, VK_SHADER_STAGE_VERTEX_BIT, 0, sizeof(GuiRectInfo), &info);
    m_cmdbuf->draw(6, 1, 0, 0);
}

void GraphicsContext::endGui()
{
    m_cmdbuf->endRenderPass();
}

void GraphicsContext::drawMaterial(Material &mat, int x, int y)
{
    drawMaterial(mat, x, y, mat.width(), mat.height());
}

void GraphicsContext::drawMaterial(Material &mat, int x, int y, int width, int height)
{
    drawMaterial(mat, 0, 0, width, height, x, y, width, height);
}

void GraphicsContext::drawMaterial(Material &mat, int sx, int sy, int sw, int sh, int dx, int dy, int dw, int dh)
{
    glm::ivec2 viewportSize = m_viewportStack.top().size;
    if (viewportSize.x < 1 || viewportSize.y < 1)
    {
        // Nothing to draw.
        return;
    }
    
    m_cmdbuf->bindDescriptorSet(VK_PIPELINE_BIND_POINT_GRAPHICS, *m_guiPipeline, GLSL_GUI_DESCRIPTOR_SET, mat.descriptorSet());
    
    GuiRectInfo info;
    info.rectMatrix = makeRectMatrix(dx, dy, dw, dh);
    info.texMatrix = glm::mat4(
        (float) sw / (float) mat.width(), 0.0f, 0.0f, 0.0f,
        0.0f, (float) sh / (float) mat.height(), 0.0f, 0.0f,
        0.0f, 0.0f, 1.0f, 0.0f,
        (float) sx / (float) mat.width(), (float) sy / (float) mat.height(), 0.0f, 1.0
    );
    info.color = color;
    drawGuiRect(info);
}

void GraphicsContext::drawRect(int x, int y, int width, int height)
{
    drawMaterial(Material::white, x, y, width, height);
}

int GraphicsContext::getWindowWidth()
{
    return m_windowWidth;
}

int GraphicsContext::getWindowHeight()
{
    return m_windowHeight;
}

int GraphicsContext::getVirtualWidth()
{
    return m_windowWidth / pixelSize;
}

int GraphicsContext::getVirtualHeight()
{
    return m_windowHeight / pixelSize;
}

void GraphicsContext::drawText(Font &font, int x, int y, std::string text, bool masked)
{
    glm::ivec2 base(x, y);

    layoutText(
        font,
        text,
        [this](glm::vec4 color)
        {
            this->color = glm::vec4(glm::vec3(color), this->color.a);
        },
        [this, &base](Font::Sheet &sheet, glm::ivec2 pos, glm::ivec2 sheetBase, glm::ivec2 size, size_t)
        {
            drawMaterial(
                sheet.mat,
                sheetBase.x, sheetBase.y,
                size.x, size.y,
                base.x + pos.x, base.y + pos.y,
                size.x, size.y
            );
        },
        masked
    );
}

glm::ivec2 GraphicsContext::measureText(Font &font, std::string text, bool masked)
{
    glm::ivec2 resultingSize(0, font.glyphHeight);

    layoutText(
        font,
        text,
        [](glm::vec4){},
        [&resultingSize](Font::Sheet &sheet, glm::ivec2 pos, glm::ivec2 sheetBase, glm::ivec2 size, size_t)
        {
            resultingSize.x = std::max(resultingSize.x, pos.x + size.x);
            resultingSize.y = std::max(resultingSize.y, pos.y + size.y);
        },
        masked
    );

    return resultingSize;
}

std::string GraphicsContext::colorstr(glm::vec4 color)
{
    int red = static_cast<int>(color.r * 255.0f);
    int green = static_cast<int>(color.g * 255.0f);
    int blue = static_cast<int>(color.b * 255.0f);

    std::ostringstream oss;
    oss << "\a"
        << std::hex << std::setw(2) << std::setfill('0') << red 
        << std::hex << std::setw(2) << std::setfill('0') << green 
        << std::hex << std::setw(2) << std::setfill('0') << blue;

    return oss.str();
}

std::string GraphicsContext::colorstr(const RGBABytes &color)
{
    std::ostringstream oss;
    oss << "\a"
        << std::hex << std::setw(2) << std::setfill('0') << (int) color.red 
        << std::hex << std::setw(2) << std::setfill('0') << (int) color.green 
        << std::hex << std::setw(2) << std::setfill('0') << (int) color.blue;

    return oss.str();
}

void GraphicsContext::pushViewport()
{
    VirtualViewport viewport = m_viewportStack.top();
    m_viewportStack.push(viewport);
}

void GraphicsContext::popViewport()
{
    m_viewportStack.pop();
    m_viewportStack.top().apply(*m_cmdbuf);
}

void GraphicsContext::clipViewport(glm::ivec2 pos, glm::ivec2 size)
{
    VirtualViewport &viewport = m_viewportStack.top();
    viewport.pos += pos;
    viewport.size.x = std::min(size.x, viewport.size.x - pos.x);
    viewport.size.y = std::min(size.y, viewport.size.y - pos.y);
    
    int clipX2 = viewport.clipPos.x + viewport.clipSize.x;
    int clipY2 = viewport.clipPos.y + viewport.clipSize.y;

    int viewportX2 = viewport.pos.x + viewport.size.x;
    int viewportY2 = viewport.pos.y + viewport.size.y;

    viewport.clipPos.x = std::max(viewport.clipPos.x, viewport.pos.x);
    viewport.clipPos.y = std::max(viewport.clipPos.y, viewport.pos.y);
    viewport.clipSize.x = std::min(viewportX2, clipX2) + viewport.clipPos.x;
    viewport.clipSize.y = std::min(viewportY2, clipY2) + viewport.clipPos.y;

    viewport.apply(*m_cmdbuf);
}

void GraphicsContext::onResized()
{
    SDL_GetWindowSize(m_sdlwin, &m_windowWidth, &m_windowHeight);
    m_swapchainDirty = true;
}

void GraphicsContext::drawToolTip(int x, int y)
{
    if (m_tooltip != nullptr)
    {
        m_tooltip->draw(*this, x, y);
    }
}

ToolTip& GraphicsContext::createToolTip()
{
    if (m_tooltip != nullptr)
    {
        delete m_tooltip;
    }

    m_tooltip = new ToolTip();
    return *m_tooltip;
}

Vk::Swapchain& GraphicsContext::swapchain()
{
    return *m_swapchain;
}

Vk::CommandBuffer& GraphicsContext::precmdbuf()
{
    return *m_precmdbuf;
}

Vk::CommandBuffer& GraphicsContext::cmdbuf()
{
    return *m_cmdbuf;
}

std::shared_ptr<Vk::CommandBuffer> GraphicsContext::createAsyncCommandBuffer()
{
    std::shared_ptr<Vk::CommandPool> pool = std::make_shared<Vk::CommandPool>(
        m_device,
        VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT | VK_COMMAND_POOL_CREATE_TRANSIENT_BIT,
        m_device->getQueueFamilies().graphicsFamily.value()
    );

    return std::make_shared<Vk::CommandBuffer>(m_device, pool, VK_COMMAND_BUFFER_LEVEL_PRIMARY);
}

std::shared_ptr<Vk::Device> GraphicsContext::getDevice()
{
    return m_device;
}

void GraphicsContext::onNextFrame(CommandQueue<GraphicsContext>::Command cmd)
{
    m_numPendingSubmissions++;
    m_asyncCommandQueue.add(cmd);
}

Vk::Framebuffer& GraphicsContext::framebuffer()
{
    return *m_framebuffer;
}

int GraphicsContext::getThisFrameIndex()
{
    return m_thisFrame;
}

int GraphicsContext::getPendingSubmissions()
{
    return m_numPendingSubmissions;
}

uint64_t GraphicsContext::getFrameCounter()
{
    return m_frameCounter;
}

size_t GraphicsContext::getTotalVramUsage()
{
    size_t total = 0;

    VmaBudget budgets[VK_MAX_MEMORY_HEAPS];
    vmaGetHeapBudgets(m_device->getAllocator(), budgets);

    uint32_t i;
    for (i=0; i<m_device->getNumMemoryHeaps(); i++)
    {
        total += budgets[i].statistics.allocationBytes;
    }

    return total;
}

glm::ivec2 GraphicsContext::getAbsPos(glm::ivec2 relPos)
{
    VirtualViewport &viewport = m_viewportStack.top();

    return viewport.pos + relPos;
}

void GraphicsContext::requestScreenshot()
{
    screenshotRequested = true;
}

void GraphicsContext::centerMouse()
{
    int logical_w, logical_h;
    SDL_GetWindowSize(m_sdlwin, &logical_w, &logical_h);

    SDL_WarpMouseInWindow(
        m_sdlwin,
        (float) logical_w / 2,
        (float) logical_h / 2
    );
}

void GraphicsContext::holdResource(std::shared_ptr<void> res)
{
    if (res == nullptr)
    {
        return;
    }

    m_framesInFlight[m_thisFrame].heldResources.push_back(res);
}

std::shared_ptr<Material> GraphicsContext::createMaterialFromPixels(glm::ivec2 size, const RGBABytes *pixels)
{
    auto material = std::make_shared<Material>();
    material->widthValue = size.x;
    material->heightValue = size.y;

    material->addComponent(m_device, *m_precmdbuf, pixels);
    material->buildMetadata(m_device, *m_precmdbuf);
    material->buildDescriptorSet(m_device);

    return material;
}

glm::ivec2 GraphicsContext::getPhysicalMousePos()
{
    float lx, ly;
    SDL_GetMouseState(&lx, &ly);   // logical coords (float)

    int lw, lh;
    SDL_GetWindowSize(m_sdlwin, &lw, &lh);

    int pw, ph;
    SDL_GetWindowSizeInPixels(m_sdlwin, &pw, &ph);

    float sx = (float)pw / (float)lw;
    float sy = (float)ph / (float)lh;

    int px = (int)(lx * sx);
    int py = (int)(ly * sy);

    return glm::ivec2(px, py);
}

void GraphicsContext::maximizeWindow()
{
    SDL_MaximizeWindow(m_sdlwin);
    SDL_GetWindowSize(m_sdlwin, &m_windowWidth, &m_windowHeight);
}

void GraphicsContext::setRelativeMouseMode(bool enabled)
{
    SDL_SetWindowRelativeMouseMode(m_sdlwin, enabled);
}