#include "queue_context.hh"
#include "device_context.hh"
#include "layer_context.hh"
#include "timestamp_pool.hh"

#include <span>
#include <vulkan/vulkan_core.h>

namespace low_latency {

QueueContext::CommandPoolOwner::CommandPoolOwner(
    const QueueContext& queue_context)
    : queue_context(queue_context) {

    const auto& device_context = this->queue_context.device_context;

    const auto cpci = VkCommandPoolCreateInfo{
        .sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO,
        .flags = VK_COMMAND_POOL_CREATE_TRANSIENT_BIT |
                 VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT,
        .queueFamilyIndex = queue_context.queue_family_index,
    };

    THROW_NON_VKSUCCESS(device_context.vtable.CreateCommandPool(
        device_context.device, &cpci, nullptr, &this->command_pool));
}

QueueContext::CommandPoolOwner::~CommandPoolOwner() {
    const auto& device_context = this->queue_context.device_context;
    device_context.vtable.DestroyCommandPool(device_context.device,
                                             this->command_pool, nullptr);
}

QueueContext::QueueContext(DeviceContext& device_context, const VkQueue& queue,
                           const std::uint32_t& queue_family_index)
    : device_context(device_context), queue(queue),
      queue_family_index(queue_family_index),
      command_pool(std::make_unique<CommandPoolOwner>(*this)) {

    // Only construct a timestamp pool if we support it!
    if (device_context.physical_device.supports_required_extensions) {
        this->timestamp_pool = std::make_unique<TimestampPool>(*this);
    }
}

QueueContext::~QueueContext() {
    this->unpresented_submissions.clear();
    this->timestamp_pool.reset();
}

void QueueContext::notify_submit(
    const present_id_t& present_id,
    const std::shared_ptr<TimestampPool::Handle> head_handle,
    const std::shared_ptr<TimestampPool::Handle> tail_handle,
    const DeviceClock::time_point_t& now) {

    // Push this submission onto our unpresented_submissions at our present_id
    // mapping (might be empty, but handled with operator[]).
    auto& submissions = this->unpresented_submissions[present_id];
    if (submissions == nullptr) {
        submissions =
            std::make_shared<std::deque<std::unique_ptr<Submission>>>();
    }

    submissions->push_back(
        std::make_unique<Submission>(Submission{.head_handle = head_handle,
                                                .tail_handle = tail_handle,
                                                .cpu_present_time = now}));

    // This is probably hit if our queue never actually presents to anything,
    // because the only time we manually evict our unpresent_submissions is
    // when we present to something.
    if (std::size(*submissions) > this->MAX_TRACKED_SUBMISSIONS) {
        submissions->pop_front();
    }
}

void QueueContext::notify_present(const VkSwapchainKHR& swapchain,
                                  const present_id_t& present_id) {

    // Notify the device that this swapchain was just presented to.
    // We're avoiding a double hash here - don't use operator[] and erase.
    auto iter = this->unpresented_submissions.try_emplace(present_id).first;
    if (iter->second == nullptr) {
        iter->second =
            std::make_shared<std::deque<std::unique_ptr<Submission>>>();
    }

    this->device_context.notify_present(swapchain, iter->second);

    // Important, we nuke the submission because now it's presented.
    this->unpresented_submissions.erase(iter);
}

bool QueueContext::should_inject_timestamps() const {
    const auto& physical_device = this->device_context.physical_device;

    if (!physical_device.supports_required_extensions) {
        return false;
    }

    // Don't bother injecting timestamps during queue submission if we
    // aren't planning on doing anything anyway.
    if (!this->device_context.was_capability_requested &&
        !physical_device.instance.layer.is_antilag_1_enabled) {

        return false;
    }

    assert(physical_device.queue_properties);
    const auto& queue_props = *physical_device.queue_properties;
    assert(this->queue_family_index < std::size(queue_props));

    const auto& props = queue_props[this->queue_family_index];
    // Probably need at least 64, don't worry about it just yet and just ensure
    // it's not zero (because that will cause a crash if we inject).
    return props.queueFamilyProperties.timestampValidBits;
}

} // namespace low_latency