#include "global.h"
#include <algorithm>
#include "Memory.h"
#include "../plot/Plot_Core.h"
#include "asio.hpp"
#include <chrono>

namespace renderive {

struct Render_Scheduler_Private {
    asio::io_context io_context;
    std::unique_ptr<asio::executor_work_guard<asio::io_context::executor_type>> work_guard;
    std::unique_ptr<asio::thread_pool> cpu_pool;
    std::thread scheduler_thread;
    std::thread::id scheduler_thread_id;
    std::vector<Plot_Core*> plots;
    std::atomic_bool running{false};
    std::atomic_bool cpu_joined{true};
};

Render_Scheduler::Render_Scheduler() : d(std::make_unique<Render_Scheduler_Private>()) {}
Global::Global() {}

namespace {
Global*& global_instance_storage() {
    static Global* value = nullptr;
    return value;
}
} // namespace

Global* Global::instance() {
    Global*& value = global_instance_storage();
    if (!value)
        value = new Global();
    return value;
}

void Global::destroy() {
    Global*& value = global_instance_storage();
    delete value;
    value = nullptr;
}

Render_Scheduler::~Render_Scheduler() {
    shutdown();
}

void Render_Scheduler::start() {
    memory_resource();
    if (d->running.exchange(true, std::memory_order_acq_rel))
        return;
    d->io_context.restart();
    d->work_guard = std::make_unique<asio::executor_work_guard<asio::io_context::executor_type>>(d->io_context.get_executor());
    unsigned cpu_count = std::max(1u, std::thread::hardware_concurrency() > 1 ? std::thread::hardware_concurrency() - 1 : 1u);
    d->cpu_pool = std::make_unique<asio::thread_pool>(cpu_count);
    d->cpu_joined.store(false, std::memory_order_release);
    d->scheduler_thread = std::thread([this]() {
        d->scheduler_thread_id = std::this_thread::get_id();
        for (Plot_Core* plot : d->plots)
            initialize_plot_on_scheduler(plot);
        d->io_context.run();
        d->running.store(false, std::memory_order_release);
    });
}

void Render_Scheduler::shutdown() {
    if (!d->running.load(std::memory_order_acquire)) {
        join_cpu_pool();
        if (d->scheduler_thread.joinable() && d->scheduler_thread.get_id() != std::this_thread::get_id())
            d->scheduler_thread.join();
        d->cpu_pool.reset();
        d->work_guard.reset();
        return;
    }
    if (is_scheduler_thread()) {
        shutdown_on_scheduler();
        d->work_guard.reset();
        join_cpu_pool();
        return;
    }
    post_and_wait([this]() {
        shutdown_on_scheduler();
    });
    join_cpu_pool();
    post([this]() {
        d->work_guard.reset();
    });
    if (d->scheduler_thread.joinable())
        d->scheduler_thread.join();
    d->cpu_pool.reset();
}

void Render_Scheduler::post(Task task) {
    asio::post(d->io_context, std::move(task));
}

void Render_Scheduler::post_at(std::uint64_t steady_time_ns, Task task) {
    if (!task)
        return;
    auto timer = std::make_shared<asio::steady_timer>(d->io_context);
    timer->expires_at(std::chrono::steady_clock::time_point(std::chrono::nanoseconds(steady_time_ns)));
    timer->async_wait([timer, task = std::move(task)](const asio::error_code& error) mutable {
        if (!error)
            task();
    });
}

void Render_Scheduler::post_cpu(Task task) {
    if (d->cpu_pool)
        asio::post(*d->cpu_pool, std::move(task));
    else
        post(std::move(task));
}

void Render_Scheduler::register_plot(Plot_Core* plot) {
    if (d->running.load(std::memory_order_acquire) && !is_scheduler_thread()) {
        post_and_wait([this, plot]() {
            register_plot_on_scheduler(plot);
        });
        return;
    }
    register_plot_on_scheduler(plot);
}

void Render_Scheduler::remove_plot(Plot_Core* plot) {
    if (d->running.load(std::memory_order_acquire) && !is_scheduler_thread()) {
        post_and_wait([this, plot]() {
            remove_plot_on_scheduler(plot);
        });
        return;
    }
    remove_plot_on_scheduler(plot);
}

bool Render_Scheduler::is_scheduler_thread() const {
    return std::this_thread::get_id() == d->scheduler_thread_id;
}

bool Render_Scheduler::is_running() const {
    return d->running.load(std::memory_order_acquire);
}

asio::io_context& Render_Scheduler::io_context() {
    return d->io_context;
}

void Render_Scheduler::join_cpu_pool() {
    if (d->cpu_pool && !d->cpu_joined.exchange(true, std::memory_order_acq_rel))
        d->cpu_pool->join();
}

void Render_Scheduler::post_and_wait(Task task) {
    if (is_scheduler_thread() || !d->running.load(std::memory_order_acquire)) {
        task();
        return;
    }
    std::mutex mutex;
    std::condition_variable cv;
    bool done = false;
    post([task = std::move(task), &mutex, &cv, &done]() mutable {
        task();
        {
            std::lock_guard<std::mutex> lock(mutex);
            done = true;
        }
        cv.notify_one();
    });
    std::unique_lock<std::mutex> lock(mutex);
    cv.wait(lock, [&done]() {
        return done;
    });
}

void Render_Scheduler::register_plot_on_scheduler(Plot_Core* plot) {
    if (std::find(d->plots.begin(), d->plots.end(), plot) == d->plots.end())
        d->plots.push_back(plot);
    if (d->running.load(std::memory_order_acquire))
        initialize_plot_on_scheduler(plot);
}

void Render_Scheduler::remove_plot_on_scheduler(Plot_Core* plot) {
    d->plots.erase(std::remove(d->plots.begin(), d->plots.end(), plot), d->plots.end());
    plot->pause_render();
}

void Render_Scheduler::shutdown_on_scheduler() {
    for (Plot_Core* plot : d->plots)
        plot->pause_render();
    d->plots.clear();
}

void Render_Scheduler::initialize_plot_on_scheduler(Plot_Core* plot) {
    plot->set_viewport_size(plot->viewport_size());
    if (plot->is_rendering())
        plot->start_render();
}

Render_Scheduler& Global::render_scheduler() {
    return render_scheduler_value;
}

void Global::start_render_scheduler() {
    render_scheduler_value.start();
}

} // namespace renderive