#include "Plot_Core.h"
#include "Plot_Frame_Pipeline.h"
#include "../architecture/Renderable.h"
#include "../architecture/Render_Time.h"
#include "../architecture/Update_Completion.h"
#include "../base/global.h"
#include "../render/Canvas.h"
#include <algorithm>
#include <memory>
#include <utility>

namespace renderive {

class Plot_Core_Private;

struct Frame_Render_Job {
    std::shared_ptr<Plot_Core_Private> owner;
    std::uint64_t generation{};
    Plot_Frame_Lease frame_lease;
    Plot_Frame* buffer{};
    std::shared_ptr<Frame_Lifecycle_Record> lifecycle;
    Plot_Render_Snapshot render_snapshot;
    Renderable_Tree_Snapshot renderables;
    std::shared_ptr<Frame_Worker_Stat> worker_stat;
};

class Plot_Core_Private final : public std::enable_shared_from_this<Plot_Core_Private> {
public:
    explicit Plot_Core_Private(std::unique_ptr<Refresh_Strategy> s)
        : strategy(std::move(s)) {}

    ~Plot_Core_Private() = default;

    bool render_enabled() const {
        return rendering.load(std::memory_order_acquire);
    }

    bool destroying() const {
        return destroying_flag.load(std::memory_order_acquire);
    }

    Refresh_Runtime_Snapshot refresh_runtime_snapshot(std::uint64_t now_ns = steady_now_ns()) const {
        return {
            render_enabled(),
            destroying(),
            pipeline.has_new_state(),
            pipeline.job_state != Plot_Render_Job_State::Idle,
            pipeline.middle_has_new_frame(),
            pipeline.paint_active(),
            now_ns
        };
    }

    bool submit_render_once() {
        if (!rendering.load(std::memory_order_acquire) || destroying_flag.load(std::memory_order_acquire))
            return false;
        if (!root)
            return false;

        Size size{
            context->render_width.load(std::memory_order_acquire),
            context->render_height.load(std::memory_order_acquire)
        };
        if (size.empty())
            return false;

        auto frame_lease = pipeline.try_begin_render();
        if (!frame_lease)
            return false;

        Plot_Frame* buffer = frame_lease.frame;
        cancel_update_states(buffer->presented_update_states);
        std::uint64_t render_role_enter_ns = buffer->metadata ? buffer->metadata->render_role_enter_ns : 0;
        buffer->metadata = std::make_shared<Frame_Lifecycle_Record>();
        auto frame_record = buffer->metadata;
        Frame_Lifecycle_Record& frame = *frame_record;
        std::uint64_t submit_time = steady_now_ns();
        frame.frame_id = ++next_frame_id;
        frame.input_version = pipeline.edit_state_version.load(std::memory_order_acquire);
        frame.render_request_ns = submit_time;
        frame.render_role_enter_ns = render_role_enter_ns ? render_role_enter_ns : submit_time;
        frame.core_data_time_ns = latest_input_time_ns.load(std::memory_order_acquire);
        frame.raw_input_count = pending_input_count.exchange(0, std::memory_order_acq_rel);
        frame.pixel_width = static_cast<std::uint64_t>(size.width);
        frame.pixel_height = static_cast<std::uint64_t>(size.height);
        frame.viewport_version = context->viewport_version.load(std::memory_order_acquire);
        frame.wait_21_ns = submit_time > frame.render_role_enter_ns ? submit_time - frame.render_role_enter_ns : 0;
        frame.render_begin_ns = submit_time;

        std::uint64_t previous_render_state_version = pipeline.render_state_version;
        pipeline.job_state = Plot_Render_Job_State::Running;
        pipeline.render_state_version = frame.input_version;

        auto job = std::make_shared<Frame_Render_Job>();
        job->owner = shared_from_this();
        job->generation = generation.load(std::memory_order_acquire);
        job->frame_lease = std::move(frame_lease);
        job->buffer = buffer;
        job->lifecycle = frame_record;
        job->worker_stat = std::make_shared<Frame_Worker_Stat>();
        job->render_snapshot.context = context;
        job->render_snapshot.size = size;
        job->render_snapshot.background_color = background_color();
        job->render_snapshot.previous_frame_time_ns = previous_frame_time_ns;
        job->render_snapshot.frame_time_ns = submit_time;
        job->render_snapshot.frame_metadata = frame_record.get();
        job->render_snapshot.frame_update_states = &buffer->presented_update_states;
        previous_frame_time_ns = submit_time;

        job->renderables = renderable_tree_snapshot();
        auto& scheduler = Global::instance()->render_scheduler();
        bool submitted = scheduler.try_submit_render_task(Render_Executor_Task{
            plot_execution_id,
            frame.frame_id,
            Render_Task_Kind::Frame,
            Task{},
            [job](tf::Taskflow& taskflow, tf::Task entry, tf::Task exit) {
                auto prepare = taskflow.emplace([job]() {
                    job->owner->prepare_frame_job(job);
                });
                auto draw = taskflow.emplace([job]() {
                    job->owner->draw_frame_job(job);
                });
                auto finalize = taskflow.emplace([job]() {
                    job->owner->finalize_frame_job(job);
                });
                entry.precede(prepare);
                prepare.precede(draw);
                draw.precede(finalize);
                finalize.precede(exit);
            },
            Task([job]() {
                Global::instance()->render_scheduler().post([job]() {
                    job->owner->complete_frame_job(job);
                });
            }),
            job->worker_stat,
            3,
            true
        });
        if (submitted)
            return true;

        frame.outcome = Frame_Outcome::Worker_Budget_Dropped;
        frame.worker.rejected_task_count++;
        job->frame_lease.reset();
        pipeline.render_state_version = previous_render_state_version;
        pipeline.job_state = Plot_Render_Job_State::Idle;
        collect_frame(frame_record);
        trigger_refresh_strategy(Refresh_Trigger::Render_Finished, &frame);
        return false;
    }

