#include "Capture.hpp"
#include <algorithm>
#include <limits>
#include <stdexcept>

Capture_Session_Id Capture_Controller::capture_next_frame() {
    return capture_frames(1);
}

Capture_Session_Id Capture_Controller::capture_frames(std::size_t count) {
    if (count == 0)
        throw std::invalid_argument("capture frame count must be positive");
    if (count > std::numeric_limits<std::uint32_t>::max())
        throw std::invalid_argument("capture frame count is too large");
    Capture_Session_Id id = next_session_id_.load(std::memory_order_relaxed);
    for (;;) {
        if (id == 0 || id == std::numeric_limits<Capture_Session_Id>::max())
            throw std::overflow_error("capture session id exhausted");
        if (next_session_id_.compare_exchange_weak(
                id, id + 1, std::memory_order_relaxed,
                std::memory_order_relaxed))
            break;
    }
    auto next = std::make_shared<Request>(id, count);
    auto current = request_.load(std::memory_order_acquire);
    for (;;) {
        if (current && current->frames.load(std::memory_order_acquire) != 0)
            throw std::logic_error("a capture session is already active");
        if (request_.compare_exchange_weak(current, next,
                                           std::memory_order_acq_rel,
                                           std::memory_order_acquire))
            break;
    }
    return id;
}

Capture_Frame_Ticket Capture_Controller::begin_frame() noexcept {
    const auto request = request_.load(std::memory_order_acquire);
    if (!request)
        return {};
    std::uint64_t state = request->frames.load(std::memory_order_acquire);
    while (Request::remaining(state) != 0) {
        const auto next = Request::pack(Request::remaining(state) - 1,
                                        Request::in_flight(state) + 1);
        if (request->frames.compare_exchange_weak(
                state, next, std::memory_order_acq_rel, std::memory_order_acquire))
            return {request->session_id, true};
    }
    return {};
}

void Capture_Controller::finish_frame(Capture_Frame_Ticket ticket,
                                      bool completed) noexcept {
    if (!ticket.capture)
        return;
    const auto request = request_.load(std::memory_order_acquire);
    if (!request || request->session_id != ticket.session_id)
        return;
    std::uint64_t state = request->frames.load(std::memory_order_acquire);
    while (Request::in_flight(state) != 0) {
        const auto next = Request::pack(
            Request::remaining(state) + (completed ? 0 : 1),
            Request::in_flight(state) - 1);
        if (request->frames.compare_exchange_weak(
                state, next, std::memory_order_acq_rel, std::memory_order_acquire))
            return;
    }
}

Capture_Controller_State Capture_Controller::state() const noexcept {
    const auto request = request_.load(std::memory_order_acquire);
    if (!request)
        return {};
    const auto state = request->frames.load(std::memory_order_acquire);
    return {request->session_id,
            Request::remaining(state) + Request::in_flight(state)};
}

void Capture_Repository::begin_session(Capture_Session_Id session_id,
                                       std::size_t requested_count) {
    if (session_id == 0 || requested_count == 0)
        throw std::invalid_argument("invalid capture session");
    std::lock_guard lock(mutex_);
    if (std::any_of(sessions_.begin(), sessions_.end(), [session_id](const auto& session) {
            return session.session_id == session_id;
        }))
        throw std::invalid_argument("capture session already exists");
    sessions_.push_back({session_id, requested_count, {}});
    sessions_.back().frames.reserve(requested_count);
}

void Capture_Repository::publish(Capture_Session_Id session_id,
                                 std::shared_ptr<const Frame_Snapshot> snapshot,
                                 Frame_Analysis analysis) {
    if (!snapshot)
        throw std::invalid_argument("captured frame snapshot is null");
    if (analysis.frame_id != snapshot->frame_id ||
        analysis.render_plan_version != snapshot->render_plan_version)
        throw std::invalid_argument("captured frame analysis does not match snapshot");
    std::lock_guard lock(mutex_);
    const auto iterator = std::find_if(sessions_.begin(), sessions_.end(),
        [session_id](const auto& session) { return session.session_id == session_id; });
    if (iterator == sessions_.end())
        throw std::invalid_argument("capture session does not exist");
    if (!iterator->active())
        throw std::logic_error("capture session is complete");
    iterator->frames.push_back({std::move(snapshot), std::move(analysis)});
}

std::optional<Capture_Session> Capture_Repository::session(
    Capture_Session_Id session_id) const {
    std::lock_guard lock(mutex_);
    const auto iterator = std::find_if(sessions_.begin(), sessions_.end(),
        [session_id](const auto& session) { return session.session_id == session_id; });
    return iterator == sessions_.end() ? std::nullopt
                                       : std::optional<Capture_Session>(*iterator);
}

std::vector<Capture_Session_Id> Capture_Repository::sessions() const {
    std::lock_guard lock(mutex_);
    std::vector<Capture_Session_Id> result;
    result.reserve(sessions_.size());
    for (const auto& session : sessions_)
        result.push_back(session.session_id);
    return result;
}
