#include <mcp/core/Control_Requests.hpp>
#include <mcp/core/Control_Service.hpp>
#include <mcp/core/Protocol_Type.hpp>
#include <benchmark/benchmark.h>
#include <render_common.hpp>
#include <algorithm>
#include <array>
#include <atomic>
#include <chrono>
#include <cstdint>
#include <limits>
#include <memory>
#include <optional>
#include <string_view>
#include <vector>

namespace aethera::mcp::benchmarks {
namespace {

[[nodiscard]] Plot_Input_Request timestamped_drag_request(
    Event_Type type = Event_Type::pointer_move) {
    Plot_Input_Request request;
    request.plot = "datoviz_point";
    request.event.type = type;
    request.event.time_milliseconds = 1'000.0;
    request.event.position = {240.0, 180.0};
    request.event.global_position = {240.0, 180.0};
    request.event.button = Mouse_Button::left;
    request.event.buttons = 1;
    return request;
}

void decode_timestamped_drag(benchmark::State& state) {
    const auto encoded = encode_protocol_value(timestamped_drag_request());
    for ([[maybe_unused]] auto iteration : state) {
        Plot_Input_Request decoded;
        decode_protocol_value(decoded, encoded);
        benchmark::DoNotOptimize(decoded.event.time_milliseconds);
        benchmark::ClobberMemory();
    }
    state.SetItemsProcessed(state.iterations());
}

void decode_timed_render(benchmark::State& state) {
    const auto encoded = encode_protocol_value(
        Plot_Render_Request{"datoviz_point", 1'000.0, 720, 420});
    for ([[maybe_unused]] auto iteration : state) {
        Plot_Render_Request decoded;
        decode_protocol_value(decoded, encoded);
        benchmark::DoNotOptimize(decoded.time_milliseconds);
        benchmark::ClobberMemory();
    }
    state.SetItemsProcessed(state.iterations());
}

class Control_Path : public benchmark::Fixture {
public:
    void SetUp(const benchmark::State&) override {
        service = Control_Service::create();
    }

    void TearDown(const benchmark::State&) override {
        service.reset();
    }

protected:
    std::shared_ptr<Control_Service> service;
};

constexpr std::array<std::string_view, 16> plot_3d_ids{
    "datoviz_point", "datoviz_splat", "datoviz_pixel", "datoviz_marker",
    "datoviz_sphere", "datoviz_segment", "datoviz_vector",
    "datoviz_primitive", "datoviz_mesh", "datoviz_spectrogram",
    "datoviz_path", "datoviz_image", "datoviz_labels", "datoviz_glyph",
    "datoviz_text", "datoviz_volume"};
constexpr std::size_t splat_plot_index{1};

enum struct Input_Workload : std::uint8_t {
    steady,
    drag,
    wheel,
    mixed
};

struct Concurrent_Workload_State {
    std::array<std::atomic_uint64_t, plot_3d_ids.size()> completed{}; /* Per-Plot publish counts for this benchmark interval. */
};

[[nodiscard]] Plot_Input_Request concurrent_input_request(
    Input_Workload workload, std::size_t plot_index,
    std::uint64_t sequence) {
    Plot_Input_Request request;
    request.plot = plot_3d_ids[plot_index];
    request.event.time_milliseconds = 1'000.0 +
        static_cast<double>(sequence) * (1'000.0 / 120.0);
    request.event.position = {
        360.0 + static_cast<double>(
            static_cast<std::int64_t>((sequence + plot_index * 7U) % 121U) - 60),
        210.0 + static_cast<double>(
            static_cast<std::int64_t>((sequence * 3U + plot_index * 11U) % 81U) - 40)};
    request.event.global_position = request.event.position;
    if (workload == Input_Workload::wheel ||
        (workload == Input_Workload::mixed && (plot_index & 1U) != 0U)) {
        request.event.type = Event_Type::wheel;
        request.event.pixel_delta_y = (sequence & 1U) != 0U ? 120.0 : -120.0;
        request.event.angle_delta_y = request.event.pixel_delta_y;
        return request;
    }
    request.event.type = Event_Type::pointer_move;
    request.event.button = Mouse_Button::left;
    request.event.buttons = 1;
    return request;
}

class Concurrent_3D : public benchmark::Fixture {
public:
    void SetUp(const benchmark::State&) override {
        if (shared_service) {
            service = shared_service;
            workload = shared_workload;
            plots = shared_plots;
            streams = shared_streams;
            return;
        }
        shared_service = Control_Service::create();
        shared_workload = std::make_shared<Concurrent_Workload_State>();
        shared_plots.reserve(plot_3d_ids.size());
        shared_streams.reserve(plot_3d_ids.size());
        for (std::size_t index = 0; index < plot_3d_ids.size(); ++index) {
            auto plot = shared_service->find_plot(plot_3d_ids[index]);
            if (!plot) throw std::logic_error("concurrent 3D benchmark Plot is unavailable");
            const auto stream = plot->subscribe(
                [workload = shared_workload, index](
                    std::shared_ptr<const web::Plot_Stream_Frame> frame) {
                    if (!frame || !frame->pixels) return;
                    workload->completed[index].fetch_add(
                        1, std::memory_order_relaxed);
                });
            plot->configure_stream(stream, 720, 420);
            shared_plots.push_back(std::move(plot));
            shared_streams.push_back(stream);
        }
        service = shared_service;
        workload = shared_workload;
        plots = shared_plots;
        streams = shared_streams;
    }

    void TearDown(const benchmark::State&) override {}

    static void shutdown() {
        if (!shared_workload) return;
        for (std::size_t index = 0; index < shared_plots.size(); ++index)
            shared_plots[index]->unsubscribe(shared_streams[index]);
        shared_plots.clear();
        shared_streams.clear();
        shared_service.reset();
        shared_workload.reset();
    }

protected:
    [[nodiscard]] bool submit_input_batch(
        Input_Workload input, std::uint64_t sequence,
        std::optional<Event_Type> pointer_type = std::nullopt) {
        for (std::size_t plot_index = 0;
             plot_index < plot_3d_ids.size(); ++plot_index) {
            const bool wheel = input == Input_Workload::wheel ||
                (input == Input_Workload::mixed && (plot_index & 1U) != 0U);
            if (pointer_type && wheel) continue;
            auto request = concurrent_input_request(
                input, plot_index, sequence);
            if (pointer_type) {
                request.event.type = *pointer_type;
                request.event.buttons = *pointer_type == Event_Type::pointer_release
                    ? 0 : 1;
            }
            const auto result = service->call_tool(
                "aethera_plot_input", encode_protocol_value(request));
            if (result.result != Tool_Call_Result::ok) return false;
        }
        return true;
    }

    void run(benchmark::State& state, Input_Workload input) {
        for (auto& count : workload->completed)
            count.store(0, std::memory_order_relaxed);
        for (const auto& plot : plots) plot->reset_diagnostics();
        if (input == Input_Workload::drag || input == Input_Workload::mixed) {
            if (!submit_input_batch(
                    input, 0, Event_Type::pointer_press)) {
                state.SkipWithError("16-Plot pointer press batch was rejected");
                return;
            }
        }
        const auto runtime_begin = service->call_tool(
            "aethera_task_runtime", nlohmann::json::object());
        const auto started = std::chrono::steady_clock::now();
        constexpr auto input_period = std::chrono::nanoseconds{
            1'000'000'000 / 120};
        auto next_input = started + input_period;
        std::uint64_t input_batches{};
        std::uint64_t input_requests{};
        std::uint32_t input_poll{};
        for ([[maybe_unused]] auto iteration : state) {
            benchmark::DoNotOptimize(
                workload->completed[splat_plot_index].load(
                    std::memory_order_relaxed));
            benchmark::ClobberMemory();
            if ((++input_poll & 0x3FFU) != 0U) continue;
            const auto now = std::chrono::steady_clock::now();
            if (input == Input_Workload::steady) continue;
            if (now < next_input) continue;
            ++input_batches;
            if (!submit_input_batch(input, input_batches)) {
                state.SkipWithError("16-Plot interaction batch was rejected");
                break;
            }
            input_requests += plot_3d_ids.size();
            do next_input += input_period; while (next_input <= now);
        }
        const auto finished = std::chrono::steady_clock::now();
        if (input == Input_Workload::drag || input == Input_Workload::mixed) {
            static_cast<void>(submit_input_batch(
                input, input_batches + 1U, Event_Type::pointer_release));
        }

        const double elapsed_seconds = std::chrono::duration<double>(
            finished - started).count();
        std::uint64_t total_completed{};
        double minimum_fps = std::numeric_limits<double>::max();
        double maximum_fps{};
        double splat_fps{};
        for (std::size_t index = 0; index < plots.size(); ++index) {
            const auto count = workload->completed[index].load(
                std::memory_order_relaxed);
            total_completed += count;
            const double fps = elapsed_seconds > 0.0
                ? static_cast<double>(count) / elapsed_seconds : 0.0;
            minimum_fps = std::min(minimum_fps, fps);
            maximum_fps = std::max(maximum_fps, fps);
            if (index == splat_plot_index) splat_fps = fps;

            const auto diagnostics = plots[index]->diagnostics();
            const auto& policy = diagnostics.at("frame_policy");
            const auto prefix = "p" + std::to_string(index) + "_";
            state.counters[prefix + "fps"] = fps;
            state.counters[prefix + "requested"] =
                policy.at("observation").at("request_count").get<double>();
            state.counters[prefix + "submitted"] =
                policy.at("observation").at("submitted_frame_count").get<double>();
            state.counters[prefix + "completed"] =
                policy.at("observation").at("completed_frame_count").get<double>();
            state.counters[prefix + "active"] =
                policy.at("observation").at("active_frame_count").get<double>();
            state.counters[prefix + "scene_rejected"] =
                policy.at("requests").at("scene_rejected").get<double>();
            state.counters[prefix + "slot_backpressure"] =
                policy.at("requests").at("frame_slot_backpressure").get<double>();
            state.counters[prefix + "completion_ms"] =
                policy.at("latency").at("average_completion_ms").get<double>();
            const auto& frame_statistics = diagnostics.at("frame_statistics");
            const auto statistic_average = [&](std::string_view name) {
                const auto found = frame_statistics.find(name);
                return found == frame_statistics.end()
                    ? 0.0 : found->at("average").get<double>();
            };
            state.counters[prefix + "backend_queue_ms"] =
                statistic_average("backend_queue_ms");
            state.counters[prefix + "backend_plan_ms"] =
                statistic_average("backend_plan_ms");
        }
        const auto runtime_end = service->call_tool(
            "aethera_task_runtime", nlohmann::json::object());
        const auto runtime_delta = [&](std::string_view key) {
            if (runtime_begin.result != Tool_Call_Result::ok ||
                runtime_end.result != Tool_Call_Result::ok) return 0.0;
            const auto begin = runtime_begin.content.at(key).get<std::uint64_t>();
            const auto end = runtime_end.content.at(key).get<std::uint64_t>();
            return static_cast<double>(end >= begin ? end - begin : 0U);
        };
        const double wall_ns = runtime_delta("observed_wall_time_ns");
        const double busy_ns = runtime_delta("worker_busy_time_ns");
        const double cpu_ns = runtime_delta("worker_cpu_time_ns");
        const double worker_count = runtime_end.result == Tool_Call_Result::ok
            ? runtime_end.content.at("worker_count").get<double>() : 0.0;
        state.counters["aggregate_fps"] = elapsed_seconds > 0.0
            ? static_cast<double>(total_completed) / elapsed_seconds : 0.0;
        state.counters["minimum_plot_fps"] = minimum_fps;
        state.counters["maximum_plot_fps"] = maximum_fps;
        state.counters["splat_fps"] = splat_fps;
        state.counters["fairness_pct"] = maximum_fps > 0.0
            ? minimum_fps / maximum_fps * 100.0 : 0.0;
        state.counters["worker_busy_pct"] = wall_ns > 0.0 && worker_count > 0.0
            ? busy_ns / (wall_ns * worker_count) * 100.0 : 0.0;
        state.counters["worker_cpu_pct"] = wall_ns > 0.0 && worker_count > 0.0
            ? cpu_ns / (wall_ns * worker_count) * 100.0 : 0.0;
        state.counters["input_batches"] = static_cast<double>(input_batches);
        state.counters["input_requests"] = static_cast<double>(input_requests);
        state.counters["input_request_rate"] = elapsed_seconds > 0.0
            ? static_cast<double>(input_requests) / elapsed_seconds : 0.0;
        state.SetItemsProcessed(static_cast<std::int64_t>(total_completed));
    }

private:
    inline static std::shared_ptr<Control_Service> shared_service{};
    inline static std::shared_ptr<Concurrent_Workload_State> shared_workload{};
    inline static std::vector<std::shared_ptr<web::Plot>> shared_plots{};
    inline static std::vector<web::Plot::Stream_Id> shared_streams{};
    std::shared_ptr<Control_Service> service{};
    std::shared_ptr<Concurrent_Workload_State> workload{};
    std::vector<std::shared_ptr<web::Plot>> plots{};
    std::vector<web::Plot::Stream_Id> streams{};
};

BENCHMARK_DEFINE_F(Concurrent_3D, Warmup)(benchmark::State& state) {
    run(state, Input_Workload::steady);
}

BENCHMARK_DEFINE_F(Concurrent_3D, Steady)(benchmark::State& state) {
    run(state, Input_Workload::steady);
}
BENCHMARK_DEFINE_F(Concurrent_3D, All_Plots_Drag)(benchmark::State& state) {
    run(state, Input_Workload::drag);
}
BENCHMARK_DEFINE_F(Concurrent_3D, All_Plots_Wheel)(benchmark::State& state) {
    run(state, Input_Workload::wheel);
}
BENCHMARK_DEFINE_F(Concurrent_3D, All_Plots_Mixed)(benchmark::State& state) {
    run(state, Input_Workload::mixed);
}

BENCHMARK_DEFINE_F(Control_Path, Timestamped_Drag_Admission)(
    benchmark::State& state) {
    auto request = timestamped_drag_request(Event_Type::pointer_press);
    auto arguments = encode_protocol_value(request);
    const auto press = service->call_tool("aethera_plot_input", arguments);
    if (press.result != Tool_Call_Result::ok) {
        state.SkipWithError("timestamped pointer press was rejected");
        return;
    }

    arguments["event"]["type"] = "pointer_move";
    double time_milliseconds = request.event.time_milliseconds;
    for ([[maybe_unused]] auto iteration : state) {
        time_milliseconds += 1'000.0 / 120.0;
        arguments["event"]["time_milliseconds"] = time_milliseconds;
        const auto result = service->call_tool(
            "aethera_plot_input", arguments);
        if (result.result != Tool_Call_Result::ok) {
            state.SkipWithError("timestamped pointer move was rejected");
            break;
        }
        benchmark::DoNotOptimize(result.result);
    }
    state.SetItemsProcessed(state.iterations());

    const auto runtime = service->call_tool(
        "aethera_task_runtime", nlohmann::json::object());
    if (runtime.result != Tool_Call_Result::ok) {
        state.SkipWithError("Taskflow runtime diagnostics were rejected");
        return;
    }
    const auto milliseconds = [](const nlohmann::json& value,
                                 std::string_view key) {
        return static_cast<double>(value.at(key).get<std::uint64_t>()) /
               1'000'000.0;
    };
    state.counters["taskflow_wall_ms"] =
        milliseconds(runtime.content, "observed_wall_time_ns");
    state.counters["worker_busy_ms"] =
        milliseconds(runtime.content, "worker_busy_time_ns");
    state.counters["worker_cpu_ms"] =
        milliseconds(runtime.content, "worker_cpu_time_ns");
    state.counters["worker_utilization_pct"] =
        runtime.content.at("worker_utilization").get<double>();
    state.counters["worker_cpu_utilization_pct"] =
        runtime.content.at("worker_cpu_utilization").get<double>();
    state.counters["active_taskflows"] =
        static_cast<double>(runtime.content.at("active_taskflows").get<std::size_t>());
    state.counters["completed_taskflows"] =
        static_cast<double>(runtime.content.at("completed_taskflows").get<std::size_t>());
    state.SetLabel(runtime.content.at("longest_task").at("name").get<std::string>());
}

BENCHMARK(decode_timestamped_drag);
BENCHMARK(decode_timed_render);
BENCHMARK_REGISTER_F(Control_Path, Timestamped_Drag_Admission)
    ->Iterations(512);
/* Fixture 静态持有 16 个 Plot；MinTime 校准不会重建 Datoviz Scene。 */
BENCHMARK_REGISTER_F(Concurrent_3D, Warmup)
    ->MinTime(10.0)->UseRealTime();
BENCHMARK_REGISTER_F(Concurrent_3D, Steady)
    ->MinTime(10.0)->UseRealTime();
BENCHMARK_REGISTER_F(Concurrent_3D, All_Plots_Drag)
    ->MinTime(10.0)->UseRealTime();
BENCHMARK_REGISTER_F(Concurrent_3D, All_Plots_Wheel)
    ->MinTime(10.0)->UseRealTime();
BENCHMARK_REGISTER_F(Concurrent_3D, All_Plots_Mixed)
    ->MinTime(10.0)->UseRealTime();

void shutdown_concurrent_3d_benchmark() {
    Concurrent_3D::shutdown();
}

}
}

int main(int argc, char** argv) {
    aethera::initialize_runtime({});
    benchmark::Initialize(&argc, argv);
    if (benchmark::ReportUnrecognizedArguments(argc, argv)) return 2;
    benchmark::RunSpecifiedBenchmarks();
    aethera::mcp::benchmarks::shutdown_concurrent_3d_benchmark();
    benchmark::Shutdown();
    return 0;
}