#include "Low_Latency_Renderable_Runtime.h"
#include "../../architecture/Renderable.h"
#include "../../base/Memory.h"
#include "Low_Latency_Roles.h"
#include "Triple_Buffer.h"
#include <algorithm>
#include <array>
#include <ostream>
#include <system_error>
#include <utility>

namespace renderive {

std::ostream& operator<<(std::ostream& stream, const Range& range) {
    stream << "Range(" << range.origin << ", " << range.target << ")";
    return stream;
}

namespace {

struct Input_Buffer_Slot {
    Input_Data* data{};
    std::uint64_t version = 0;
    bool pending = false;
    std::shared_ptr<Update_State> completion_state;
};

std::uint8_t state_role_index(Renderable_State_Role role) {
    switch (role) {
    case Renderable_State_Role::Edit:
        return State_Edit;
    case Renderable_State_Role::Ready:
        return State_Ready;
    case Renderable_State_Role::Render:
        return State_Render;
    }
    return State_Render;
}

std::uint8_t input_role_index(Renderable_Input_Role role) {
    switch (role) {
    case Renderable_Input_Role::Edit:
        return Input_Edit;
    case Renderable_Input_Role::Ready:
        return Input_Ready;
    case Renderable_Input_Role::Render:
        return Input_Render;
    }
    return Input_Render;
}

} // namespace

struct Low_Latency_Renderable_Runtime_Data {
    Renderable* owner{};
    Renderable_Runtime::Copy_Render_State_Function copy_render_state;
    Triple_Role_Buffer_Control state_control;
    std::array<Render_State*, 3> state_buffers{};
    std::array<Render_Object_Owner, 3> state_buffer_owners{};
    Triple_Role_Buffer_Control input_control;
    std::array<Input_Buffer_Slot, 3> input_buffers{};
    std::array<Render_Object_Owner, 3> input_buffer_owners{};
    std::pmr::vector<std::shared_ptr<Update_State>> committed_update_states{memory_resource(Memory_Domain::Update_Completion)};
    std::uint64_t input_version = 0;

    Render_State* state_by_index(std::uint8_t index) const {
        return index < state_buffers.size() ? state_buffers[index] : nullptr;
    }

    Render_State* state(std::uint8_t role) const {
        return state_by_index(state_control.read_view()[role]);
    }

    Input_Data* input_data_by_index(std::uint8_t index) const {
        return index < input_buffers.size() ? input_buffers[index].data : nullptr;
    }

    Input_Buffer_Slot& input_slot_by_index(std::uint8_t index) {
        return input_buffers[index];
    }

    const Input_Buffer_Slot& input_slot_by_index(std::uint8_t index) const {
        return input_buffers[index];
    }

    Input_Buffer_Slot& input_slot(std::uint8_t role) {
        return input_slot_by_index(input_control.read_view()[role]);
    }

    const Input_Buffer_Slot& input_slot(std::uint8_t role) const {
        return input_slot_by_index(input_control.read_view()[role]);
    }

    Input_Data* input_data(std::uint8_t role) const {
        return input_slot(role).data;
    }

    Triple_Buffer_Lease try_mark_state_role(std::uint8_t role) {
        auto record = state_control.try_acquire_role(role);
        return record.result == Triple_Buffer_Result::Success ? record.lease : Triple_Buffer_Lease{};
    }

    Triple_Buffer_Lease mark_state_role(std::uint8_t role) {
        auto record = state_control.acquire_role(role);
        return record.result == Triple_Buffer_Result::Success ? record.lease : Triple_Buffer_Lease{};
    }

    Triple_Buffer_Lease try_mark_input_role(std::uint8_t role) {
        auto record = input_control.try_acquire_role(role);
        return record.result == Triple_Buffer_Result::Success ? record.lease : Triple_Buffer_Lease{};
    }

    void mark_input_dirty(Triple_Buffer_Lease lease) {
        Input_Buffer_Slot& slot = input_slot_by_index(lease.index);
        slot.version = ++input_version;
        slot.pending = true;
    }

    void set_input_completion(Triple_Buffer_Lease lease, std::shared_ptr<Update_State> state) {
        Input_Buffer_Slot& slot = input_slot_by_index(lease.index);
        if (slot.completion_state && slot.completion_state != state)
            slot.completion_state->complete_all({}, Update_Outcome::Superseded);
        slot.completion_state = std::move(state);
    }