    static bool frame_job_valid(const std::shared_ptr<Frame_Render_Job>& job) {
        return job && job->buffer && job->lifecycle && job->render_snapshot.context && !job->render_snapshot.destroying();
    }

    void prepare_frame_job(const std::shared_ptr<Frame_Render_Job>& job) {
        if (!frame_job_valid(job))
            return;
        Frame_Lifecycle_Record& frame = *job->lifecycle;
        frame.prepare_begin_ns = steady_now_ns();
        for (const auto& node : job->renderables) {
            const auto& renderable = node.renderable;
            if (!renderable || !renderable->should_prepare())
                continue;
            renderable->prepare_data(job->render_snapshot);
            ++frame.prepared_output_count;
        }
        frame.prepare_end_ns = steady_now_ns();
        job->render_snapshot.context->first_prepare_data.store(false, std::memory_order_release);
    }

    void draw_frame_job(const std::shared_ptr<Frame_Render_Job>& job) {
        if (!frame_job_valid(job) || !job->lifecycle->prepare_end_ns)
            return;
        Frame_Lifecycle_Record& frame = *job->lifecycle;
        job->buffer->image.resize(job->render_snapshot.size.width, job->render_snapshot.size.height);
        job->buffer->image.fill(job->render_snapshot.background_color);
        Canvas canvas(job->buffer->image);

        frame.draw_begin_ns = steady_now_ns();
        if (!job->renderables.empty() && job->renderables.front().renderable)
            job->renderables.front().renderable->render_tree_snapshot(canvas, job->render_snapshot, job->renderables, 0);
        frame.draw_end_ns = steady_now_ns();
        frame.render_end_ns = steady_now_ns();
    }

    void finalize_frame_job(const std::shared_ptr<Frame_Render_Job>& job) {
        if (!frame_job_valid(job) || !job->lifecycle->render_end_ns)
            return;
        Frame_Lifecycle_Record& frame = *job->lifecycle;
        job->buffer->version.store(frame.frame_id, std::memory_order_release);
        for (const auto& update_state : job->buffer->presented_update_states) {
            if (update_state)
                update_state->complete_until(Update_Stage::Rendered);
        }
    }

    void execute_frame_job(const std::shared_ptr<Frame_Render_Job>& job) {
        if (!job || !job->buffer || !job->lifecycle || !job->render_snapshot.context || job->render_snapshot.destroying())
            return;
        prepare_frame_job(job);
        draw_frame_job(job);
        finalize_frame_job(job);
    }

