#include "Scene_Base.hpp"
#include <algorithm>
#include <memory_resource>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <utility>
#include <taskflow/taskflow.hpp>
#include "renderive/renderable/color/Color_Cache.hpp"
static_assert(TF_VERSION == 400100, "Renderive requires Taskflow 4.1.0");
class Scene_Base::Execution_Context {
public:
    static tf::Executor& executor() {
        static tf::Executor executor;
        return executor;
    }
};
class Scene_Base::Render_Execution_Scope {
public:
    explicit Render_Execution_Scope(Scene_Base& scene) noexcept : previous_(std::exchange(active_execution_scene_, &scene)) {}
    ~Render_Execution_Scope() {
        active_execution_scene_ = previous_;
    }
private:
    Scene_Base* previous_{};
};
Scene_Base::Scene_Base() : Scene_Base(*std::pmr::get_default_resource()) {}
Scene_Base::Scene_Base(std::pmr::memory_resource& upstream_memory_resource)
    : scene_lifetime_(std::make_shared<Scene_Lifetime>(*this)), memory_domain_(std::allocate_shared<Scene_Memory_Domain>(std::pmr::polymorphic_allocator<Scene_Memory_Domain>(&upstream_memory_resource), upstream_memory_resource)), renderable_states_{Renderable_List(&memory_domain_->resource()), Renderable_List(&memory_domain_->resource())}, render_renderables_(&renderable_states_[0]), cache_renderables_(&renderable_states_[1]), task_(memory_domain_->resource()) {
    worker_ = std::thread([this] {
        render_loop();
    });
}
Scene_Base::~Scene_Base() {
    shutdown();
}
void Scene_Base::render() {
    if (active_submitted_observer_scene_ == this) {
        std::lock_guard<std::recursive_mutex> lock(task_mutex_);
        ++deferred_render_count_;
        return;
    }
    auto task_lock = lock_render_idle();
    if (pending_exception_observed_) {
        pending_exception_ = nullptr;
        pending_exception_observed_ = false;
    }
    if (!pending_exception_ && current_completion_ && current_completion_->completed && current_completion_->exception && !current_completion_->observed) {
        pending_exception_ = current_completion_->exception;
    }
    Render_Task task(memory_resource());
    {
        std::lock_guard<std::mutex> renderable_lock(renderable_mutex_);
        std::swap(render_renderables_, cache_renderables_);
        *cache_renderables_ = *render_renderables_;
        prepare_render_task(task, *render_renderables_);
    }
    auto& strategy = frame_control_strategy();
    strategy.swap();
    task.frame_control_state = strategy.frame_control_state();
    task.scene_state_revision = acquire_scene_state();
    task.render_sequence = ++render_sequence_;
    task.completion = std::make_shared<Render_Completion>();
    current_completion_ = task.completion;
    const Observation submitted_observation{Observation_Event::render_submitted, observer_now_ns(), task.render_sequence, task.scene_state_revision, task.render_order.size()};
    task_ = std::move(task);
    task_pending_ = true;
    task_lock.unlock();
    Scene_Base* previous_observer_scene = std::exchange(active_submitted_observer_scene_, this);
    observe_scene(submitted_observation);
    active_submitted_observer_scene_ = previous_observer_scene;
    task_ready_.notify_one();
    std::size_t deferred_render_count{};
    {
        std::lock_guard<std::recursive_mutex> lock(task_mutex_);
        deferred_render_count = std::exchange(deferred_render_count_, 0);
    }
    for (std::size_t index = 0; index < deferred_render_count; ++index) {
        render();
    }
}
void Scene_Base::wait_for_render() {
    if (is_render_worker_thread() || active_submitted_observer_scene_ == this) {
        return;
    }
    std::unique_lock<std::recursive_mutex> lock(task_mutex_);
    const auto completion = current_completion_;
    if (!completion) {
        return;
    }
    render_completed_.wait(lock, [&completion] {
        return completion->completed;
    });
    std::exception_ptr exception;
    if (pending_exception_) {
        exception = pending_exception_;
        pending_exception_observed_ = true;
        if (!completion->exception) {
            completion->observed = true;
        }
    } else {
        exception = completion->exception;
        completion->observed = true;
    }
    lock.unlock();
    if (exception) {
        std::rethrow_exception(exception);
    }
}
void Scene_Base::attach_renderable(Renderable renderable) {
    if (!renderable) {
        throw std::invalid_argument("renderable is null");
    }
    auto task_lock = lock_render_idle();
    validate_renderable_scene(*renderable);
    std::lock_guard<std::mutex> lock(renderable_mutex_);
    if (is_renderable_attached_locked(*renderable)) {
        return;
    }
    cache_renderables_->reserve(cache_renderables_->size() + 1);
    on_renderable_attached(*renderable);
    cache_renderables_->push_back(std::move(renderable));
}
void Scene_Base::detach_renderable(Renderable_Base& renderable) {
    auto task_lock = lock_render_idle();
    validate_renderable_scene(renderable);
    std::lock_guard<std::mutex> lock(renderable_mutex_);
    if (!is_renderable_attached_locked(renderable)) {
        return;
    }
    on_renderable_detached(renderable);
    const auto detached = [&renderable](const Renderable& value) {
        return value.get() == &renderable;
    };
    std::erase_if(*cache_renderables_, detached);
    std::erase_if(*render_renderables_, detached);
}
void Scene_Base::set_display_parent(Renderable_Base& renderable, Renderable_Base* parent) {
    auto task_lock = lock_render_idle();
    validate_renderable_scene(renderable);
    if (parent) {
        validate_renderable_scene(*parent);
    }
    std::lock_guard<std::mutex> lock(renderable_mutex_);
    validate_renderable_attached_locked(renderable);
    if (parent) {
        validate_renderable_attached_locked(*parent);
    }
    auto& node = layer_node(renderable);
    auto* target = parent ? &layer_node(*parent) : display_root_node();
    if (node.parent() == target) {
        return;
    }
    if (target) {
        target->append_child(node);
    } else {
        node.detach();
    }
}
void Scene_Base::set_dependency_parent(Renderable_Base& renderable, Renderable_Base* parent) {
    auto task_lock = lock_render_idle();
    validate_renderable_scene(renderable);
    if (parent) {
        validate_renderable_scene(*parent);
    }
    std::lock_guard<std::mutex> lock(renderable_mutex_);
    validate_renderable_attached_locked(renderable);
    if (parent) {
        validate_renderable_attached_locked(*parent);
    }
    auto& node = dependency_node(renderable);
    auto* target = parent ? &dependency_node(*parent) : dependency_root_node();
    if (node.parent() == target) {
        return;
    }
    if (target) {
        target->append_child(node);
    } else {
        node.detach();
    }
    renderable.invalidate_cache();
}
void Scene_Base::set_renderable_configuration(Renderable_Base& renderable, Renderable_Configuration configuration) {
    auto task_lock = lock_render_idle();
    validate_renderable_scene(renderable);
    std::lock_guard<std::mutex> lock(renderable_mutex_);
    renderable.set_configuration(configuration);
}
std::size_t Scene_Base::renderable_count() const {
    std::lock_guard<std::mutex> lock(renderable_mutex_);
    return cache_renderables_->size();
}
Scene_Base::Topology_Snapshot Scene_Base::topology_snapshot() const {
    std::lock_guard<std::mutex> lock(renderable_mutex_);
    Topology_Snapshot snapshot;
    std::unordered_map<const Renderable_Base*, Const_Renderable> renderables;
    renderables.reserve(cache_renderables_->size());
    snapshot.renderables.reserve(cache_renderables_->size());
    snapshot.display.reserve(cache_renderables_->size());
    snapshot.dependency.reserve(cache_renderables_->size());
    for (const Renderable& renderable : *cache_renderables_) {
        Const_Renderable value = renderable;
        renderables.emplace(renderable.get(), value);
        snapshot.renderables.push_back(std::move(value));
    }
    for (const Renderable& renderable : *cache_renderables_) {
        const auto* display_parent_node = renderable->layer_node_.parent();
        const auto* dependency_parent_node = renderable->dependency_node_.parent();
        const auto* display_parent = display_parent_node ? display_parent_node->owner() : nullptr;
        const auto* dependency_parent = dependency_parent_node ? dependency_parent_node->owner() : nullptr;
        snapshot.display.push_back({renderables.at(renderable.get()), display_parent ? renderables.at(display_parent) : Const_Renderable{}});
        snapshot.dependency.push_back({renderables.at(renderable.get()), dependency_parent ? renderables.at(dependency_parent) : Const_Renderable{}});
    }
    return snapshot;
}
std::pmr::memory_resource& Scene_Base::memory_resource() const noexcept {
    return memory_domain_->resource();
}
std::pmr::memory_resource& Scene_Base::upstream_memory_resource() const noexcept {
    return memory_domain_->upstream();
}
Frame_Control_Strategy_Base& Scene_Base::frame_control_strategy() {
    return frame_control_strategy_impl();
}
const Frame_Control_Strategy_Base& Scene_Base::frame_control_strategy() const {
    return frame_control_strategy_impl();
}
Frame_Control_Strategy_Base& Scene_Base::frame_control_strategy_impl() {
    throw std::logic_error("frame control strategy is not available");
}
const Frame_Control_Strategy_Base& Scene_Base::frame_control_strategy_impl() const {
    throw std::logic_error("frame control strategy is not available");
}
Renderable_Base::Layer_Node& Scene_Base::layer_node(Renderable_Base& renderable) noexcept {
    return renderable.layer_node_;
}
Renderable_Base::Dependency_Node& Scene_Base::dependency_node(Renderable_Base& renderable) noexcept {
    return renderable.dependency_node_;
}
Renderable_Base::Layer_Node* Scene_Base::display_root_node() noexcept {
    return nullptr;
}
Renderable_Base::Dependency_Node* Scene_Base::dependency_root_node() noexcept {
    return nullptr;
}
void Scene_Base::on_renderable_attached(Renderable_Base&) {}
void Scene_Base::on_renderable_detached(Renderable_Base&) {}
void Scene_Base::prepare_render_task(Render_Task& task, const Renderable_List& renderables) {
    task.render_order = renderables;
    task.display_order = renderables;
}
Color_Cache* Scene_Base::prepare_renderable_cache(Renderable_Base&, bool) {
    return nullptr;
}
void Scene_Base::generate_final_color_cache(const Scene_Render_Context&, const Renderable_List&) {}
std::uint64_t Scene_Base::acquire_scene_state() {
    return 0;
}
void Scene_Base::observe_scene(const Observation&) noexcept {}
std::uint64_t Scene_Base::observer_now_ns() const noexcept {
    return 0;
}
std::unique_lock<std::recursive_mutex> Scene_Base::lock_render_idle() {
    if (is_render_worker_thread()) {
        throw std::logic_error("scene control mutation is not allowed during render execution");
    }
    if (active_submitted_observer_scene_ == this) {
        throw std::logic_error("scene control mutation is not allowed during render submission observation");
    }
    std::unique_lock<std::recursive_mutex> lock(task_mutex_);
    render_completed_.wait(lock, [this] {
        return !task_pending_ && !rendering_;
    });
    return lock;
}
bool Scene_Base::is_render_worker_thread() const noexcept {
    return active_execution_scene_ == this || (worker_.joinable() && std::this_thread::get_id() == worker_.get_id());
}
void Scene_Base::shutdown() noexcept {
    {
        std::unique_lock<std::recursive_mutex> lock(task_mutex_);
        render_completed_.wait(lock, [this] {
            return !task_pending_ && !rendering_;
        });
        stop_ = true;
    }
    task_ready_.notify_one();
    if (worker_.joinable()) {
        worker_.join();
    }
    scene_lifetime_->invalidate();
}
void Scene_Base::render_loop() {
    for (;;) {
        Render_Task task(memory_resource());
        {
            std::unique_lock<std::recursive_mutex> lock(task_mutex_);
            task_ready_.wait(lock, [this] {
                return stop_ || task_pending_;
            });
            if (stop_) {
                return;
            }
            task = std::move(task_);
            task_pending_ = false;
            rendering_ = true;
        }
        observe_scene({Observation_Event::render_started, observer_now_ns(), task.render_sequence, task.scene_state_revision, task.render_order.size()});
        std::exception_ptr exception;
        try {
            execute_taskflow(task);
            observe_scene({Observation_Event::render_completed, observer_now_ns(), task.render_sequence, task.scene_state_revision, task.render_order.size()});
        } catch (...) {
            exception = std::current_exception();
            observe_scene({Observation_Event::render_failed, observer_now_ns(), task.render_sequence, task.scene_state_revision, task.render_order.size()});
        }
        {
            std::lock_guard lock(task_mutex_);
            task.completion->exception = exception;
            task.completion->completed = true;
            rendering_ = false;
        }
        render_completed_.notify_all();
    }
}
void Scene_Base::execute_taskflow(Render_Task& task) {
    struct Module {
        tf::Taskflow graph;
        tf::Task module_task;
        bool render{};
    };
    std::pmr::monotonic_buffer_resource scratch_resource(&memory_resource());
    tf::Taskflow scene_graph;
    std::pmr::unordered_map<Renderable_Base*, std::size_t> indices(&scratch_resource);
    indices.reserve(task.render_order.size());
    std::pmr::vector<bool> render_flags(task.render_order.size(), false, &scratch_resource);
    for (std::size_t index = 0; index < task.render_order.size(); ++index) {
        indices.emplace(task.render_order[index].get(), index);
        render_flags[index] = task.render_order[index]->requires_render();
    }
    std::pmr::vector<std::size_t> dirty_queue(&scratch_resource);
    dirty_queue.reserve(task.render_order.size());
    for (std::size_t index = 0; index < task.render_order.size(); ++index) {
        if (render_flags[index]) {
            dirty_queue.push_back(index);
        }
    }
    for (std::size_t queue_index = 0; queue_index < dirty_queue.size(); ++queue_index) {
        auto& node = dependency_node(*task.render_order[dirty_queue[queue_index]]);
        for (std::size_t child_index = 0; child_index < node.child_count(); ++child_index) {
            Renderable_Base* child = node.child(child_index).owner();
            auto child_iterator = indices.find(child);
            if (child_iterator == indices.end() || render_flags[child_iterator->second]) {
                continue;
            }
            render_flags[child_iterator->second] = true;
            dirty_queue.push_back(child_iterator->second);
        }
    }
    std::pmr::vector<Module> modules(&scratch_resource);
    modules.resize(task.render_order.size());
    for (std::size_t index = 0; index < task.render_order.size(); ++index) {
        Renderable_Base& renderable = *task.render_order[index];
        Module& module = modules[index];
        module.render = render_flags[index];
        if (module.render) {
            const std::uint64_t render_revision = renderable.cache_revision();
            Color_Cache* color_cache = prepare_renderable_cache(renderable, true);
            const Scene_Render_Context context{this, task.frame_control_state, task.scene_state_revision, task.render_sequence, &renderable, color_cache};
            const auto task_graph = renderable.task_graph();
            const auto& nodes = task_graph->nodes();
            std::pmr::vector<tf::Task> internal_tasks(&scratch_resource);
            internal_tasks.reserve(nodes.size());
            for (const auto& node : nodes) {
                internal_tasks.push_back(module.graph.emplace([this, function = node.function, context] {
                    Render_Execution_Scope scope(*this);
                    function(context);
                }).name(std::string(node.name)));
            }
            for (std::size_t node_index = 0; node_index < nodes.size(); ++node_index) {
                for (std::size_t successor : nodes[node_index].successors) {
                    internal_tasks[node_index].precede(internal_tasks.at(successor));
                }
            }
            tf::Task mark_rendered = module.graph.emplace([&renderable, render_revision] {
                renderable.mark_rendered(render_revision);
            }).name("cache_publish");
            if (internal_tasks.empty()) {
                module.graph.emplace([] {}).precede(mark_rendered);
            } else {
                for (std::size_t node_index = 0; node_index < nodes.size(); ++node_index) {
                    if (nodes[node_index].successors.empty()) {
                        internal_tasks[node_index].precede(mark_rendered);
                    }
                }
            }
        } else {
            module.graph.emplace([] {}).name("cache_hit");
        }
        module.module_task = scene_graph.composed_of(module.graph).name("renderable");
    }
    for (std::size_t index = 0; index < task.render_order.size(); ++index) {
        auto& node = dependency_node(*task.render_order[index]);
        Renderable_Base* parent = node.parent() ? node.parent()->owner() : nullptr;
        auto parent_iterator = indices.find(parent);
        if (parent_iterator != indices.end()) {
            modules[parent_iterator->second].module_task.precede(modules[index].module_task);
        }
    }
    const Scene_Render_Context final_context{this, task.frame_control_state, task.scene_state_revision, task.render_sequence, nullptr, nullptr};
    tf::Task final_task = scene_graph.emplace([this, final_context, &task] {
        Render_Execution_Scope scope(*this);
        generate_final_color_cache(final_context, task.display_order);
    }).name("final_color_cache");
    if (modules.empty()) {
        scene_graph.emplace([] {}).precede(final_task);
    } else {
        for (Module& module : modules) {
            module.module_task.precede(final_task);
        }
    }
    Execution_Context::executor().run(scene_graph).get();
}
void Scene_Base::validate_renderable_scene(const Renderable_Base& renderable) const {
    if (renderable.real_time_data_state_->scene_lifetime.get() != scene_lifetime_.get()) {
        throw std::invalid_argument("renderable belongs to another scene");
    }
}
bool Scene_Base::is_renderable_attached_locked(const Renderable_Base& renderable) const {
    return std::find_if(cache_renderables_->begin(), cache_renderables_->end(), [&renderable](const Renderable& value) {
        return value.get() == &renderable;
    }) != cache_renderables_->end();
}
void Scene_Base::validate_renderable_attached_locked(const Renderable_Base& renderable) const {
    if (!is_renderable_attached_locked(renderable)) {
        throw std::invalid_argument("renderable is not attached to this scene");
    }
}