#include "Plot_Core.h"

#include "../plottable/Performance_Overlay.h"
#include "../render/Blend2D_Cache.h"
#include "../renderable/Renderable.h"

#include <renderive/frame_control/strategy/low_latency/Low_Latency_Strategy.hpp>
#include <renderive/frame_control/strategy/manual/Manual_Refresh_Strategy.hpp>
#include <renderive/frame_control/strategy/flow/Flow_Refresh_Strategy.hpp>
#include <renderive/base/observer/Observer.hpp>
#include <renderive/scene/Scene2D_Context.hpp>

#include <algorithm>
#include <chrono>
#include <cmath>
#include <limits>
#include <mutex>
#include <stdexcept>
#include <type_traits>
#include <utility>
#include <variant>

namespace renderive {
namespace {

struct Frame_Observer_Data {
    mutable std::mutex mutex;
    std::uint32_t last_event{};
    std::uint64_t observation_count{};
    std::uint64_t produced_frame_count{};
    std::uint64_t consumed_frame_count{};
    std::uint64_t dropped_frame_count{};
    std::uint64_t failed_operation_count{};
    std::uint64_t pending_frame_count{};
    std::uint64_t latest_sequence{};
    std::uint64_t paint_duration_ns{};
    std::uint64_t render_duration_ns{};
    std::uint64_t paint_lease_wait_ns{};
    std::uint64_t paint_state_wait_ns{};
    std::uint64_t publish_state_wait_ns{};
    std::uint64_t ready_wait_ns{};
    std::uint64_t frame_age_at_render_ns{};
    std::uint64_t render_lease_wait_ns{};
    std::uint64_t render_state_wait_ns{};
    std::uint64_t render_finish_state_wait_ns{};
    std::uint64_t queue_wait_ns{};
};

struct Frame_Observer {
    static constexpr bool enabled = true;
    std::shared_ptr<Frame_Observer_Data> data;

