#include "Taskflow_Trace_Json.hpp"
#include <magic_enum/magic_enum.hpp>
#include <nlohmann/json.hpp>
#include <string>
#include <unordered_map>

namespace aethera::web {

nlohmann::json taskflow_trace_json(
    const Taskflow_Frame_Trace& trace,
    const nlohmann::json& captured_components,
    const nlohmann::json& captured_backend) {
    nlohmann::json markers = nlohmann::json::object();
    for (const auto& marker : trace.markers)
        markers[magic_enum::enum_name(marker.marker)] =
            static_cast<double>(marker.elapsed_ns) / 1'000'000.0;
    nlohmann::json measurements = nlohmann::json::object();
    for (const auto& measurement : trace.measurements)
        measurements[magic_enum::enum_name(measurement.measurement)] =
            static_cast<double>(measurement.value_ns) / 1'000'000.0;
    nlohmann::json graphs = nlohmann::json::array();
    std::unordered_map<std::uint64_t, std::string> node_ids;
    for (const auto& graph : trace.graphs) {
        nlohmann::json nodes = nlohmann::json::array();
        for (const auto& node : graph.nodes) {
            node_ids.emplace(node.native_id, node.node_id);
            nlohmann::json predecessors = nlohmann::json::array();
            for (const auto native_id : node.predecessors)
                predecessors.push_back(std::to_string(native_id));
            nlohmann::json successors = nlohmann::json::array();
            for (const auto native_id : node.successors)
                successors.push_back(std::to_string(native_id));
            nlohmann::json attributes = nlohmann::json::object();
            for (const auto& [key, value] : node.attributes)
                attributes[key] = value;
            nlohmann::json encoded{
                {"native_id", std::to_string(node.native_id)},
                {"id", node.node_id}, {"parent_id", node.parent_node_id},
                {"name", node.name}, {"type", node.type},
                {"predecessors", std::move(predecessors)},
                {"successors", std::move(successors)},
                {"attributes", std::move(attributes)}};
            const auto owner = encoded["attributes"].value(
                "owner_component", std::string{});
            if (!owner.empty() && captured_components.contains(owner)) {
                const auto& captured = captured_components.at(owner);
                encoded["owner"] = {
                    {"component", owner},
                    {"label", captured.value("label", owner)},
                    {"kind", captured.value("kind", std::string{})}};
                encoded["prop"] = captured.value(
                    "prop", nlohmann::json::object());
                encoded["state"] = captured.value(
                    "state", nlohmann::json::object());
            }
            nodes.push_back(std::move(encoded));
        }
        graphs.push_back({
            {"stage", graph.stage}, {"name", graph.taskflow_name},
            {"submitted_ms", graph.submitted_ms},
            {"finished_ms", graph.finished_ms},
            {"completed", graph.completed}, {"nodes", std::move(nodes)}});
    }
    nlohmann::json executions = nlohmann::json::array();
    for (const auto& task : trace.tasks) {
        const auto found = node_ids.find(task.native_id);
        nlohmann::json waits = nlohmann::json::array();
        for (const auto& wait : task.cooperative_waits)
            waits.push_back({{"started_ms", wait.started_ms},
                             {"finished_ms", wait.finished_ms}});
        executions.push_back({
            {"native_id", std::to_string(task.native_id)},
            {"node_id", found == node_ids.end() ? std::string{} : found->second},
            {"worker_id", task.worker_id},
            {"worker_queue_size", task.worker_queue_size},
            {"worker_queue_capacity", task.worker_queue_capacity},
            {"ready_ms", task.ready_ms}, {"entered_ms", task.entered_ms},
            {"started_ms", task.started_ms}, {"finished_ms", task.finished_ms},
            {"completed_ms", task.completed_ms},
            {"duration_ms", task.duration_ms},
            {"cooperative_wait_ms", task.cooperative_wait_ms},
            {"cooperative_waits", std::move(waits)},
            {"observer_entry_ms", task.observer_entry_ms},
            {"observer_exit_ms", task.observer_exit_ms},
            {"queue_wait_ms", task.queue_wait_ms}});
    }
    nlohmann::json result{
        {"sequence", trace.identity.sequence},
        {"correlation_id", trace.identity.correlation_id},
        {"request_source", magic_enum::enum_name(trace.request_source)},
        {"created_time_unix_ns", trace.created_time_unix_ns},
        {"worker_count", trace.worker_count},
        {"markers", std::move(markers)},
        {"measurements", std::move(measurements)},
        {"graphs", std::move(graphs)},
        {"executions", std::move(executions)}};
    if (!captured_backend.empty()) result["datoviz"] = captured_backend;
    return result;
}

}