#include <mcp/core/Control_Requests.hpp>
#include <mcp/core/Control_Service.hpp>
#include <mcp/core/Protocol_Type.hpp>
#include <mcp/core/runtime/Gallery_Plots.hpp>
#include <benchmark/benchmark.h>
#include <magic_enum/magic_enum.hpp>
#include <task_flow/src/Task_Runtime.hpp>
#include <algorithm>
#include <array>
#include <atomic>
#include <charconv>
#include <chrono>
#include <cstddef>
#include <cstdint>
#include <iomanip>
#include <iostream>
#include <memory>
#include <stdexcept>
#include <string>
#include <string_view>
#include <vector>

namespace aethera::mcp::event_latency_benchmarks {
namespace {

using Clock = std::chrono::steady_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{12};
    bool three_d_only{};
    bool manual_policy{};
};

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

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

    void add(double value) { samples.push_back(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 mcp_accept_ms{};
    Distribution accepted_to_frame_ms{};
    Distribution end_to_end_ms{};
};

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

struct Plot_Probe {
    std::atomic_uint64_t expected_correlation{};
    std::atomic_uint64_t completed_steady_ns{};
};

struct Active_Plot {
    std::string id{};
    std::shared_ptr<Plot_Probe> probe{};
};

struct Batch_Timing {
    std::array<std::uint64_t, event_count> submitted_ns{};
    std::array<std::uint64_t, event_count> accepted_ns{};
};

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

[[nodiscard]] std::uint64_t steady_time_ns() noexcept {
    return static_cast<std::uint64_t>(
        std::chrono::duration_cast<std::chrono::nanoseconds>(
            Clock::now().time_since_epoch()).count());
}

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

[[nodiscard]] nlohmann::json make_event(
    Event_Type type, std::size_t plot_index, std::uint64_t batch) {
    const nlohmann::json position{
        {"x", 120.0 + static_cast<double>(
            (batch * 13U + plot_index * 7U) % 400U)},
        {"y", 80.0 + static_cast<double>(
            (batch * 11U + plot_index * 5U) % 240U)}};
    const auto wheel_delta = (batch & 1U) == 0U ? 120.0 : -120.0;
    return {
        {"type", magic_enum::enum_name(type)},
        {"time_milliseconds",
         static_cast<double>(steady_time_ns()) / 1'000'000.0},
        {"position", position}, {"global_position", position},
        {"button", "left"},
        {"buttons", type == Event_Type::pointer_release ? 0U : 1U},
        {"modifiers", static_cast<std::uint8_t>(Keyboard_Modifier::control)},
        {"pixel_delta_x", 3.0}, {"pixel_delta_y", wheel_delta},
        {"angle_delta_x", 3.0}, {"angle_delta_y", wheel_delta},
        {"key", "space"}, {"native_key", 32U},
        {"auto_repeat", (batch & 1U) != 0U}
    };
}

[[nodiscard]] std::uint64_t correlation_id(
    std::uint64_t batch, std::size_t plot_index) noexcept {
    return (batch + 1U) * 1'000U +
        static_cast<std::uint64_t>(plot_index + 1U);
}

void add_distribution(Distribution& destination,
                      const Distribution& source) {
    destination.samples.insert(destination.samples.end(),
                               source.samples.begin(), source.samples.end());
}

class Event_Latency_Reporter final : public benchmark::ConsoleReporter {
public:
    bool ReportContext(const Context& context) override {
        return ConsoleReporter::ReportContext(context);
    }

    void ReportRuns(const std::vector<Run>& reports) override {
        auto& output = GetOutputStream();
        for (const auto& report : reports) {
            output << report.benchmark_name();
            if (report.skipped != benchmark::internal::NotSkipped)
                output << " ERROR: " << report.skip_message << '\n';
            else
                output << "  elapsed=" << report.real_accumulated_time
                       << " s\n";
        }
        if (published_results.empty()) return;
        const auto flags = output.flags();
        const auto precision = output.precision();
        output << std::fixed << std::setprecision(3)
               << "\nPure backend event timeline (ms)\n"
               << "MCP = emit -> aethera_plot_input accepted; "
                  "Frame = accepted -> correlated backend frame complete.\n\n";

        output << std::left << std::setw(24) << "Plot"
               << std::setw(6) << "Dim"
               << std::setw(18) << "Worst event"
               << std::right << std::setw(11) << "MCP P95"
               << std::setw(11) << "Frame P95"
               << std::setw(11) << "E2E P95" << '\n';
        for (const auto& plot : published_results) {
            std::size_t worst_index{};
            double worst_e2e{};
            for (std::size_t index = 0; index < event_count; ++index) {
                const auto p95 = plot.events[index].end_to_end_ms.summarize().p95;
                if (p95 <= worst_e2e) continue;
                worst_e2e = p95;
                worst_index = index;
            }
            const auto& event = plot.events[worst_index];
            output << std::left << std::setw(24) << plot.id
                   << std::setw(6) << web::plot_dimension_name(plot.dimension)
                   << std::setw(18)
                   << magic_enum::enum_name(event_types[worst_index])
                   << std::right
                   << std::setw(11) << event.mcp_accept_ms.summarize().p95
                   << std::setw(11)
                   << event.accepted_to_frame_ms.summarize().p95
                   << std::setw(11) << worst_e2e << '\n';
        }

        output << "\nAll plots by event type\n"
               << std::left << std::setw(18) << "Event"
               << std::right << std::setw(11) << "MCP P50"
               << std::setw(11) << "MCP P95"
               << std::setw(11) << "Frame P95"
               << std::setw(11) << "E2E P50"
               << std::setw(11) << "E2E P95"
               << std::setw(11) << "E2E P99" << '\n';
        for (std::size_t event_index = 0; event_index < event_count;
             ++event_index) {
            Event_Result aggregate{};
            for (const auto& plot : published_results) {
                const auto& event = plot.events[event_index];
                add_distribution(aggregate.mcp_accept_ms, event.mcp_accept_ms);
                add_distribution(
                    aggregate.accepted_to_frame_ms,
                    event.accepted_to_frame_ms);
                add_distribution(aggregate.end_to_end_ms, event.end_to_end_ms);
            }
            const auto mcp = aggregate.mcp_accept_ms.summarize();
            const auto frame = aggregate.accepted_to_frame_ms.summarize();
            const auto end_to_end = aggregate.end_to_end_ms.summarize();
            output << std::left << std::setw(18)
                   << magic_enum::enum_name(event_types[event_index])
                   << std::right << std::setw(11) << mcp.p50
                   << std::setw(11) << mcp.p95
                   << std::setw(11) << frame.p95
                   << std::setw(11) << end_to_end.p50
                   << std::setw(11) << end_to_end.p95
                   << std::setw(11) << end_to_end.p99 << '\n';
        }
        output.flags(flags);
        output.precision(precision);
    }
};

void run_event_latency(benchmark::State& state) {
    published_results.clear();
    const auto definitions = web::gallery_plot_definitions();
    auto active_owner = std::make_shared<std::vector<Active_Plot>>();
    auto& active = *active_owner;
    active.reserve(definitions.size());
    published_results.reserve(definitions.size());

    std::shared_ptr<Control_Service> service;
    const auto stop_service = [&] {
        if (!service) return;
        std::atomic_bool stopped{};
        service->stop([&] { stopped.store(true, std::memory_order_release); stopped.notify_one(); });
        /* This is the benchmark executable's terminal boundary, not a production worker wait. */
        stopped.wait(false, std::memory_order_acquire);
        service.reset();
    };

    try {
        for (const auto& definition : definitions) {
            if (configuration.three_d_only &&
                definition.dimension != web::Plot_Dimension::three_d)
                continue;
            auto probe = std::make_shared<Plot_Probe>();
            active.push_back({std::string{definition.id}, std::move(probe)});
            published_results.push_back({
                std::string{definition.id}, definition.dimension, {}});
        }
        service = Control_Service::create(Gallery_Output{
            .width = configuration.width,
            .height = configuration.height,
            .publish = [active_owner](
                std::string_view id,
                std::shared_ptr<const web::Gallery_Frame> frame) {
                auto& active = *active_owner;
                const auto found = std::ranges::find(active, id,
                                                     &Active_Plot::id);
                if (found == active.end() || !frame || !frame->pixels)
                    return web::Gallery_Frame_Publication::ignored;
                const auto expected = found->probe->expected_correlation.load(
                    std::memory_order_acquire);
                if (expected != 0 &&
                    frame->pixels->correlation_id == expected) {
                    std::uint64_t incomplete{};
                    static_cast<void>(found->probe->completed_steady_ns.
                        compare_exchange_strong(
                            incomplete, steady_time_ns(),
                            std::memory_order_release,
                            std::memory_order_relaxed));
                }
                return web::Gallery_Frame_Publication::completed;
            }});
        std::vector<Batch_Timing> batch_timing(active.size());
        std::string failure;
        Task_Graph coordinator{"mcp.event-latency"};
        const auto measurement = coordinator.add("measure.all-plots", [&] {
            for (std::uint64_t batch = 0; batch < configuration.samples;
                 ++batch) {
                const auto deadline = Clock::now() + std::chrono::seconds{30};
                for (std::size_t plot_index = 0;
                     plot_index < active.size(); ++plot_index) {
                    auto& item = active[plot_index];
                    auto& timing = batch_timing[plot_index];
                    const auto correlation = correlation_id(batch, plot_index);
                    item.probe->completed_steady_ns.store(
                        0, std::memory_order_relaxed);
                    item.probe->expected_correlation.store(
                        correlation, std::memory_order_release);
                    for (std::size_t event_index = 0;
                         event_index < event_count; ++event_index) {
                        auto request = Plot_Input_Request{
                            .plot = published_results[plot_index].id,
                            .event = make_event(
                                event_types[event_index], plot_index, batch)
                        };
                        timing.submitted_ns[event_index] = steady_time_ns();
                        const auto result = service->call_tool(
                            "aethera_plot_input",
                            encode_protocol_value(request));
                        timing.accepted_ns[event_index] = steady_time_ns();
                        if (result.result != Tool_Call_Result::ok) {
                            failure = "aethera_plot_input rejected " +
                                published_results[plot_index].id + "/" +
                                std::string{magic_enum::enum_name(
                                    event_types[event_index])} + ": " +
                                result.message;
                            return;
                        }
                    }
                    service->request_frame(item.id, web::Gallery_Frame_Request{
                        .issued_at = Clock::now(),
                        .time_milliseconds =
                            static_cast<double>(steady_time_ns()) / 1'000'000.0,
                        .correlation_id = correlation,
                        .width = configuration.width,
                        .height = configuration.height
                    });
                }

                /* Benchmark-only semantic wait: the calling Taskflow worker cooperatively executes available work while waiting for correlated publications. */
                aethera::detail::corun_taskflow_until([&] {
                    const bool complete = std::ranges::all_of(
                        active, [](const Active_Plot& item) {
                            return item.probe->completed_steady_ns.load(
                                std::memory_order_acquire) != 0;
                        });
                    if (complete) return true;
                    if (Clock::now() < deadline) return false;
                    failure = "timed out waiting for a correlated backend frame";
                    return true;
                });
                if (!failure.empty()) return;

                for (std::size_t plot_index = 0;
                     plot_index < active.size(); ++plot_index) {
                    const auto completed = active[plot_index].probe->
                        completed_steady_ns.load(std::memory_order_acquire);
                    for (std::size_t event_index = 0;
                         event_index < event_count; ++event_index) {
                        auto& result = published_results[plot_index].
                            events[event_index];
                        const auto& timing = batch_timing[plot_index];
                        result.mcp_accept_ms.add(elapsed_ms(
                            timing.submitted_ns[event_index],
                            timing.accepted_ns[event_index]));
                        result.accepted_to_frame_ms.add(elapsed_ms(
                            timing.accepted_ns[event_index], completed));
                        result.end_to_end_ms.add(elapsed_ms(
                            timing.submitted_ns[event_index], completed));
                    }
                    active[plot_index].probe->expected_correlation.store(
                        0, std::memory_order_release);
                }
            }
        });
        if (!measurement) throw std::logic_error("event latency coordinator task is invalid");

        for ([[maybe_unused]] auto iteration : state) {
            std::atomic_bool completed{};
            std::exception_ptr task_failure;
            if (run_taskflow(coordinator, [&](std::exception_ptr value) { task_failure = std::move(value); completed.store(true, std::memory_order_release); completed.notify_one(); }) != Run_Taskflow_Result::submitted) throw std::runtime_error("event latency Taskflow submission failed");
            /* Google Benchmark owns this non-worker thread and requires the iteration result before return. This test-only terminal wait never occupies a Taskflow, Scene, Frame Policy, timer, or publication worker. */
            completed.wait(false, std::memory_order_acquire);
            if (task_failure) std::rethrow_exception(task_failure);
        }
        if (!failure.empty()) state.SkipWithError(failure);

        Distribution all_end_to_end;
        for (const auto& plot : published_results) {
            Distribution plot_end_to_end;
            for (std::size_t event_index = 0;
                 event_index < event_count; ++event_index) {
                const auto event_name = std::string{
                    magic_enum::enum_name(event_types[event_index])};
                const auto& event = plot.events[event_index];
                add_distribution(plot_end_to_end, event.end_to_end_ms);
                add_distribution(all_end_to_end, event.end_to_end_ms);
                const auto prefix = plot.id + "/" + event_name + "/";
                state.counters[prefix + "e2e_p95_ms"] =
                    event.end_to_end_ms.summarize().p95;
            }
            state.counters[plot.id + "/e2e_p95_ms"] =
                plot_end_to_end.summarize().p95;
        }
        const auto end_to_end = all_end_to_end.summarize();
        state.counters["all/e2e_p50_ms"] = end_to_end.p50;
        state.counters["all/e2e_p95_ms"] = end_to_end.p95;
        state.counters["all/e2e_p99_ms"] = end_to_end.p99;
        state.SetItemsProcessed(static_cast<std::int64_t>(
            configuration.samples * active.size() * event_count));
        stop_service();
    }
    catch (const std::exception& error) {
        stop_service();
        state.SkipWithError(error.what());
    }
}

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="};
        constexpr std::string_view policy_prefix{"--aethera_policy="};
        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)) {
            const auto value = argument.substr(plots_prefix.size());
            if (value == "all") configuration.three_d_only = false;
            else if (value == "3d") configuration.three_d_only = true;
            else {
                std::cerr << "--aethera_plots requires all or 3d\n";
                return false;
            }
            continue;
        }
        if (argument.starts_with(policy_prefix)) {
            const auto value = argument.substr(policy_prefix.size());
            if (value == "current") configuration.manual_policy = false;
            else {
                std::cerr << "--aethera_policy requires current\n";
                return false;
            }
            continue;
        }
        argv[retained++] = argv[index];
    }
    argc = retained;
    return true;
}

} // namespace
} // namespace aethera::mcp::event_latency_benchmarks

