#include <web_server/src/Graph_WebSocket.hpp>
#include <web_server/src/media/Gallery_Video_Stream.hpp>
#include <mcp/core/Control_Service.hpp>
#include <mcp/core/runtime/Gallery_Plots.hpp>
#include <benchmark/benchmark.h>
#include <magic_enum/magic_enum.hpp>
#include <render_common.hpp>
#include <algorithm>
#include <array>
#include <atomic>
#include <charconv>
#include <chrono>
#include <cstddef>
#include <cstdint>
#include <iomanip>
#include <iostream>
#include <memory>
#include <ranges>
#include <stdexcept>
#include <string>
#include <string_view>
#include <vector>
#include <unordered_map>

namespace aethera::web::event_latency_benchmarks {
namespace {

using Steady_Clock = std::chrono::steady_clock;
using System_Clock = std::chrono::system_clock;
inline constexpr auto event_types = magic_enum::enum_values<Event_Type>();
inline constexpr std::size_t event_count = event_types.size();

struct Configuration {
    std::uint32_t width{720};
    std::uint32_t height{420};
    std::uint32_t samples{3};
    std::string plots{"all"};
};

struct Summary {
    double p50{};
    double p95{};
    double p99{};
    double maximum{};
};

struct Distribution {
    std::vector<double> samples{};

    void add(double value) { samples.push_back(std::max(0.0, value)); }