    void cancel_pending_completions() {
        for (Input_Buffer_Slot& slot : input_buffers) {
            if (!slot.completion_state)
                continue;
            slot.completion_state->complete_all(std::make_error_code(std::errc::operation_canceled), Update_Outcome::Cancelled);
            slot.completion_state.reset();
        }
        cancel_update_states(committed_update_states);
    }

    void clear_state_buffers() {
        state_control.close_and_wait();
        for (std::size_t i = 0; i < state_buffers.size(); ++i) {
            state_buffer_owners[i].reset();
            state_buffers[i] = nullptr;
        }
        state_control.reset();
    }

    void clear_input_buffers() {
        input_control.close_and_wait();
        cancel_pending_completions();
        for (std::size_t i = 0; i < input_buffers.size(); ++i) {
            input_buffer_owners[i].reset();
            input_buffers[i].data = nullptr;
            input_buffers[i].version = 0;
            input_buffers[i].pending = false;
            input_buffers[i].completion_state.reset();
        }
        input_version = 0;
        input_control.reset();
    }

    void initialize_buffers(Renderable& renderable) {
        clear_state_buffers();
        clear_input_buffers();
        state_buffer_owners[State_Edit] = renderable.create_render_state();
        state_buffer_owners[State_Ready] = renderable.create_render_state();
        state_buffer_owners[State_Render] = renderable.create_render_state();
        state_buffers[State_Edit] = state_buffer_owners[State_Edit].get<Render_State>();
        state_buffers[State_Ready] = state_buffer_owners[State_Ready].get<Render_State>();
        state_buffers[State_Render] = state_buffer_owners[State_Render].get<Render_State>();
        input_buffer_owners[Input_Edit] = renderable.create_input_data();
        input_buffer_owners[Input_Ready] = renderable.create_input_data();
        input_buffer_owners[Input_Render] = renderable.create_input_data();
        if (input_buffer_owners[Input_Edit])
            input_buffers[Input_Edit].data = input_buffer_owners[Input_Edit].get<Input_Data>();
        if (input_buffer_owners[Input_Ready])
            input_buffers[Input_Ready].data = input_buffer_owners[Input_Ready].get<Input_Data>();
        if (input_buffer_owners[Input_Render])
            input_buffers[Input_Render].data = input_buffer_owners[Input_Render].get<Input_Data>();
        if (state_buffers[State_Edit])
            state_buffers[State_Edit]->version = 1;
    }

    void sync_state_pipeline() {
        auto edit_record = state_control.try_acquire_role(State_Edit);
        if (edit_record.result != Triple_Buffer_Result::Success)
            return;
        auto ready_record = state_control.try_acquire_role(State_Ready);
        if (ready_record.result != Triple_Buffer_Result::Success) {
            state_control.unmark_use(edit_record.lease);
            return;
        }
        Render_State* ready = state_by_index(ready_record.lease.index);
        Render_State* edit = state_by_index(edit_record.lease.index);
        if (ready && edit && buffer_version_newer(edit->version, ready->version)) {
            if (copy_render_state)
                copy_render_state(*ready, *edit);
            else
                ready->version = edit->version;
        }
        state_control.unmark_use(ready_record.lease);
        state_control.unmark_use(edit_record.lease);
        state_control.try_swap_role(State_Ready, State_Render, [this](const Triple_Buffer_View& view) {
            auto* ready = state_by_index(view[State_Ready]);
            auto* render = state_by_index(view[State_Render]);
            return ready && render && buffer_version_newer(ready->version, render->version);
        });
        input_control.try_swap_role(Input_Edit, Input_Ready, [this](const Triple_Buffer_View& view) {
            const Input_Buffer_Slot& edit = input_slot_by_index(view[Input_Edit]);
            const Input_Buffer_Slot& ready = input_slot_by_index(view[Input_Ready]);
            return edit.pending && buffer_version_newer(edit.version, ready.version);
        });
        input_control.try_swap_role(Input_Ready, Input_Render, [this](const Triple_Buffer_View& view) {
            const Input_Buffer_Slot& ready = input_slot_by_index(view[Input_Ready]);
            const Input_Buffer_Slot& render = input_slot_by_index(view[Input_Render]);
            return ready.pending && buffer_version_newer(ready.version, render.version);
        });
    }