    void complete_frame_job(const std::shared_ptr<Frame_Render_Job>& job) {
        if (!job || !job->buffer || !job->lifecycle)
            return;
        Frame_Lifecycle_Record& frame = *job->lifecycle;
        if (job->worker_stat)
            frame.worker = *job->worker_stat;
        bool valid_generation = job->generation == generation.load(std::memory_order_acquire);
        bool valid_context = !destroying_flag.load(std::memory_order_acquire) && context && !context->destroying.load(std::memory_order_acquire);
        job->frame_lease.reset();
        if (!valid_generation || !valid_context || !frame.render_end_ns) {
            frame.outcome = Frame_Outcome::Cancelled;
            pipeline.job_state = Plot_Render_Job_State::Idle;
            supersede_update_states(job->buffer->presented_update_states);
            trigger_refresh_strategy(Refresh_Trigger::Render_Finished, &frame);
            return;
        }

        auto publish_result = pipeline.try_publish_rendered(steady_now_ns());
        if (publish_result.superseded_middle) {
            trigger_refresh_strategy(Refresh_Trigger::Frame_Returned, publish_result.superseded_middle.get());
            collect_frame(publish_result.superseded_middle);
        }
        if (publish_result.dropped_frame) {
            collect_frame(publish_result.dropped_frame);
        }

        pipeline.job_state = Plot_Render_Job_State::Idle;
        if (publish_result.request_present)
            request_present(publish_result.dirty_rect);
        trigger_refresh_strategy(Refresh_Trigger::Render_Finished, &frame);
    }

    void post_scheduler_action(Task task) {
        if (!task)
            return;
        auto guard = context;
        auto action = [guard, task = std::move(task)]() mutable {
            if (!guard || guard->destroying.load(std::memory_order_acquire))
                return;
            task();
        };
        auto& scheduler = Global::instance()->render_scheduler();
        if (scheduler.is_scheduler_thread()) {
            action();
            return;
        }
        scheduler.post(std::move(action));
    }

    void trigger_refresh_strategy(Refresh_Trigger trigger, Frame_Lifecycle_Record* frame = nullptr) {
        if (!strategy)
            return;
        std::uint64_t now_ns = steady_now_ns();
        Refresh_Runtime_Snapshot runtime = refresh_runtime_snapshot(now_ns);
        Refresh_Action action = strategy->on_trigger(trigger, runtime, frame);
        apply_refresh_action(action, runtime);
    }

    void notify_input_changed(std::uint64_t now_ns) {
        if (!strategy)
            return;
        Refresh_Runtime_Snapshot runtime = refresh_runtime_snapshot(now_ns);
        Refresh_Action action = strategy->on_trigger(Refresh_Trigger::Input_Dirty, runtime, nullptr);
        apply_refresh_action(action, runtime);
    }

    void notify_started() {
        if (!strategy_started) {
            strategy_started = true;
            trigger_refresh_strategy(Refresh_Trigger::Started);
        }
        else {
            trigger_refresh_strategy(Refresh_Trigger::Input_Dirty);
        }
    }

    void notify_paused() {
        invalidate_refresh_deadline();
        if (!strategy_started)
            return;
        strategy_started = false;
        trigger_refresh_strategy(Refresh_Trigger::Paused);
    }

    void notify_frame_returned(std::shared_ptr<Frame_Lifecycle_Record> frame) {
        if (!frame)
            return;
        trigger_refresh_strategy(Refresh_Trigger::Frame_Returned, frame.get());
        collect_frame(frame);
    }

    void apply_refresh_action(Refresh_Action action, const Refresh_Runtime_Snapshot& runtime) {
        if (action.request_paint) {
            auto update_result = pipeline.request_paint_update(steady_now_ns());
            if (update_result.request_present)
                request_present(update_result.dirty_rect);
        }
        if (action.try_render) {
            invalidate_refresh_deadline();
            submit_render_once();
            return;
        }
        if (action.render_deadline_ns) {
            schedule_refresh_deadline(action.render_deadline_ns);
            return;
        }
        if (!runtime.enabled || runtime.destroying || !runtime.dirty)
            invalidate_refresh_deadline();
    }