    template <class Observation>
    void observe(const Observation& observation) noexcept {
        std::lock_guard lock(data->mutex);
        data->last_event = static_cast<std::uint32_t>(observation.event);
        ++data->observation_count;
        if (observation.statistics.sequence != 0)
            data->latest_sequence = observation.statistics.sequence;
        if constexpr (requires { observation.statistics.timing.paint_duration_ns; }) {
            if (observation.statistics.sequence != 0) {
                const auto& timing = observation.statistics.timing;
                data->paint_lease_wait_ns = timing.paint_lease_wait_ns;
                data->paint_state_wait_ns = timing.paint_state_wait_ns;
                data->publish_state_wait_ns = timing.publish_state_wait_ns;
                data->ready_wait_ns = timing.ready_wait_ns;
                data->frame_age_at_render_ns = timing.frame_age_at_render_ns;
                data->render_lease_wait_ns = timing.render_lease_wait_ns;
                data->render_state_wait_ns = timing.render_state_wait_ns;
                data->render_finish_state_wait_ns = timing.render_finish_state_wait_ns;
            }
        } else {
            if (observation.statistics.sequence != 0) {
                data->paint_duration_ns = observation.statistics.paint_duration_ns;
                data->render_duration_ns = observation.statistics.render_duration_ns;
            }
        }
        if constexpr (requires { observation.statistics.queue_wait_ns; }) {
            if (observation.statistics.sequence != 0)
                data->queue_wait_ns = observation.statistics.queue_wait_ns;
        }

        if constexpr (requires { observation.state.prepared_frame_count; }) {
            data->produced_frame_count = observation.state.prepared_frame_count;
            data->consumed_frame_count = observation.state.render_count;
            data->dropped_frame_count = observation.state.replaced_prepared_frame_count +
                                        observation.state.discarded_prepared_frame_count;
            data->failed_operation_count = observation.state.failed_refresh_count;
            data->pending_frame_count = observation.state.pending_frame ? 1U : 0U;
        } else if constexpr (requires { observation.state.limit_state; }) {
            data->produced_frame_count = observation.statistics.counters.published_frame_count;
            data->consumed_frame_count = observation.state.completed_lifecycle_count;
            data->dropped_frame_count =
                observation.statistics.counters.abandoned_frame_count +
                observation.statistics.counters.manually_discarded_frame_count;
            data->failed_operation_count = observation.statistics.counters.swap_failure_count;
            const auto retired_frame_count =
                observation.state.completed_lifecycle_count +
                observation.statistics.counters.abandoned_frame_count +
                observation.statistics.counters.manually_discarded_frame_count;
            data->pending_frame_count =
                observation.statistics.counters.published_frame_count > retired_frame_count ? 1U : 0U;
        } else {
            data->produced_frame_count = observation.state.enqueued_frame_count;
            data->consumed_frame_count = observation.state.rendered_frame_count;
            data->failed_operation_count = observation.state.empty_acquire_count;
            data->pending_frame_count = observation.state.pending_frame_count;
        }
    }
};

using Frame_Observer_State = ::Observer_State<Frame_Observer>;
using Manual_Frame_Control =
    ::Manual_Refresh_Strategy<::Scene2D_Frame_Data, std::mutex, Frame_Observer_State>;
using Low_Latency_Frame_Control =
    ::Low_Latency_Strategy<::Scene2D_Frame_Data, std::mutex, Frame_Observer_State>;
using Playback_Frame_Control =
    ::Flow_Refresh_Strategy<::Scene2D_Frame_Data, std::mutex, Frame_Observer_State>;
using Manual_Plot_Scene =
    ::Scene2D_Context<Manual_Frame_Control, detail::Blend2D_Color_Cache>;
using Low_Latency_Plot_Scene =
    ::Scene2D_Context<Low_Latency_Frame_Control, detail::Blend2D_Color_Cache>;
using Playback_Plot_Scene =
    ::Scene2D_Context<Playback_Frame_Control, detail::Blend2D_Color_Cache>;
using Plot_Scene = std::variant<std::unique_ptr<Manual_Plot_Scene>,
                                std::unique_ptr<Low_Latency_Plot_Scene>,
                                std::unique_ptr<Playback_Plot_Scene>>;

template <class Function>
decltype(auto) with_scene(Plot_Scene& scene, Function&& function) {
    return std::visit(
        [&function](auto& concrete) -> decltype(auto) {
            return std::forward<Function>(function)(*concrete);
        },
        scene);
}

template <class Function>
decltype(auto) with_scene(const Plot_Scene& scene, Function&& function) {
    return std::visit(
        [&function](const auto& concrete) -> decltype(auto) {
            return std::forward<Function>(function)(*concrete);
        },
        scene);
}

std::string observer_event_name(Frame_Control_Mode mode, std::uint32_t event) {
    if (mode == Frame_Control_Mode::Manual) {
        static constexpr const char* names[]{"prepared", "prepared_replaced", "refresh_succeeded",
                                              "refresh_failed", "rendered", "real_time_data_updated",
                                              "manually_discarded"};
        return event < std::size(names) ? names[event] : "none";
    }
    if (mode == Frame_Control_Mode::Playback) {
        static constexpr const char* names[]{"enqueued", "dequeued", "queue_empty", "rendered",
                                              "real_time_data_updated"};
        return event < std::size(names) ? names[event] : "none";
    }
    static constexpr const char* names[]{"published", "abandoned", "manually_discarded",
                                          "swap_failed", "rendered", "lifecycle_completed",
                                          "real_time_data_updated"};
    return event < std::size(names) ? names[event] : "none";
}

std::string limit_state_name(bool available, std::uint32_t state) {
    if (!available)
        return "not_applicable";
    static constexpr const char* names[]{"frequency_limited", "paint_limited", "render_limited",
                                          "consumer_limited", "unlimited"};
    return state < std::size(names) ? names[state] : "unknown";
}

Rect full_rect(Size size) {
    return {0, 0, std::max(0, size.width), std::max(0, size.height)};
}

class Renderable_Group final : public Renderable {
public:
    using Renderable::Renderable;

private:
    void paint(detail::Painter&) override {}
};

} // namespace

struct Plot_Core::Impl {
    explicit Impl(Frame_Control_Mode selected_mode)
        : mode(selected_mode), observer(std::make_shared<Frame_Observer_Data>()),
          scene(make_scene(selected_mode, observer)) {}