    [[nodiscard]] Summary summarize() const {
        if (samples.empty()) return {};
        auto ordered = samples;
        std::ranges::sort(ordered);
        const auto percentile = [&](double probability) {
            const auto rank = probability *
                static_cast<double>(ordered.size() - 1U);
            const auto lower = static_cast<std::size_t>(rank);
            const auto upper = std::min(lower + 1U, ordered.size() - 1U);
            const auto fraction = rank - static_cast<double>(lower);
            return ordered[lower] +
                (ordered[upper] - ordered[lower]) * fraction;
        };
        return {percentile(0.50), percentile(0.95), percentile(0.99),
                ordered.back()};
    }
};

struct Event_Result {
    Distribution websocket_receive_ms{};
    Distribution scene_queue_ms{};
    Distribution scene_dispatch_ms{};
    Distribution scene_total_ms{};
    Distribution receive_to_pixel_ms{};
    Distribution media_pipeline_ms{};
    Distribution end_to_end_ms{};
};

struct Plot_Result {
    std::string id{};
    Plot_Dimension dimension{Plot_Dimension::two_d};
    std::array<Event_Result, event_count> events{};
};

struct Media_Probe {
    std::atomic_uint64_t frame_count{};
    std::atomic_int64_t source_time_unix_ns{};
    std::atomic_int64_t published_time_unix_ns{};
};

struct Active_Plot {
    std::string id{};
    std::shared_ptr<Graph_WebSocket> input{};
    std::shared_ptr<media::Gallery_Video_Stream> video{};
    media::Gallery_Video_Stream::Stream_Id video_subscription{};
    std::shared_ptr<Media_Probe> probe{};
};

struct Event_Timing {
    std::int64_t submitted_unix_ns{};
    std::int64_t accepted_unix_ns{};
    std::uint64_t encoded_count_before{};
    std::uint64_t statistic_count_before{};
};

Configuration configuration{};
std::vector<Plot_Result> published_results{};

[[nodiscard]] std::int64_t system_time_ns() noexcept {
    return std::chrono::duration_cast<std::chrono::nanoseconds>(
        System_Clock::now().time_since_epoch()).count();
}

[[nodiscard]] double milliseconds(std::int64_t begin,
                                  std::int64_t end) noexcept {
    return end >= begin
        ? static_cast<double>(end - begin) / 1'000'000.0 : 0.0;
}

[[nodiscard]] const nlohmann::json* find_event_statistic(
    const nlohmann::json& diagnostics, Event_Type type,
    std::string_view statistic) {
    const auto input = diagnostics.find("input_statistics");
    if (input == diagnostics.end() || !input->is_object()) return nullptr;
    const auto event = input->find(
        std::string{magic_enum::enum_name(type)});
    if (event == input->end() || !event->is_object()) return nullptr;
    const auto value = event->find(std::string{statistic});
    return value == event->end() || !value->is_object()
        ? nullptr : &*value;
}

[[nodiscard]] std::uint64_t event_statistic_count(
    const mcp::Control_Service& control, std::string_view id,
    Event_Type type) {
    const auto diagnostics = control.plot_diagnostics(id);
    const auto* statistic = find_event_statistic(
        diagnostics, type, "total_ms");
    return statistic ? statistic->value("count", 0ULL) : 0ULL;
}

[[nodiscard]] double event_statistic_latest(
    const nlohmann::json& diagnostics, Event_Type type,
    std::string_view statistic) {
    const auto* value = find_event_statistic(diagnostics, type, statistic);
    if (!value)
        throw std::runtime_error(
            "missing Scene event statistic " +
            std::string{magic_enum::enum_name(type)} + "/" +
            std::string{statistic});
    return value->at("latest").get<double>();
}

[[nodiscard]] std::string input_message(Event_Type type,
                                        std::size_t plot_index,
                                        std::uint64_t sample) {
    const auto x = 120.0 + static_cast<double>(
        (sample * 13U + plot_index * 7U) % 400U);
    const auto y = 80.0 + static_cast<double>(
        (sample * 11U + plot_index * 5U) % 240U);
    return nlohmann::json{
        {"kind", "input"},
        {"event", {
            {"type", magic_enum::enum_name(type)},
            {"time_milliseconds", static_cast<double>(
                std::chrono::duration_cast<std::chrono::nanoseconds>(
                    Steady_Clock::now().time_since_epoch()).count()) /
                    1'000'000.0},
            {"position", {{"x", x}, {"y", y}}},
            {"global_position", {{"x", x}, {"y", y}}},
            {"button", "left"},
            {"buttons", type == Event_Type::pointer_release ? 0 : 1},
            {"modifiers", static_cast<int>(Keyboard_Modifier::control)},
            {"pixel_delta_x", 3.0},
            {"pixel_delta_y", (sample & 1U) == 0U ? 120.0 : -120.0},
            {"angle_delta_x", 3.0},
            {"angle_delta_y", (sample & 1U) == 0U ? 120.0 : -120.0},
            {"key", "space"},
            {"native_key", 32},
            {"auto_repeat", (sample & 1U) != 0U}
        }}
    }.dump();
}

void configure_plot(mcp::Control_Service& control, std::string_view id) {
    if (!control.write_plot_property(id,
            "frame-analysis", "pacing_mode", "manual").value(
                "success", false))
        throw std::runtime_error("failed to select manual Plot policy");
}

[[nodiscard]] bool manual_pixel_policy_applied(
    const mcp::Control_Service& control, std::string_view id) {
    const auto diagnostics = control.plot_diagnostics(id);
    const auto policy = diagnostics.find("frame_policy");
    if (policy == diagnostics.end() || !policy->is_object()) return false;
    const auto state = policy->find("configuration");
    if (state == policy->end() || !state->is_object()) return false;
    const auto mode = state->find("mode");
    const auto pixels = state->find("pixel_delivery_enabled");
    return mode != state->end() && mode->is_string() &&
        *mode == "manual" && pixels != state->end() &&
        pixels->is_boolean() && pixels->get<bool>();
}

void cooperatively_wait(std::string_view operation,
                        const std::function<bool()>& completed,
                        std::chrono::seconds timeout) {
    std::string failure;
    Task_Graph graph{"web.event-latency.await"};
    graph.add("await", [&] {
        const auto deadline = Steady_Clock::now() + timeout;
        Task_Graph::corun_until([&] {
            if (completed()) return true;
            if (Steady_Clock::now() < deadline) return false;
            failure = std::string{operation} + " timed out";
            return true;
        });
    });
    aethera::detail::run_taskflow(graph);
    if (!failure.empty()) throw std::runtime_error(failure);
}

[[nodiscard]] std::vector<Plot_Definition> selected_definitions() {
    const auto definitions = gallery_plot_definitions();
    if (configuration.plots == "all")
        return {definitions.begin(), definitions.end()};
    std::vector<Plot_Definition> selected;
    std::size_t begin{};
    while (begin <= configuration.plots.size()) {
        const auto end = configuration.plots.find(',', begin);
        const auto token = std::string_view{configuration.plots}.substr(
            begin, end == std::string::npos
                ? std::string::npos : end - begin);
        for (const auto& definition : definitions) {
            const bool matches_group =
                (token == "2d" &&
                 definition.dimension == Plot_Dimension::two_d) ||
                (token == "3d" &&
                 definition.dimension == Plot_Dimension::three_d);
            if (!matches_group && token != definition.id) continue;
            if (std::ranges::find(selected, definition.id,
                                  &Plot_Definition::id) == selected.end())
                selected.push_back(definition);
        }
        if (end == std::string::npos) break;
        begin = end + 1U;
    }
    if (selected.empty())
        throw std::invalid_argument(
            "--aethera_plots did not select a Gallery Plot");
    return selected;
}

void print_results() {
    std::cout << std::fixed << std::setprecision(3)
              << "\nWeb event -> H.264 backend timeline (milliseconds)\n"
              << "WS = JSON decode + Plot submit; Pixel = WS return to Plot "
                 "pixel entering Gallery; Media = Gallery queue/atlas/FFmpeg/"
                 "publish; E2E excludes network, WebCodecs and Canvas.\n\n"
              << std::left << std::setw(24) << "Plot"
              << std::setw(5) << "Dim"
              << std::setw(18) << "Event"
              << std::right << std::setw(10) << "WS P95"
              << std::setw(11) << "Queue P95"
              << std::setw(11) << "Disp P95"
              << std::setw(11) << "Pixel P95"
              << std::setw(11) << "Media P95"
              << std::setw(11) << "E2E P50"
              << std::setw(11) << "E2E P95"
              << std::setw(11) << "E2E P99" << '\n';
    for (const auto& plot : published_results) {
        for (std::size_t index = 0; index < event_count; ++index) {
            const auto& result = plot.events[index];
            const auto websocket = result.websocket_receive_ms.summarize();
            const auto queue = result.scene_queue_ms.summarize();
            const auto dispatch = result.scene_dispatch_ms.summarize();
            const auto pixel = result.receive_to_pixel_ms.summarize();
            const auto media = result.media_pipeline_ms.summarize();
            const auto end_to_end = result.end_to_end_ms.summarize();
            std::cout << std::left << std::setw(24) << plot.id
                      << std::setw(5) << plot_dimension_name(plot.dimension)
                      << std::setw(18) << magic_enum::enum_name(
                          event_types[index])
                      << std::right << std::setw(10) << websocket.p95
                      << std::setw(11) << queue.p95
                      << std::setw(11) << dispatch.p95
                      << std::setw(11) << pixel.p95
                      << std::setw(11) << media.p95
                      << std::setw(11) << end_to_end.p50
                      << std::setw(11) << end_to_end.p95
                      << std::setw(11) << end_to_end.p99 << '\n';
        }
    }
}

void run_event_latency(benchmark::State& state) {
    published_results.clear();
    std::vector<Active_Plot> active;

    try {
        const auto definitions = selected_definitions();
        auto control_slot = std::make_shared<std::atomic<
            std::shared_ptr<mcp::Control_Service>>>();
        auto videos = std::make_shared<std::unordered_map<
            std::string, std::shared_ptr<media::Gallery_Video_Stream>>>();
        active.reserve(definitions.size());
        published_results.reserve(definitions.size());
        for (const auto& definition : definitions) {
            const auto id = std::string{definition.id};
            auto probe = std::make_shared<Media_Probe>();
            auto video = media::Gallery_Video_Stream::create({
                {id, [control_slot, id](Frame_Publication_Feedback feedback) {
                    if (const auto control = control_slot->load())
                        control->submit_publication_feedback(
                            id, std::move(feedback));
                }}});
            videos->emplace(id, video);
            const auto video_subscription = video->subscribe(
                "web-event-latency-" + std::string{definition.id},
                [probe](media::Gallery_Stream_Frame frame) {
                    if (!frame.video) return false;
                    probe->source_time_unix_ns.store(
                        frame.video->source_time_unix.count(),
                        std::memory_order_relaxed);
                    probe->published_time_unix_ns.store(
                        system_time_ns(), std::memory_order_relaxed);
                    probe->frame_count.fetch_add(
                        1, std::memory_order_release);
                    return true;
                },
                [] { return true; },
                [] { return nlohmann::json::object(); });
            active.push_back({id, {},
                              std::move(video), video_subscription,
                              std::move(probe)});
            published_results.push_back({
                std::string{definition.id}, definition.dimension, {}});
        }
        auto control = mcp::Control_Service::create(mcp::Gallery_Output{
            .width = configuration.width,
            .height = configuration.height,
            .publish = [videos](
                std::string_view id,
                std::shared_ptr<const Gallery_Frame> frame) {
                const auto found = videos->find(std::string{id});
                return found == videos->end()
                    ? Gallery_Frame_Publication::ignored
                    : found->second->accept_frame(id, std::move(frame));
            }});
        control_slot->store(control);
        for (auto& item : active) {
            configure_plot(*control, item.id);
            item.input = std::make_shared<Graph_WebSocket>(control, item.id);
            item.input->start();
            item.input->receive(nlohmann::json{
                {"kind", "stream"},
                {"viewport", {{"width", configuration.width},
                               {"height", configuration.height}}}
            }.dump());
        }

        cooperatively_wait("manual Plot policies", [&] {
            return std::ranges::all_of(active, [&control](const Active_Plot& item) {
                return manual_pixel_policy_applied(*control, item.id);
            });
        }, std::chrono::seconds{10});

        std::vector<std::uint64_t> warmup_counts;
        warmup_counts.reserve(active.size());
        for (auto& item : active) {
            warmup_counts.push_back(item.probe->frame_count.load(
                std::memory_order_acquire));
            item.input->receive(R"({"kind":"manual_render"})");
        }
        cooperatively_wait("media warm-up", [&] {
            for (std::size_t index = 0; index < active.size(); ++index)
                if (active[index].probe->frame_count.load(
                        std::memory_order_acquire) <= warmup_counts[index])
                    return false;
            return true;
        }, std::chrono::seconds{60});
        for (auto& item : active)
            control->reset_plot_diagnostics(item.id);

        std::vector<Event_Timing> timings(active.size());
        std::string failure;
        Task_Graph coordinator{"web.event-latency.measure"};
        coordinator.add("measure", [&] {
            for (std::uint64_t sample = 0;
                 sample < configuration.samples; ++sample) {
                for (std::size_t event_index = 0;
                     event_index < event_count; ++event_index) {
                    const auto type = event_types[event_index];
                    for (std::size_t plot_index = 0;
                         plot_index < active.size(); ++plot_index) {
                        auto& item = active[plot_index];
                        auto& timing = timings[plot_index];
                        timing.encoded_count_before =
                            item.probe->frame_count.load(
                                std::memory_order_acquire);
                        timing.statistic_count_before = event_statistic_count(
                            *control, item.id, type);
                        const auto message = input_message(
                            type, plot_index, sample);
                        timing.submitted_unix_ns = system_time_ns();
                        item.input->receive(message);
                        timing.accepted_unix_ns = system_time_ns();
                        item.input->receive(
                            R"({"kind":"manual_render"})");
                    }

                    const auto deadline = Steady_Clock::now() +
                        std::chrono::seconds{60};
                    Task_Graph::corun_until([&] {
                        bool complete{true};
                        for (std::size_t index = 0;
                             index < active.size(); ++index) {
                            const auto encoded = active[index].probe->
                                frame_count.load(std::memory_order_acquire) >
                                timings[index].encoded_count_before;
                            const auto observed = event_statistic_count(
                                *control, active[index].id, type) >
                                timings[index].statistic_count_before;
                            complete = complete && encoded && observed;
                        }
                        if (complete) return true;
                        if (Steady_Clock::now() < deadline) return false;
                        failure = "timed out waiting for " +
                            std::string{magic_enum::enum_name(type)} +
                            " H.264 frames";
                        return true;
                    });
                    if (!failure.empty()) return;

                    for (std::size_t plot_index = 0;
                         plot_index < active.size(); ++plot_index) {
                        const auto& timing = timings[plot_index];
                        const auto source = active[plot_index].probe->
                            source_time_unix_ns.load(std::memory_order_acquire);
                        const auto published = active[plot_index].probe->
                            published_time_unix_ns.load(
                                std::memory_order_acquire);
                        const auto diagnostics =
                            control->plot_diagnostics(
                                active[plot_index].id);
                        auto& result = published_results[plot_index].
                            events[event_index];
                        result.websocket_receive_ms.add(milliseconds(
                            timing.submitted_unix_ns,
                            timing.accepted_unix_ns));
                        result.receive_to_pixel_ms.add(milliseconds(
                            timing.accepted_unix_ns, source));
                        result.media_pipeline_ms.add(milliseconds(
                            source, published));
                        result.end_to_end_ms.add(milliseconds(
                            timing.submitted_unix_ns, published));
                        result.scene_queue_ms.add(event_statistic_latest(
                            diagnostics, type, "queue_wait_ms"));
                        result.scene_dispatch_ms.add(event_statistic_latest(
                            diagnostics, type, "dispatch_ms"));
                        result.scene_total_ms.add(event_statistic_latest(
                            diagnostics, type, "total_ms"));
                    }
                }
            }
        });

        for ([[maybe_unused]] auto iteration : state)
            aethera::detail::run_taskflow(coordinator);
        if (!failure.empty()) state.SkipWithError(failure);

        Distribution aggregate;
        for (const auto& plot : published_results)
            for (const auto& event : plot.events)
                aggregate.samples.insert(
                    aggregate.samples.end(), event.end_to_end_ms.samples.begin(),
                    event.end_to_end_ms.samples.end());
        const auto summary = aggregate.summarize();
        state.counters["all/e2e_p50_ms"] = summary.p50;
        state.counters["all/e2e_p95_ms"] = summary.p95;
        state.counters["all/e2e_p99_ms"] = summary.p99;
        state.SetItemsProcessed(static_cast<std::int64_t>(
            configuration.samples * active.size() * event_count));
        print_results();
    }
    catch (const std::exception& error) {
        state.SkipWithError(error.what());
    }

    for (auto& item : active) {
        item.input->close();
        item.video->unsubscribe(item.video_subscription);
        item.video->shutdown();
    }
}

BENCHMARK(run_event_latency)->Iterations(1)->UseRealTime();

[[nodiscard]] bool parse_unsigned(std::string_view value,
                                  std::uint32_t& output) {
    const auto* begin = value.data();
    const auto* end = begin + value.size();
    const auto parsed = std::from_chars(begin, end, output);
    return parsed.ec == std::errc{} && parsed.ptr == end;
}

[[nodiscard]] bool configure(int& argc, char** argv) {
    int retained{1};
    for (int index = 1; index < argc; ++index) {
        const std::string_view argument{argv[index]};
        constexpr std::string_view samples_prefix{
            "--aethera_event_samples="};
        constexpr std::string_view size_prefix{"--aethera_size="};
        constexpr std::string_view plots_prefix{"--aethera_plots="};
        if (argument.starts_with(samples_prefix)) {
            if (!parse_unsigned(argument.substr(samples_prefix.size()),
                                configuration.samples) ||
                configuration.samples == 0 || configuration.samples > 600) {
                std::cerr << "--aethera_event_samples must be in [1, 600]\n";
                return false;
            }
            continue;
        }
        if (argument.starts_with(size_prefix)) {
            const auto value = argument.substr(size_prefix.size());
            const auto separator = value.find('x');
            if (separator == std::string_view::npos ||
                !parse_unsigned(value.substr(0, separator),
                                configuration.width) ||
                !parse_unsigned(value.substr(separator + 1U),
                                configuration.height) ||
                configuration.width == 0 || configuration.height == 0) {
                std::cerr << "--aethera_size requires WIDTHxHEIGHT\n";
                return false;
            }
            continue;
        }
        if (argument.starts_with(plots_prefix)) {
            configuration.plots = argument.substr(plots_prefix.size());
            if (configuration.plots.empty()) {
                std::cerr << "--aethera_plots requires all, 2d, 3d or IDs\n";
                return false;
            }
            continue;
        }
        argv[retained++] = argv[index];
    }
    argc = retained;
    return true;
}

} // namespace
} // namespace aethera::web::event_latency_benchmarks

int main(int argc, char** argv) {
    if (!aethera::web::event_latency_benchmarks::configure(argc, argv))
        return 2;
    aethera::initialize_runtime({});
    benchmark::Initialize(&argc, argv);
    if (benchmark::ReportUnrecognizedArguments(argc, argv)) return 2;
    benchmark::AddCustomContext(
        "aethera_plots",
        aethera::web::event_latency_benchmarks::configuration.plots);
    benchmark::AddCustomContext(
        "aethera_event_samples",
        std::to_string(
            aethera::web::event_latency_benchmarks::configuration.samples));
    benchmark::AddCustomContext(
        "aethera_size",
        std::to_string(
            aethera::web::event_latency_benchmarks::configuration.width) +
        "x" + std::to_string(
            aethera::web::event_latency_benchmarks::configuration.height));
    benchmark::RunSpecifiedBenchmarks();
    benchmark::Shutdown();
    return 0;
}