    void schedule_refresh_deadline(std::uint64_t deadline_ns) {
        if (deadline_ns <= steady_now_ns()) {
            invalidate_refresh_deadline();
            trigger_refresh_strategy(Refresh_Trigger::Timer_Fired);
            return;
        }
        if (scheduled_refresh_deadline_ns && scheduled_refresh_deadline_ns <= deadline_ns)
            return;
        scheduled_refresh_deadline_ns = deadline_ns;
        std::uint64_t generation = ++refresh_deadline_generation;
        auto guard = context;
        auto self = shared_from_this();
        Global::instance()->render_scheduler().post_at(deadline_ns, [guard, self, generation, deadline_ns]() {
            if (!guard || guard->destroying.load(std::memory_order_acquire))
                return;
            if (self->refresh_deadline_generation != generation || self->scheduled_refresh_deadline_ns != deadline_ns)
                return;
            self->scheduled_refresh_deadline_ns = 0;
            self->trigger_refresh_strategy(Refresh_Trigger::Timer_Fired);
        });
    }

    void invalidate_refresh_deadline() {
        scheduled_refresh_deadline_ns = 0;
        ++refresh_deadline_generation;
    }

    void dispatch_event(const Event& event) {
        if (!root)
            return;
        switch (event.type) {
        case Event_Type::Pointer_Move: {
            const auto& pointer_event = static_cast<const Pointer_Event&>(event);
            update_hover_targets(pointer_event);
            dispatch_to_renderables(root->hit_renderables(pointer_event.position), event);
            break;
        }
        case Event_Type::Pointer_Press:
        case Event_Type::Pointer_Release: {
            const auto& pointer_event = static_cast<const Pointer_Event&>(event);
            dispatch_to_renderables(root->hit_renderables(pointer_event.position), event);
            break;
        }
        case Event_Type::Wheel: {
            const auto& wheel_event = static_cast<const Wheel_Event&>(event);
            dispatch_to_renderables(root->hit_renderables(wheel_event.position), event);
            break;
        }
        case Event_Type::Key_Press:
        case Event_Type::Key_Release:
            dispatch_to_renderables(reverse_renderable_snapshot(), event);
            break;
        case Event_Type::Leave:
            clear_hover_targets();
            dispatch_to_renderables(renderable_snapshot(), event);
            break;
        default:
            dispatch_to_renderables(renderable_snapshot(), event);
            break;
        }
    }

    std::vector<std::shared_ptr<Renderable>> renderable_snapshot() const {
        std::vector<std::shared_ptr<Renderable>> result;
        append_renderable_snapshot(result, root);
        return result;
    }

    Renderable_Tree_Snapshot renderable_tree_snapshot() const {
        Renderable_Tree_Snapshot result;
        if (root)
            root->append_tree_snapshot(result);
        return result;
    }

    std::vector<std::shared_ptr<Renderable>> reverse_renderable_snapshot() const {
        std::vector<std::shared_ptr<Renderable>> result = renderable_snapshot();
        std::reverse(result.begin(), result.end());
        return result;
    }

    void append_renderable_snapshot(std::vector<std::shared_ptr<Renderable>>& result, const std::shared_ptr<Renderable>& renderable) const {
        if (!renderable)
            return;
        result.push_back(renderable);
        for (const auto& child : renderable->children_snapshot())
            append_renderable_snapshot(result, child);
    }

    void dispatch_to_renderables(const std::vector<std::shared_ptr<Renderable>>& renderables, const Event& event) {
        for (const auto& renderable : renderables) {
            if (!renderable || event.is_accepted())
                break;
            renderable->plot_event(event);
        }
    }

    void update_hover_targets(const Pointer_Event& event) {
        Point pos{static_cast<int>(event.position.x), static_cast<int>(event.position.y)};
        std::vector<std::shared_ptr<Renderable>> hits = root->hit_renderables(event.position);
        for (const auto& old_target : hover_targets) {
            auto target = old_target.lock();
            if (!target)
                continue;
            if (std::find(hits.begin(), hits.end(), target) == hits.end())
                target->set_hover_state(pos, false);
        }
        hover_targets.clear();
        hover_targets.reserve(hits.size());
        for (const auto& target : hits) {
            if (!target)
                continue;
            target->set_hover_state(pos, true);
            hover_targets.push_back(target);
        }
    }