    void clear_input_role(Renderable_Input_Role role) {
        Input_Buffer_Slot& slot = input_slot(input_role_index(role));
        if (slot.completion_state) {
            if (role == Renderable_Input_Role::Render) {
                slot.completion_state->complete_until(Update_Stage::Committed);
                committed_update_states.push_back(std::move(slot.completion_state));
            }
            else {
                slot.completion_state->complete_all({}, Update_Outcome::Superseded);
                slot.completion_state.reset();
            }
        }
        if (slot.data)
            slot.data->clear();
        slot.pending = false;
    }

    void consume_input_lease(Triple_Buffer_Lease lease) {
        if (!lease)
            return;
        Input_Buffer_Slot& slot = input_slot_by_index(lease.index);
        if (slot.completion_state) {
            slot.completion_state->complete_until(Update_Stage::Committed);
            committed_update_states.push_back(std::move(slot.completion_state));
        }
        if (slot.data)
            slot.data->consume();
        slot.pending = false;
    }
};

namespace {

struct Low_Latency_State_Read_Handle final : Render_State_Read_Handle {
    Low_Latency_State_Read_Handle(Low_Latency_Renderable_Runtime_Data* runtime, Triple_Buffer_Lease lease)
        : runtime(runtime), lease(lease) {}
    ~Low_Latency_State_Read_Handle() override {
        if (runtime && lease)
            runtime->state_control.unmark_use(lease);
    }
    Low_Latency_Renderable_Runtime_Data* runtime{};
    Triple_Buffer_Lease lease{};
};

struct Low_Latency_Input_Frame_Handle final : Render_Input_Frame_Handle {
    Low_Latency_Input_Frame_Handle(Low_Latency_Renderable_Runtime_Data* runtime, Triple_Buffer_Lease lease)
        : runtime(runtime), lease(lease) {}
    ~Low_Latency_Input_Frame_Handle() override {
        if (runtime && lease)
            runtime->input_control.unmark_use(lease);
    }
    void consume_input(Input_Data*) override {
        if (consumed || !runtime || !lease)
            return;
        consumed = true;
        runtime->consume_input_lease(lease);
    }
    Low_Latency_Renderable_Runtime_Data* runtime{};
    Triple_Buffer_Lease lease{};
    bool consumed = false;
};

struct Low_Latency_State_Write_Handle final : Render_State_Write_Handle {
    Low_Latency_State_Write_Handle(Low_Latency_Renderable_Runtime_Data* runtime, Triple_Buffer_Lease lease)
        : runtime(runtime), lease(lease) {}
    ~Low_Latency_State_Write_Handle() override {
        if (!runtime || !lease)
            return;
        if (auto* state = runtime->state_by_index(lease.index))
            ++state->version;
        runtime->state_control.unmark_use(lease);
        if (runtime->owner)
            runtime->owner->mark_render_dirty(Dirty_Reason::Config);
    }
    Low_Latency_Renderable_Runtime_Data* runtime{};
    Triple_Buffer_Lease lease{};
};

struct Low_Latency_Input_Write_Handle final : Render_Input_Write_Handle {
    Low_Latency_Input_Write_Handle(Low_Latency_Renderable_Runtime_Data* runtime, Triple_Buffer_Lease lease, bool notify_render)
        : runtime(runtime), lease(lease), notify_render(notify_render) {}
    ~Low_Latency_Input_Write_Handle() override {
        if (!runtime || !lease)
            return;
        runtime->mark_input_dirty(lease);
        runtime->input_control.unmark_use(lease);
        if (notify_render && runtime->owner)
            runtime->owner->mark_render_dirty();
    }
    void set_completion(std::shared_ptr<Update_State> state) override {
        if (runtime && lease)
            runtime->set_input_completion(lease, std::move(state));
    }
    Low_Latency_Renderable_Runtime_Data* runtime{};
    Triple_Buffer_Lease lease{};
    bool notify_render{};
};

} // namespace

Low_Latency_Renderable_Runtime::Low_Latency_Renderable_Runtime(Renderable& owner)
    : d(std::make_unique<Low_Latency_Renderable_Runtime_Data>()) {
    d->owner = &owner;
    d->initialize_buffers(owner);
}

Low_Latency_Renderable_Runtime::~Low_Latency_Renderable_Runtime() {
    d->clear_state_buffers();
    d->clear_input_buffers();
}

void Low_Latency_Renderable_Runtime::set_render_state_copier(Copy_Render_State_Function copier) {
    d->copy_render_state = std::move(copier);
}

void Low_Latency_Renderable_Runtime::cancel_pending_completions() {
    d->cancel_pending_completions();
}

void Low_Latency_Renderable_Runtime::drain_committed_update_states(std::pmr::vector<std::shared_ptr<Update_State>>& output) {
    for (auto& state : d->committed_update_states)
        output.push_back(std::move(state));
    d->committed_update_states.clear();
}

Render_State_Read_Guard Low_Latency_Renderable_Runtime::acquire_state_read(Renderable_State_Role role) {
    Triple_Buffer_Lease lease = d->mark_state_role(state_role_index(role));
    if (!lease)
        return {};
    return Render_State_Read_Guard(d->state_by_index(lease.index), std::make_unique<Low_Latency_State_Read_Handle>(d.get(), lease));
}

Render_State_Read_Guard Low_Latency_Renderable_Runtime::acquire_render_state_read() {
    Triple_Buffer_Lease lease = d->try_mark_state_role(State_Render);
    if (!lease)
        return {};
    return Render_State_Read_Guard(d->state_by_index(lease.index), std::make_unique<Low_Latency_State_Read_Handle>(d.get(), lease));
}

Render_State_Read_Guard Low_Latency_Renderable_Runtime::acquire_edit_state_read() {
    Triple_Buffer_Lease lease = d->mark_state_role(State_Edit);
    if (!lease)
        return {};
    return Render_State_Read_Guard(d->state_by_index(lease.index), std::make_unique<Low_Latency_State_Read_Handle>(d.get(), lease));
}

Render_State_Write_Guard Low_Latency_Renderable_Runtime::acquire_edit_state() {
    Triple_Buffer_Lease lease = d->mark_state_role(State_Edit);
    if (!lease)
        return {};
    return Render_State_Write_Guard(d->state_by_index(lease.index), std::make_unique<Low_Latency_State_Write_Handle>(d.get(), lease));
}

Render_Input_Write_Guard Low_Latency_Renderable_Runtime::acquire_edit_input(bool notify_render) {
    auto record = d->input_control.acquire_role(Input_Edit);
    if (record.result != Triple_Buffer_Result::Success)
        return {};
    Input_Data* input = d->input_data_by_index(record.lease.index);
    if (!input) {
        d->input_control.unmark_use(record.lease);
        return {};
    }
    return Render_Input_Write_Guard(input, std::make_unique<Low_Latency_Input_Write_Handle>(d.get(), record.lease, notify_render));
}

Render_Input_Frame_Guard Low_Latency_Renderable_Runtime::acquire_input_read(Renderable_Input_Role role) {
    Triple_Buffer_Lease lease = d->try_mark_input_role(input_role_index(role));
    if (!lease)
        return {};
    Input_Data* input = d->input_data_by_index(lease.index);
    if (!input) {
        d->input_control.unmark_use(lease);
        return {};
    }
    return Render_Input_Frame_Guard(input, std::make_unique<Low_Latency_Input_Frame_Handle>(d.get(), lease));
}

Renderable_Frame_View Low_Latency_Renderable_Runtime::capture_frame_view() {
    d->sync_state_pipeline();
    Triple_Buffer_Lease state_lease = d->try_mark_state_role(State_Render);
    if (!state_lease)
        return {};
    auto state_guard = Render_State_Read_Guard(
        d->state_by_index(state_lease.index),
        std::make_unique<Low_Latency_State_Read_Handle>(d.get(), state_lease));

    Render_Input_Frame_Guard input_guard;
    Triple_Buffer_Lease input_lease = d->try_mark_input_role(Input_Render);
    if (input_lease) {
        Input_Data* input = d->input_data_by_index(input_lease.index);
        if (input) {
            input_guard = Render_Input_Frame_Guard(
                input,
                std::make_unique<Low_Latency_Input_Frame_Handle>(d.get(), input_lease));
        }
        else {
            d->input_control.unmark_use(input_lease);
        }
    }
    return Renderable_Frame_View(std::move(state_guard), std::move(input_guard));
}

Render_State* Low_Latency_Renderable_Runtime::state(Renderable_State_Role role) {
    return d->state(state_role_index(role));
}

Input_Data* Low_Latency_Renderable_Runtime::input_data(Renderable_Input_Role role) {
    return d->input_data(input_role_index(role));
}

void Low_Latency_Renderable_Runtime::clear_render_input() {
    d->clear_input_role(Renderable_Input_Role::Render);
}

std::unique_ptr<Renderable_Runtime> make_low_latency_renderable_runtime(Renderable& owner) {
    return std::make_unique<Low_Latency_Renderable_Runtime>(owner);
}

} // namespace renderive