    static Plot_Scene make_scene(Frame_Control_Mode mode,
                                 const std::shared_ptr<Frame_Observer_Data>& observer) {
        auto& resource = *memory_resource(Memory_Domain::Plot_Frame);
        Frame_Observer_State observer_state(Frame_Observer{observer});
        switch (mode) {
        case Frame_Control_Mode::Manual:
            return std::make_unique<Manual_Plot_Scene>(resource, std::move(observer_state));
        case Frame_Control_Mode::Playback:
            return std::make_unique<Playback_Plot_Scene>(resource, std::move(observer_state));
        case Frame_Control_Mode::Low_Latency:
            return std::make_unique<Low_Latency_Plot_Scene>(resource, std::move(observer_state));
        }
        throw std::invalid_argument("unknown frame control mode");
    }

    Frame_Control_Mode mode;
    std::shared_ptr<Frame_Observer_Data> observer;
    Plot_Scene scene;
    mutable std::mutex mutex;
    std::mutex initialization_mutex;
    std::weak_ptr<Presentation_Sink> presentation_sink;
    std::shared_ptr<Performance_Overlay> performance;
    Color background = Color::black();
    Size viewport;
    std::atomic_bool active{};
    std::atomic_bool dirty{true};
};

Plot_Core::Plot_Core()
    : Plot_Core(Frame_Control_Mode::Low_Latency) {}

Plot_Core::Plot_Core(Frame_Control_Mode mode)
    : impl_(std::make_unique<Impl>(mode)) {}

Plot_Core::~Plot_Core() = default;

void Plot_Core::init() {
    std::lock_guard initialization_lock(impl_->initialization_mutex);
    if (root_renderable())
        return;
    auto root = renderive::make_shared<Renderable_Group>(Memory_Domain::Renderable, *this, true);
    root->set_object_name("root");
    with_scene(impl_->scene, [&root](auto& scene) { scene.attach_renderable(root); });
    notify_model_dirty();
}

std::shared_ptr<Renderable> Plot_Core::root_renderable() const {
    const auto topology = with_scene(impl_->scene, [](const auto& scene) {
        return scene.topology_snapshot();
    });
    for (const auto& relationship : topology.display) {
        if (relationship.parent)
            continue;
        auto base = std::const_pointer_cast<::Renderable_Base>(relationship.child);
        if (auto renderable = std::dynamic_pointer_cast<Renderable>(base))
            return renderable;
    }
    return {};
}

std::shared_ptr<Renderable> Plot_Core::create_renderable_node(
    const std::shared_ptr<Renderable>& parent, std::string object_name) {
    auto renderable = renderive::make_shared<Renderable_Group>(
        Memory_Domain::Renderable, *this, true);
    renderable->set_object_name(std::move(object_name));
    attach_renderable(renderable, parent);
    return renderable;
}

void Plot_Core::attach_renderable(const std::shared_ptr<Renderable>& renderable,
                                  const std::shared_ptr<Renderable>& parent) {
    if (!renderable)
        throw std::invalid_argument("renderable is null");
    Renderable* parent_pointer = parent.get();
    if (!parent_pointer) {
        auto root = root_renderable();
        if (!root)
            throw std::logic_error("Plot_Core::init must be called before adding renderables");
        parent_pointer = root.get();
    }
    with_scene(impl_->scene, [&renderable](auto& scene) { scene.attach_renderable(renderable); });
    try {
        with_scene(impl_->scene, [&renderable, parent_pointer](auto& scene) {
            scene.set_display_parent(*renderable, parent_pointer);
            scene.set_dependency_parent(*renderable, parent_pointer);
        });
    } catch (...) {
        with_scene(impl_->scene, [&renderable](auto& scene) { scene.detach_renderable(*renderable); });
        throw;
    }
    notify_model_dirty();
}

void Plot_Core::remove_renderable(const std::shared_ptr<Renderable>& renderable) {
    if (!renderable || renderable == root_renderable())
        return;
    with_scene(impl_->scene, [&renderable](auto& scene) { scene.detach_renderable(*renderable); });
    notify_model_dirty();
}

void Plot_Core::set_background_color(Color color) {
    {
        std::lock_guard lock(impl_->mutex);
        if (impl_->background == color)
            return;
        impl_->background = color;
    }
    notify_model_dirty();
}

Color Plot_Core::background_color() const noexcept {
    std::lock_guard lock(impl_->mutex);
    return impl_->background;
}

void Plot_Core::set_viewport_size(Size size) {
    size.width = std::max(0, size.width);
    size.height = std::max(0, size.height);
    {
        std::lock_guard lock(impl_->mutex);
        if (impl_->viewport == size)
            return;
        impl_->viewport = size;
    }
    const auto topology = with_scene(impl_->scene, [](const auto& scene) {
        return scene.topology_snapshot();
    });
    for (const auto& renderable : topology.renderables)
        const_cast<::Renderable_Base&>(*renderable).invalidate_cache();
    notify_model_dirty();
}

Size Plot_Core::viewport_size() const noexcept {
    std::lock_guard lock(impl_->mutex);
    return impl_->viewport;
}

void Plot_Core::dispatch_event(const Event& event) {
    const auto topology = with_scene(impl_->scene, [](const auto& scene) {
        return scene.topology_snapshot();
    });
    for (auto iterator = topology.display.rbegin(); iterator != topology.display.rend(); ++iterator) {
        auto base = std::const_pointer_cast<::Renderable_Base>(iterator->child);
        if (auto renderable = std::dynamic_pointer_cast<Renderable>(base)) {
            if (renderable->is_visible()) {
                if (auto handler = std::dynamic_pointer_cast<Event_Handler>(base))
                    handler->handle_event(event);
            }
            if (event.is_accepted())
                break;
        }
    }
}

void Plot_Core::notify_model_dirty() noexcept {
    impl_->dirty.store(true, std::memory_order_release);
}

bool Plot_Core::prepare_frame() {
    return with_scene(impl_->scene, [](auto& scene) {
        auto paint_frame = scene.frame_control.acquire_painter();
        if (!paint_frame)
            return false;
        auto& scene_state = static_cast<typename std::remove_reference_t<decltype(scene)>::Scene_State_Strategy&>(scene);
        scene_state.template set<&::Scene2D_State::revision>(scene_state.state_revision() + 1);
        scene_state.publish();
        return true;
    });
}

bool Plot_Core::refresh_manual_frame() {
    if (impl_->mode != Frame_Control_Mode::Manual)
        return false;
    return std::get<std::unique_ptr<Manual_Plot_Scene>>(impl_->scene)->frame_control.refresh();
}

bool Plot_Core::render_prepared_frame() {
    return with_scene(impl_->scene, [](auto& scene) {
        auto render_frame = scene.frame_control.acquire_renderer();
        if (!render_frame)
            return false;
        scene.render();
        scene.wait_for_render();
        return true;
    });
}

bool Plot_Core::discard_pending_frame() {
    return with_scene(impl_->scene, [](auto& scene) {
        using Scene = std::remove_reference_t<decltype(scene)>;
        if constexpr (std::is_same_v<Scene, Playback_Plot_Scene>)
            return false;
        else
            return scene.frame_control.discard_pending_frame();
    });
}

bool Plot_Core::render_frame(bool force) {
    if (!view_active() && !force)
        return false;
    const Size viewport = viewport_size();
    if (viewport.empty())
        return false;
    if (!impl_->dirty.exchange(false, std::memory_order_acq_rel) && !force)
        return false;

    const auto started = std::chrono::steady_clock::now();
    try {
        if (!prepare_frame())
            return false;
        if (impl_->mode == Frame_Control_Mode::Manual && !refresh_manual_frame())
            return false;
        if (!render_prepared_frame())
            return false;
    } catch (...) {
        notify_model_dirty();
        throw;
    }
    const auto finished = std::chrono::steady_clock::now();
    const double duration_ms = std::chrono::duration<double, std::milli>(finished - started).count();

    std::shared_ptr<Presentation_Sink> sink;
    std::shared_ptr<Performance_Overlay> overlay;
    {
        std::lock_guard lock(impl_->mutex);
        sink = impl_->presentation_sink.lock();
        overlay = impl_->performance;
    }
    if (overlay)
        overlay->record_frame(duration_ms, diagnostics(),
                              with_scene(impl_->scene, [](const auto& scene) {
                                  return scene.renderable_count();
                              }),
                              viewport);
    if (sink)
        sink->request_present(full_rect(viewport));
    return true;
}

void Plot_Core::with_frame(const std::function<void(Image_View)>& consumer) {
    if (!consumer)
        return;
    with_scene(impl_->scene, [&consumer](auto& scene) {
        scene.with_final_color_cache([&consumer](const detail::Blend2D_Color_Cache& cache) {
            consumer(cache.view());
        });
    });
}

void Plot_Core::activate_view() noexcept {
    impl_->active.store(true, std::memory_order_release);
    notify_model_dirty();
}

void Plot_Core::deactivate_view() noexcept {
    impl_->active.store(false, std::memory_order_release);
}

bool Plot_Core::view_active() const noexcept {
    return impl_->active.load(std::memory_order_acquire);
}

void Plot_Core::set_max_render_fps(double fps) {
    if (!std::isfinite(fps) || fps <= 0.0)
        throw std::invalid_argument("maximum render FPS must be finite and positive");
    if (impl_->mode != Frame_Control_Mode::Low_Latency)
        throw std::logic_error("maximum render FPS is only available in low-latency mode");
    std::get<std::unique_ptr<Low_Latency_Plot_Scene>>(impl_->scene)->frame_control.set_frequency_hz(fps);
    notify_model_dirty();
}

void Plot_Core::clear_max_render_fps() {
    if (impl_->mode != Frame_Control_Mode::Low_Latency)
        throw std::logic_error("maximum render FPS is only available in low-latency mode");
    std::get<std::unique_ptr<Low_Latency_Plot_Scene>>(impl_->scene)->frame_control.clear_frequency_limit();
    notify_model_dirty();
}

void Plot_Core::set_consumer_feedback(Frame_Consumer_Feedback feedback) {
    if (impl_->mode != Frame_Control_Mode::Low_Latency)
        throw std::logic_error("consumer feedback is only available in low-latency mode");
    std::get<std::unique_ptr<Low_Latency_Plot_Scene>>(impl_->scene)->frame_control.set_consumer_feedback(feedback);
}

void Plot_Core::clear_consumer_feedback() {
    if (impl_->mode != Frame_Control_Mode::Low_Latency)
        throw std::logic_error("consumer feedback is only available in low-latency mode");
    std::get<std::unique_ptr<Low_Latency_Plot_Scene>>(impl_->scene)->frame_control.clear_consumer_feedback();
}

double Plot_Core::max_render_fps() const noexcept {
    if (impl_->mode != Frame_Control_Mode::Low_Latency)
        return std::numeric_limits<double>::quiet_NaN();
    return std::get<std::unique_ptr<Low_Latency_Plot_Scene>>(impl_->scene)->frame_control.state().frequency_hz;
}

Low_Latency_Diagnostics Plot_Core::diagnostics() const {
    if (impl_->mode == Frame_Control_Mode::Low_Latency) {
        const auto state = std::get<std::unique_ptr<Low_Latency_Plot_Scene>>(impl_->scene)->frame_control.state();
        return {view_active(),
                {state.frequency_hz, state.completed_lifecycle_count,
                 state.next_refresh_interval_ns}};
    }
    const auto observed = frame_observer_snapshot();
    return {view_active(),
            {std::numeric_limits<double>::quiet_NaN(), observed.consumed_frame_count, 0}};
}

Refresh_Control_Snapshot Plot_Core::refresh_feedback_snapshot() const {
    return diagnostics().refresh;
}

Frame_Control_Mode Plot_Core::frame_control_mode() const noexcept {
    return impl_->mode;
}

Frame_Observer_Snapshot Plot_Core::frame_observer_snapshot() const {
    Frame_Observer_Snapshot snapshot;
    {
        std::lock_guard lock(impl_->observer->mutex);
        snapshot.mode = impl_->mode;
        snapshot.last_event = impl_->observer->observation_count == 0 ?
                                  std::string("none") :
                                  observer_event_name(impl_->mode, impl_->observer->last_event);
        snapshot.limit_state = "not_applicable";
        snapshot.observation_count = impl_->observer->observation_count;
        snapshot.produced_frame_count = impl_->observer->produced_frame_count;
        snapshot.consumed_frame_count = impl_->observer->consumed_frame_count;
        snapshot.dropped_frame_count = impl_->observer->dropped_frame_count;
        snapshot.failed_operation_count = impl_->observer->failed_operation_count;
        snapshot.pending_frame_count = impl_->observer->pending_frame_count;
        snapshot.latest_sequence = impl_->observer->latest_sequence;
        snapshot.paint_duration_ns = impl_->observer->paint_duration_ns;
        snapshot.render_duration_ns = impl_->observer->render_duration_ns;
        snapshot.paint_lease_wait_ns = impl_->observer->paint_lease_wait_ns;
        snapshot.paint_state_wait_ns = impl_->observer->paint_state_wait_ns;
        snapshot.publish_state_wait_ns = impl_->observer->publish_state_wait_ns;
        snapshot.ready_wait_ns = impl_->observer->ready_wait_ns;
        snapshot.frame_age_at_render_ns = impl_->observer->frame_age_at_render_ns;
        snapshot.render_lease_wait_ns = impl_->observer->render_lease_wait_ns;
        snapshot.render_state_wait_ns = impl_->observer->render_state_wait_ns;
        snapshot.render_finish_state_wait_ns = impl_->observer->render_finish_state_wait_ns;
        snapshot.queue_wait_ns = impl_->observer->queue_wait_ns;
    }
    if (impl_->mode == Frame_Control_Mode::Low_Latency) {
        const auto state = std::get<std::unique_ptr<Low_Latency_Plot_Scene>>(impl_->scene)->frame_control.state();
        snapshot.limit_state = limit_state_name(true, static_cast<std::uint32_t>(state.limit_state));
        snapshot.frequency_hz = state.frequency_hz;
        snapshot.paint_duration_ns = state.paint_duration_ns;
        snapshot.render_duration_ns = state.render_duration_ns;
        snapshot.target_interval_ns = state.target_interval_ns;
        snapshot.frequency_limit_enabled = state.frequency_limit_enabled;
        snapshot.consumer_feedback_enabled = state.consumer_feedback_enabled;
        snapshot.bottleneck_duration_ns = state.bottleneck_duration_ns;
        snapshot.consumer_sample_interval_ns = state.consumer_sample_interval_ns;
        snapshot.consumer_smoothed_interval_ns = state.consumer_smoothed_interval_ns;
        snapshot.consumer_variation_ns = state.consumer_variation_ns;
        snapshot.consumer_safety_interval_ns = state.consumer_safety_interval_ns;
        snapshot.consumer_interval_ns = state.consumer_interval_ns;
        snapshot.next_refresh_interval_ns = state.next_refresh_interval_ns;
        snapshot.end_to_end_ns = state.end_to_end_ns;
    }
    return snapshot;
}

void Plot_Core::set_presentation_sink(std::weak_ptr<Presentation_Sink> sink) {
    std::lock_guard lock(impl_->mutex);
    impl_->presentation_sink = std::move(sink);
}

std::shared_ptr<Performance_Overlay> Plot_Core::performance_overlay() const {
    std::lock_guard lock(impl_->mutex);
    return impl_->performance;
}

void Plot_Core::set_performance_overlay(std::shared_ptr<Performance_Overlay> overlay) {
    {
        std::lock_guard lock(impl_->mutex);
        impl_->performance = std::move(overlay);
    }
    notify_model_dirty();
}

::Scene_Base& Plot_Core::kernel_scene() const noexcept {
    return with_scene(impl_->scene, [](auto& scene) -> ::Scene_Base& { return scene; });
}

void Plot_Core::set_renderable_cache(Renderable& renderable, Renderable_Cache_Mode mode) {
    with_scene(impl_->scene, [&renderable, mode](auto& scene) {
        scene.set_renderable_configuration(
            renderable, {.cache_enabled = mode == Renderable_Cache_Mode::Local_Pixel});
    });
    notify_model_dirty();
}

} // namespace renderive