    void clear_hover_targets() {
        Point pos{};
        for (const auto& old_target : hover_targets) {
            if (auto target = old_target.lock())
                target->set_hover_state(pos, false);
        }
        hover_targets.clear();
    }

    void request_present(Rect dirty_rect) {
        if (auto target = sink.lock())
            target->request_present(dirty_rect);
    }

    void collect_frame(std::shared_ptr<const Frame_Lifecycle_Record> frame) {
        if (!frame)
            return;
        auto observer = std::atomic_load_explicit(&context->frame_lifecycle_observer, std::memory_order_acquire);
        if (observer && observer->enabled())
            observer->collect_frame(frame);
    }

    Color background_color() const {
        Pixel p = context->background_rgba.load(std::memory_order_acquire);
        return premultiplied_to_color(p);
    }

    std::unique_ptr<Refresh_Strategy> strategy;
    std::shared_ptr<Plot_Render_Context> context = std::make_shared<Plot_Render_Context>();
    std::shared_ptr<Renderable> root;
    std::vector<std::weak_ptr<Renderable>> hover_targets;
    Plot_Frame_Pipeline pipeline;
    std::weak_ptr<Presentation_Sink> sink;
    std::atomic_bool rendering{false};
    std::atomic_bool destroying_flag{false};
    std::atomic<std::uint64_t> latest_input_time_ns{0};
    std::atomic<std::uint64_t> pending_input_count{0};
    bool strategy_started{};
    std::uint64_t refresh_deadline_generation{};
    std::uint64_t scheduled_refresh_deadline_ns{};
    std::uint64_t next_frame_id{};
    std::uint64_t previous_frame_time_ns{};
    Plot_Execution_Id plot_execution_id{};
    std::atomic<std::uint64_t> generation{1};
};

namespace {
Plot_Execution_Id next_plot_execution_id() {
    static std::atomic<Plot_Execution_Id> next_id{1};
    return next_id.fetch_add(1, std::memory_order_acq_rel);
}
} // namespace

Plot_Core::Plot_Core(std::unique_ptr<Refresh_Strategy> strategy)
    : d(std::make_shared<Plot_Core_Private>(std::move(strategy))) {
    d->plot_execution_id = next_plot_execution_id();
    d->context->owner.store(this, std::memory_order_release);
    Global::instance()->render_scheduler().register_plot(this);
}

Plot_Core::~Plot_Core() {
    d->destroying_flag.store(true, std::memory_order_release);
    d->context->destroying.store(true, std::memory_order_release);
    d->generation.fetch_add(1, std::memory_order_acq_rel);
    Global::instance()->render_scheduler().remove_plot(this);
    d->pipeline.cancel_all_update_states();
    d->context->owner.store(nullptr, std::memory_order_release);
}

void Plot_Core::init() {
    if (!d->root) {
        d->root = std::make_shared<Renderable>();
        d->root->set_object_name("Root_Renderable");
        d->root->init_root(this);
    }
    d->context->first_resize.store(true, std::memory_order_release);
}

std::shared_ptr<Plot_Render_Context> Plot_Core::render_context() const {
    return d->context;
}

std::shared_ptr<Renderable> Plot_Core::root_renderable() const {
    return d->root;
}

std::shared_ptr<Renderable> Plot_Core::create_renderable_node(
    const std::shared_ptr<Renderable>& parent, std::string object_name) {
    if (!parent)
        return {};
    auto node = std::make_shared<Renderable>();
    node->set_object_name(std::move(object_name));
    node->init(parent);
    return node;
}

void Plot_Core::remove_renderable(const std::shared_ptr<Renderable>& renderable) {
    if (!renderable)
        return;
    if (renderable == d->root)
        d->root.reset();
    else
        renderable->detach_from_parent();
    mark_render_state_dirty();
}

void Plot_Core::set_background_color(Color color) {
    d->context->background_rgba.store(premultiply(color), std::memory_order_release);
    mark_render_state_dirty();
}

Color Plot_Core::background_color() const {
    return d->background_color();
}

void Plot_Core::set_viewport_size(Size size) {
    Size previous{
        d->context->render_width.load(std::memory_order_acquire),
        d->context->render_height.load(std::memory_order_acquire)
    };
    d->context->render_width.store(size.width, std::memory_order_release);
    d->context->render_height.store(size.height, std::memory_order_release);
    if (previous != size)
        d->context->viewport_version.fetch_add(1, std::memory_order_acq_rel);
    if (!size.empty())
        d->context->first_resize.store(false, std::memory_order_release);
}

Size Plot_Core::viewport_size() const {
    return {
        d->context->render_width.load(std::memory_order_acquire),
        d->context->render_height.load(std::memory_order_acquire)
    };
}

void Plot_Core::dispatch_event(const Event& event) {
    d->dispatch_event(event);
}

void Plot_Core::mark_render_state_dirty() {
    if (d->destroying_flag.load(std::memory_order_acquire))
        return;
    std::uint64_t now_ns = steady_now_ns();
    d->pipeline.mark_edit_state();
    d->latest_input_time_ns.store(now_ns, std::memory_order_release);
    d->pending_input_count.fetch_add(1, std::memory_order_relaxed);
    d->post_scheduler_action([data = d, now_ns]() {
        data->notify_input_changed(now_ns);
    });
}

void Plot_Core::request_manual_refresh() {
    if (d->destroying_flag.load(std::memory_order_acquire))
        return;
    d->post_scheduler_action([data = d]() {
        data->trigger_refresh_strategy(Refresh_Trigger::Manual_Request);
    });
}

void Plot_Core::start_render() {
    if (d->destroying_flag.load(std::memory_order_acquire))
        return;
    d->rendering.store(true, std::memory_order_release);
    d->post_scheduler_action([data = d]() {
        data->notify_started();
    });
}

void Plot_Core::pause_render() {
    d->rendering.store(false, std::memory_order_release);
    d->post_scheduler_action([data = d]() {
        data->notify_paused();
    });
}

bool Plot_Core::is_rendering() const {
    return d->rendering.load(std::memory_order_acquire);
}

void Plot_Core::set_presentation_sink(std::weak_ptr<Presentation_Sink> sink) {
    d->sink = std::move(sink);
}

Present_Frame Plot_Core::begin_present() {
    Present_Frame result;
    if (!d->root)
        return result;
    auto consumed = d->pipeline.try_begin_paint(steady_now_ns());
    if (consumed.returned_frame) {
        d->post_scheduler_action([data = d, frame = consumed.returned_frame]() mutable {
            data->notify_frame_returned(frame);
        });
    }
    result.is_new_frame = consumed.new_frame;
    result.update_was_pending = consumed.update_was_pending;
    if (consumed.lease)
        result.lease_ = std::make_unique<Plot_Frame_Lease>(std::move(consumed.lease));
    if (!result.new_frame() && result.request_was_pending()) {
        d->post_scheduler_action([data = d]() {
            data->trigger_refresh_strategy(Refresh_Trigger::Paint_Finished);
        });
    }
    return result;
}

void Plot_Core::end_present(Present_Frame&& frame, std::uint64_t paint_begin_ns, std::uint64_t paint_end_ns) {
    if (frame.lease_ && frame.lease_->frame) {
        auto* rendered_frame = frame.lease_->frame;
        if (rendered_frame->metadata) {
            rendered_frame->metadata->paint_begin_ns = paint_begin_ns;
            rendered_frame->metadata->paint_end_ns = paint_end_ns;
            rendered_frame->metadata->wait_23_ns = paint_begin_ns > rendered_frame->metadata->transition_12_ns ? paint_begin_ns - rendered_frame->metadata->transition_12_ns : 0;
        }
        complete_update_states(rendered_frame->presented_update_states, Update_Stage::Presented);
    }
    frame.lease_.reset();
    d->post_scheduler_action([data = d]() mutable {
        data->trigger_refresh_strategy(Refresh_Trigger::Paint_Finished);
    });
}

} // namespace renderive