#include "Frame_Target_Pool.h"

#include <stdexcept>
#include <utility>

namespace renderive::detail {
namespace {

template <class Resource, class Allocate>
Resource* allocate_wrapper(Allocate allocate, const char* name) {
    Resource* resource = allocate();
    if (resource == nullptr)
        throw std::runtime_error(name);
    return resource;
}

} // namespace

Frame_Target::Frame_Target(
    DvzGpuCtx* gpu_context, Extent_2D extent, std::uint64_t generation)
    : gpu_context_(gpu_context), extent_(extent), generation_(generation) {
    if (gpu_context == nullptr || extent.empty())
        throw std::invalid_argument("invalid Render_3D frame target configuration");
    DvzDevice* device = dvz_gpu_ctx_device(gpu_context);
    DvzVma* allocator = dvz_gpu_ctx_alloc(gpu_context);
    DvzQueue* queue = dvz_gpu_ctx_queue(gpu_context, DVZ_QUEUE_MAIN);
    if (device == nullptr || allocator == nullptr || queue == nullptr)
        throw std::invalid_argument("Render_3D frame target requires a complete GPU context");

    try {
        image_ = allocate_wrapper<DvzImages>(dvz_images_create_wrapper, "failed to allocate Datoviz image wrapper");
        dvz_images(device, allocator, VK_IMAGE_TYPE_2D, 1, image_);
        dvz_images_format(image_, VK_FORMAT_R8G8B8A8_UNORM);
        dvz_images_size(image_, extent.width, extent.height, 1);
        dvz_images_tiling(image_, VK_IMAGE_TILING_OPTIMAL);
        dvz_images_usage(image_, VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT);
        dvz_images_alloc_flags(image_, DVZ_ALLOC_FLAGS_NONE);
        if (dvz_images_create(image_) != 0)
            throw std::runtime_error("failed to create Render_3D color image");

        image_view_ = allocate_wrapper<DvzImageViews>(dvz_image_views_create_wrapper, "failed to allocate Datoviz image-view wrapper");
        dvz_image_views(image_, image_view_);
        dvz_image_views_type(image_view_, VK_IMAGE_VIEW_TYPE_2D);
        dvz_image_views_aspect(image_view_, VK_IMAGE_ASPECT_COLOR_BIT);
        dvz_image_views_mip(image_view_, 0, 1);
        dvz_image_views_layers(image_view_, 0, 1);
        if (dvz_image_views_create(image_view_) != 0)
            throw std::runtime_error("failed to create Render_3D color image view");

        commands_ = allocate_wrapper<DvzCommands>(dvz_commands_create_wrapper, "failed to allocate Datoviz command wrapper");
        dvz_commands(device, queue, 1, commands_);
        if (dvz_commands_handle(commands_) == VK_NULL_HANDLE)
            throw std::runtime_error("failed to create Render_3D command buffer");

        fence_ = allocate_wrapper<DvzFence>(dvz_fence_create_wrapper, "failed to allocate Datoviz fence wrapper");
        dvz_fence(device, true, fence_);
        if (dvz_fence_handle(fence_) == VK_NULL_HANDLE)
            throw std::runtime_error("failed to create Render_3D fence");

        submit_ = allocate_wrapper<DvzSubmit>(dvz_submit_create_wrapper, "failed to allocate Datoviz submit wrapper");
    } catch (...) {
        destroy();
        throw;
    }
}

Frame_Target::~Frame_Target() {
    destroy();
}

bool Frame_Target::begin() {
    wait();
    dvz_cmd_reset(commands_);
    if (dvz_cmd_begin_result(commands_) != 0)
        return false;
    recording_ = true;
    if (completed_layout_ == VK_IMAGE_LAYOUT_UNDEFINED)
        transition_to_color_attachment();
    return true;
}

DvzStreamFrame Frame_Target::stream_frame() const {
    DvzStreamFrame frame{};
    frame.image = dvz_image_handle(image_, 0);
    frame.command_buffer = dvz_commands_handle(commands_);
    frame.image_view = dvz_image_views_handle(image_view_, 0);
    frame.extent = {extent_.width, extent_.height};
    frame.color_format = VK_FORMAT_R8G8B8A8_UNORM;
    frame.image_layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
    frame.usage = DVZ_STREAM_FRAME_USAGE_RENDER_TARGET | DVZ_STREAM_FRAME_USAGE_COPY_SRC;
    frame.command_buffer_recording = recording_;
    frame.image_borrowed = true;
    frame.image_view_borrowed = true;
    frame.command_buffer_borrowed = true;
    frame.handles_dirty = true;
    frame.resource_generation = generation_;
    frame.image_valid = frame.image != VK_NULL_HANDLE && frame.image_view != VK_NULL_HANDLE;
    frame.memory_fd = -1;
    frame.wait_semaphore_fd = -1;
    return frame;
}

bool Frame_Target::submit_and_wait() {
    if (!recording_ || dvz_cmd_end_result(commands_) != 0) {
        abort();
        return false;
    }
    recording_ = false;
    dvz_fence_reset(fence_);
    dvz_submit(submit_);
    dvz_submit_command(submit_, dvz_commands_handle(commands_));
    DvzQueue* queue = dvz_gpu_ctx_queue(gpu_context_, DVZ_QUEUE_MAIN);
    if (dvz_submit_send(submit_, dvz_queue_handle(queue), dvz_fence_handle(fence_)) != VK_SUCCESS)
        return false;
    submitted_ = true;
    if (!dvz_fence_wait(fence_)) {
        dvz_device_wait(dvz_gpu_ctx_device(gpu_context_));
        submitted_ = false;
        return false;
    }
    submitted_ = false;
    completed_layout_ = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
    return true;
}

void Frame_Target::abort() noexcept {
    if (submitted_)
        wait();
    if (recording_ && commands_ != nullptr)
        dvz_cmd_reset(commands_);
    recording_ = false;
}

void Frame_Target::wait() noexcept {
    if (submitted_ && fence_ != nullptr) {
        if (!dvz_fence_wait(fence_))
            dvz_device_wait(dvz_gpu_ctx_device(gpu_context_));
        submitted_ = false;
        completed_layout_ = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
    }
}

void Frame_Target::transition_to_color_attachment() {
    DvzBarriers barriers{};
    dvz_barriers(&barriers);
    auto* barrier = dvz_barriers_image(&barriers, dvz_image_handle(image_, 0));
    if (barrier == nullptr)
        throw std::runtime_error("failed to allocate Render_3D image barrier");
    dvz_barrier_image_stage(barrier, VK_PIPELINE_STAGE_2_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT);
    dvz_barrier_image_access(barrier, 0, VK_ACCESS_2_COLOR_ATTACHMENT_READ_BIT | VK_ACCESS_2_COLOR_ATTACHMENT_WRITE_BIT);
    dvz_barrier_image_layout(barrier, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL);
    dvz_barrier_image_aspect(barrier, VK_IMAGE_ASPECT_COLOR_BIT);
    dvz_barrier_image_mip(barrier, 0, 1);
    dvz_barrier_image_layers(barrier, 0, 1);
    dvz_cmd_barriers(commands_, &barriers);
}

void Frame_Target::destroy() noexcept {
    wait();
    if (fence_ != nullptr) {
        dvz_fence_destroy(fence_);
        dvz_fence_free(fence_);
        fence_ = nullptr;
    }
    if (submit_ != nullptr) {
        dvz_submit_free(submit_);
        submit_ = nullptr;
    }
    if (commands_ != nullptr) {
        dvz_commands_destroy(commands_);
        dvz_commands_free(commands_);
        commands_ = nullptr;
    }
    if (image_view_ != nullptr) {
        dvz_image_views_destroy(image_view_);
        dvz_image_views_free(image_view_);
        image_view_ = nullptr;
    }
    if (image_ != nullptr) {
        dvz_images_destroy(image_);
        dvz_images_free(image_);
        image_ = nullptr;
    }
}

Frame_Target_Pool::Frame_Target_Pool(DvzGpuCtx* gpu_context, std::size_t slot_count)
    : gpu_context_(gpu_context), slot_count_(slot_count) {
    if (gpu_context == nullptr || slot_count < 2)
        throw std::invalid_argument("Render_3D requires at least two valid GPU frame slots");
}

Frame_Target_Pool::~Frame_Target_Pool() {
    wait_all();
}

Frame_Target& Frame_Target_Pool::acquire(
    std::uint64_t frame_sequence, Extent_2D extent) {
    if (frame_sequence == 0 || extent.empty())
        throw std::invalid_argument("cannot acquire an empty Render_3D target");
    if (targets_.empty() || targets_.front()->extent_ != extent) {
        auto replacement = make_generation(extent);
        wait_all();
        for (auto& target : targets_)
            retired_targets_.push_back(std::move(target));
        targets_ = std::move(replacement);
    }
    Frame_Target& target = *targets_[(frame_sequence - 1) % targets_.size()];
    target.wait();
    return target;
}

void Frame_Target_Pool::external_target_replaced() {
    retired_targets_.clear();
}

void Frame_Target_Pool::runtime_released_targets() {
    retired_targets_.clear();
}

void Frame_Target_Pool::wait_all() noexcept {
    for (auto& target : targets_)
        target->wait();
    for (auto& target : retired_targets_)
        target->wait();
}

std::vector<std::unique_ptr<Frame_Target>> Frame_Target_Pool::make_generation(Extent_2D extent) {
    const std::uint64_t next_generation = generation_ + 1;
    std::vector<std::unique_ptr<Frame_Target>> result;
    result.reserve(slot_count_);
    for (std::size_t index = 0; index < slot_count_; ++index)
        result.push_back(std::make_unique<Frame_Target>(gpu_context_, extent, next_generation));
    generation_ = next_generation;
    return result;
}

} // namespace renderive::detail
