#include "Control_Service.hpp"
#include "Control_Service.ipp"
#include "Control_Requests.hpp"
#include "Protocol_Type.hpp"
#include "runtime/Gallery_Plots.hpp"
#include <render_common.hpp>
#include <array>
#include <chrono>
#include <stdexcept>
#include <utility>

namespace aethera::mcp {
namespace {

using Invoke = Tool_Call_Output (*)(Control_Service&, const nlohmann::json&);
using Schema = nlohmann::json (*)();

struct Operation {
    std::string_view name;                                  /* MCP 稳定 tool name。 */
    std::string_view description;                           /* 模型选择工具时使用的业务说明。 */
    Schema schema;                                          /* PFR 自动生成的输入结构。 */
    Invoke invoke;                                          /* 协议无关的业务调用入口。 */
};

template <typename Request>
[[nodiscard]] nlohmann::json request_schema() {
    return describe_protocol_type<Request>();
}

template <typename Request, typename Callback>
[[nodiscard]] Tool_Call_Output decode_and_call(
    const nlohmann::json& arguments, Callback&& callback) {
    try {
        Request request{};
        decode_protocol_value(request, arguments);
        return std::forward<Callback>(callback)(request);
    } catch (const nlohmann::json::exception& failure) {
        return {Tool_Call_Result::invalid_arguments, {}, failure.what()};
    } catch (const std::invalid_argument& failure) {
        return {Tool_Call_Result::invalid_arguments, {}, failure.what()};
    }
}

[[nodiscard]] Tool_Call_Output list_plots(
    Control_Service& service, const nlohmann::json& arguments) {
    return decode_and_call<Empty_Request>(arguments, [&service](const auto&) {
        return Tool_Call_Output{Tool_Call_Result::ok, service.plot_catalog(), {}};
    });
}

[[nodiscard]] Tool_Call_Output plot_schema(
    Control_Service& service, const nlohmann::json& arguments) {
    return decode_and_call<Plot_Request>(arguments, [&service](const auto& request) {
        const auto plot = service.find_plot(request.plot);
        if (!plot) return Tool_Call_Output{Tool_Call_Result::unknown_plot, {}, "unknown plot"};
        return Tool_Call_Output{Tool_Call_Result::ok, plot->schema(), {}};
    });
}

[[nodiscard]] Tool_Call_Output plot_diagnostics(
    Control_Service& service, const nlohmann::json& arguments) {
    return decode_and_call<Plot_Request>(arguments, [&service](const auto& request) {
        const auto plot = service.find_plot(request.plot);
        if (!plot) return Tool_Call_Output{Tool_Call_Result::unknown_plot, {}, "unknown plot"};
        return Tool_Call_Output{Tool_Call_Result::ok, plot->diagnostics(), {}};
    });
}

[[nodiscard]] Tool_Call_Output reset_plot_diagnostics(
    Control_Service& service, const nlohmann::json& arguments) {
    return decode_and_call<Plot_Request>(arguments, [&service](const auto& request) {
        const auto plot = service.find_plot(request.plot);
        if (!plot) return Tool_Call_Output{Tool_Call_Result::unknown_plot, {}, "unknown plot"};
        plot->reset_diagnostics();
        return Tool_Call_Output{Tool_Call_Result::ok, {{"accepted", true}}, {}};
    });
}

[[nodiscard]] Tool_Call_Output component_state(
    Control_Service& service, const nlohmann::json& arguments) {
    return decode_and_call<Component_Request>(arguments, [&service](const auto& request) {
        const auto plot = service.find_plot(request.plot);
        if (!plot) return Tool_Call_Output{Tool_Call_Result::unknown_plot, {}, "unknown plot"};
        auto result = plot->component_state(request.component);
        if (!result.value("success", true))
            return Tool_Call_Output{Tool_Call_Result::rejected, std::move(result), "unknown component"};
        return Tool_Call_Output{Tool_Call_Result::ok, std::move(result), {}};
    });
}

[[nodiscard]] Tool_Call_Output write_property(
    Control_Service& service, const nlohmann::json& arguments) {
    return decode_and_call<Write_Property_Request>(arguments, [&service](const auto& request) {
        const auto plot = service.find_plot(request.plot);
        if (!plot) return Tool_Call_Output{Tool_Call_Result::unknown_plot, {}, "unknown plot"};
        auto result = plot->write_prop(
            request.component, request.property, request.value);
        const bool success = result.value("success", false);
        return Tool_Call_Output{success ? Tool_Call_Result::ok : Tool_Call_Result::rejected,
                                std::move(result), success ? "" : "property write rejected"};
    });
}

[[nodiscard]] Tool_Call_Output generate_data(
    Control_Service& service, const nlohmann::json& arguments) {
    return decode_and_call<Generate_Data_Request>(arguments, [&service](const auto& request) {
        const auto plot = service.find_plot(request.plot);
        if (!plot) return Tool_Call_Output{Tool_Call_Result::unknown_plot, {}, "unknown plot"};
        return Tool_Call_Output{Tool_Call_Result::ok,
                                plot->generate_data(request.input), {}};
    });
}

[[nodiscard]] Tool_Call_Output submit_plot_input(
    Control_Service& service, const nlohmann::json& arguments) {
    return decode_and_call<Plot_Input_Request>(
        arguments, [&service](auto request) {
            const auto plot = service.find_plot(request.plot);
            if (!plot)
                return Tool_Call_Output{
                    Tool_Call_Result::unknown_plot, {}, "unknown plot"};
            const auto time_milliseconds = request.event.time_milliseconds;
            plot->submit_input(std::move(request.event));
            return Tool_Call_Output{
                Tool_Call_Result::ok,
                {{"accepted", true},
                 {"plot", request.plot},
                 {"time_milliseconds", time_milliseconds}}, {}};
        });
}

[[nodiscard]] Tool_Call_Output render_plot_at(
    Control_Service& service, const nlohmann::json& arguments) {
    return decode_and_call<Plot_Render_Request>(
        arguments, [&service](const auto& request) {
            const auto plot = service.find_plot(request.plot);
            if (!plot)
                return Tool_Call_Output{
                    Tool_Call_Result::unknown_plot, {}, "unknown plot"};
            plot->schedule_render(Frame_Request{
                .issued_at = std::chrono::steady_clock::now(),
                .time_milliseconds = request.time_milliseconds,
                .width = request.width,
                .height = request.height,
                .source = Frame_Request_Source::immediate});
            return Tool_Call_Output{
                Tool_Call_Result::ok,
                {{"accepted", true},
                 {"plot", request.plot},
                 {"time_milliseconds", request.time_milliseconds},
                 {"width", request.width}, {"height", request.height}}, {}};
        });
}

[[nodiscard]] Tool_Call_Output begin_frame_trace(
    Control_Service& service, const nlohmann::json& arguments) {
    return decode_and_call<Frame_Trace_Request>(
        arguments, [&service](const auto& request) {
            const auto plot = service.find_plot(request.plot);
            if (!plot)
                return Tool_Call_Output{
                    Tool_Call_Result::unknown_plot, {}, "unknown plot"};
            plot->request_taskflow_trace(request.frame_count);
            return Tool_Call_Output{
                Tool_Call_Result::ok,
                {{"accepted", true}, {"plot", request.plot},
                 {"frame_count", request.frame_count},
                 {"status_tool", "aethera_frame_trace_read"}}, {}};
        });
}

[[nodiscard]] Tool_Call_Output read_frame_trace(
    Control_Service& service, const nlohmann::json& arguments) {
    return decode_and_call<Plot_Request>(
        arguments, [&service](const auto& request) {
            const auto plot = service.find_plot(request.plot);
            if (!plot)
                return Tool_Call_Output{
                    Tool_Call_Result::unknown_plot, {}, "unknown plot"};
            return Tool_Call_Output{
                Tool_Call_Result::ok, plot->taskflow_trace(), {}};
        });
}

[[nodiscard]] Tool_Call_Output begin_benchmark(
    Control_Service& service, const nlohmann::json& arguments) {
    return decode_and_call<Plot_Request>(arguments, [&service](const auto& request) {
        const auto plot = service.find_plot(request.plot);
        if (!plot) return Tool_Call_Output{Tool_Call_Result::unknown_plot, {}, "unknown plot"};
        plot->reset_diagnostics();
        plot->render_once();
        return Tool_Call_Output{Tool_Call_Result::ok,
            {{"accepted", true}, {"plot", request.plot},
             {"status_tool", "aethera_benchmark_read"}}, {}};
    });
}

[[nodiscard]] nlohmann::json task_runtime_json() {
    const auto state = task_runtime_state();
    nlohmann::json workers = nlohmann::json::array();
    for (const auto& worker : state.workers)
        workers.push_back({
            {"id", worker.id}, {"task_count", worker.task_count},
            {"entry_queue_size", worker.current_queue_size},
            {"entry_queue_capacity", worker.current_queue_capacity},
            {"peak_queue_size", worker.peak_observed_queue_size},
            {"max_queue_capacity", worker.max_observed_queue_capacity},
            {"active_task", {{"native_id", std::to_string(worker.active_task_hash)},
                              {"type", worker.active_task_type},
                              {"time_ns", worker.active_task_time_ns}}},
            {"task_time_ns", worker.task_time_ns},
            {"busy_time_ns", worker.busy_time_ns},
            {"cpu_time_ns", worker.cpu_time_ns},
            {"cooperative_wait_count", worker.cooperative_wait_count},
            {"cooperative_wait_time_ns", worker.cooperative_wait_time_ns},
            {"idle_time_ns", worker.idle_time_ns},
            {"min_task_time_ns", worker.min_task_time_ns},
            {"max_task_time_ns", worker.max_task_time_ns},
            {"utilization", worker.utilization},
            {"cpu_utilization", worker.cpu_utilization}});
    nlohmann::json task_types = nlohmann::json::array();
    for (const auto& type : state.task_types)
        task_types.push_back({
            {"name", type.name}, {"count", type.count},
            {"total_time_ns", type.total_time_ns},
            {"min_time_ns", type.min_time_ns},
            {"max_time_ns", type.max_time_ns}});
    return {
        {"protocol", "aethera.taskflow.runtime"}, {"version", 2},
        {"worker_count", state.worker_count},
        {"active_topologies", state.active_topology_count},
        {"active_taskflows", state.active_taskflow_count},
        {"peak_active_taskflows", state.peak_active_taskflow_count},
        {"completed_taskflows", state.completed_taskflow_count},
        {"failed_taskflows", state.failed_taskflow_count},
        {"active_tasks", state.active_task_count},
        {"peak_active_tasks", state.peak_active_task_count},
        {"active_workers", state.active_worker_count},
        {"peak_active_workers", state.peak_active_worker_count},
        {"observed_tasks", state.observed_task_count},
        {"named_tasks", state.named_task_count},
        {"peak_worker_queue_size", state.peak_observed_worker_queue_size},
        {"max_worker_queue_capacity", state.max_observed_worker_queue_capacity},
        {"longest_task", {{"native_id", std::to_string(state.longest_task_hash)},
                           {"name", state.longest_task_name},
                           {"type", state.longest_task_type},
                           {"time_ns", state.longest_task_time_ns}}},
        {"total_task_time_ns", state.total_task_time_ns},
        {"worker_busy_time_ns", state.worker_busy_time_ns},
        {"worker_cpu_time_ns", state.worker_cpu_time_ns},
        {"cooperative_wait_count", state.cooperative_wait_count},
        {"cooperative_wait_time_ns", state.cooperative_wait_time_ns},
        {"observed_wall_time_ns", state.observed_wall_time_ns},
        {"worker_utilization", state.worker_utilization},
        {"worker_cpu_utilization", state.worker_cpu_utilization},
        {"task_types", std::move(task_types)}, {"workers", std::move(workers)}};
}

[[nodiscard]] Tool_Call_Output task_runtime(
    Control_Service&, const nlohmann::json& arguments) {
    return decode_and_call<Empty_Request>(arguments, [](const auto&) {
        return Tool_Call_Output{Tool_Call_Result::ok, task_runtime_json(), {}};
    });
}

constexpr std::array operations{
    Operation{"aethera_plot_list", "List every 2D and 3D gallery plot.",
              &request_schema<Empty_Request>, &list_plots},
    Operation{"aethera_plot_schema", "Describe a plot and its editable render components.",
              &request_schema<Plot_Request>, &plot_schema},
    Operation{"aethera_plot_diagnostics", "Read current frame, scene, GPU and Datoviz diagnostics.",
              &request_schema<Plot_Request>, &plot_diagnostics},
    Operation{"aethera_plot_diagnostics_reset", "Reset the plot's authoritative diagnostic counters.",
              &request_schema<Plot_Request>, &reset_plot_diagnostics},
    Operation{"aethera_component_state", "Read one render component's published state.",
              &request_schema<Component_Request>, &component_state},
    Operation{"aethera_component_write", "Write one editable render component property.",
              &request_schema<Write_Property_Request>, &write_property},
    Operation{"aethera_data_generate", "Generate input data for a gallery plot.",
              &request_schema<Generate_Data_Request>, &generate_data},
    Operation{"aethera_plot_input", "Submit one timestamped input event without waiting for rendering.",
              &request_schema<Plot_Input_Request>, &submit_plot_input},
    Operation{"aethera_plot_render_at", "Request one diagnostic frame at a caller-owned timeline time.",
              &request_schema<Plot_Render_Request>, &render_plot_at},
    Operation{"aethera_frame_trace_begin", "Capture existing per-frame Taskflow and Datoviz observations.",
              &request_schema<Frame_Trace_Request>, &begin_frame_trace},
    Operation{"aethera_frame_trace_read", "Read the current physical-frame trace capture.",
              &request_schema<Plot_Request>, &read_frame_trace},
    Operation{"aethera_benchmark_begin", "Reset diagnostics and asynchronously start rendering a plot.",
              &request_schema<Plot_Request>, &begin_benchmark},
    Operation{"aethera_benchmark_read", "Read benchmark results from the plot's current diagnostics.",
              &request_schema<Plot_Request>, &plot_diagnostics},
    Operation{"aethera_task_runtime", "Read Taskflow executor utilization and queue diagnostics.",
              &request_schema<Empty_Request>, &task_runtime},
};

}

Control_Service::Control_Service() : d(std::make_unique<Private>()) {}
Control_Service::~Control_Service() = default;

std::shared_ptr<Control_Service> Control_Service::create() {
    auto service = std::shared_ptr<Control_Service>(new Control_Service);
    service->d->plots.reserve(web::gallery_plot_definitions().size());
    for (const auto& definition : web::gallery_plot_definitions())
        service->d->plots.emplace(definition.id, definition.create());
    return service;
}

std::shared_ptr<web::Plot> Control_Service::find_plot(
    std::string_view id) const {
    const auto found = d->plots.find(std::string{id});
    return found == d->plots.end() ? nullptr : found->second;
}

nlohmann::json Control_Service::plot_catalog() const {
    nlohmann::json result = nlohmann::json::array();
    for (const auto& definition : web::gallery_plot_definitions())
        result.push_back({
            {"id", definition.id}, {"title", definition.title},
            {"category", definition.category},
            {"description", definition.description},
            {"dimension", web::plot_dimension_name(definition.dimension)}});
    return result;
}

nlohmann::json Control_Service::tool_catalog() const {
    nlohmann::json result = nlohmann::json::array();
    for (const auto& operation : operations)
        result.push_back({{"name", operation.name},
                          {"description", operation.description},
                          {"inputSchema", operation.schema()}});
    return result;
}

Tool_Call_Output Control_Service::call_tool(
    std::string_view name, const nlohmann::json& arguments) {
    for (const auto& operation : operations)
        if (operation.name == name) return operation.invoke(*this, arguments);
    return {Tool_Call_Result::unknown_tool, {}, "unknown tool"};
}

}