int main(int argc, char** argv) {
    if (!aethera::mcp::event_latency_benchmarks::configure(argc, argv))
        return 2;
    if (aethera::initialize_task_runtime() != aethera::Initialize_Task_Runtime_Result::initialized) return 2;
    benchmark::Initialize(&argc, argv);
    if (benchmark::ReportUnrecognizedArguments(argc, argv)) return 2;
    benchmark::AddCustomContext(
        "aethera_plots",
        aethera::mcp::event_latency_benchmarks::configuration.three_d_only
            ? "3d" : "all");
    benchmark::AddCustomContext(
        "aethera_policy",
        aethera::mcp::event_latency_benchmarks::configuration.manual_policy
            ? "manual" : "current");
    benchmark::AddCustomContext(
        "aethera_event_samples",
        std::to_string(
            aethera::mcp::event_latency_benchmarks::configuration.samples));
    benchmark::AddCustomContext(
        "aethera_size",
        std::to_string(
            aethera::mcp::event_latency_benchmarks::configuration.width) +
        "x" + std::to_string(
            aethera::mcp::event_latency_benchmarks::configuration.height));
    aethera::mcp::event_latency_benchmarks::Event_Latency_Reporter reporter;
    benchmark::RunSpecifiedBenchmarks(&reporter);
    benchmark::Shutdown();
    return 0;
}