#include "Web_Server.hpp"
#include "Graph_WebSocket.hpp"
#include "Gallery_Plots.hpp"
#include <asio/thread_pool.hpp>
#include <drogon/drogon.h>
#include <nlohmann/json.hpp>
#include <algorithm>
#include <functional>
#include <memory>
#include <string>
#include <thread>
#include <unordered_map>

namespace aethera::web {
namespace {
using Plot_Map = std::unordered_map<std::string, std::shared_ptr<Plot>>;

drogon::HttpResponsePtr json_response(nlohmann::json value) {
    auto response = drogon::HttpResponse::newHttpResponse();
    response->setContentTypeCode(drogon::CT_APPLICATION_JSON);
    response->setBody(value.dump());
    return response;
}

drogon::HttpResponsePtr error_response(drogon::HttpStatusCode status, std::string message) {
    auto response = json_response({{"error", std::move(message)}});
    response->setStatusCode(status);
    return response;
}

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

int run_web_server(std::uint16_t port, const std::filesystem::path& asset_root) {
    const auto hardware_threads = std::max(2U, std::thread::hardware_concurrency());
    auto graph_pool = std::make_shared<asio::thread_pool>(std::min(8U, hardware_threads));
    const auto executor = graph_pool->get_executor();

    auto plots = std::make_shared<Plot_Map>();
    plots->emplace("spectrum", make_spectrum_plot(executor));
    plots->emplace("frequency_trace", make_frequency_trace_plot(executor));
    plots->emplace("sweep_spectrum", make_sweep_spectrum_plot(executor));
    plots->emplace("afterglow", make_afterglow_plot(executor));
    plots->emplace("waterfall", make_waterfall_plot(executor));
    plots->emplace("constellation", make_constellation_plot(executor));
    plots->emplace("selection_overlay", make_selection_overlay_plot(executor));
    plots->emplace("datoviz_point", make_datoviz_point_plot(executor));

    auto resolve_plot = [plots](std::string_view id) { return find_plot(*plots, id); };
    auto websocket = std::make_shared<Graph_WebSocket_Controller>(resolve_plot);
    auto& app = drogon::app();

    app.registerHandler("/plot", [plots](const drogon::HttpRequestPtr&,
                                         std::function<void(const drogon::HttpResponsePtr&)>&& callback) {
        nlohmann::json result = nlohmann::json::array();
        for (const auto& [id, plot] : *plots) {
            static_cast<void>(plot);
            result.push_back({
                {"id", id}, {"title", id}, {"category", "Plots"}, {"description", ""},
                {"dimension", ""}, {"websocket", "/ws/plot/" + id},
                {"schema", "/plot/" + id + "/schema"}});
        }
        callback(json_response(std::move(result)));
    }, {drogon::Get});

    app.registerHandler("/plot/{1}/schema", [plots](
        const drogon::HttpRequestPtr&,
        std::function<void(const drogon::HttpResponsePtr&)>&& callback,
        std::string plot_id) {
        auto plot = find_plot(*plots, plot_id);
        if (!plot) {
            callback(error_response(drogon::k404NotFound, "unknown plot"));
            return;
        }
        auto output = std::make_shared<std::function<void(const drogon::HttpResponsePtr&)>>(
            std::move(callback));
        plot->async_schema([output](nlohmann::json schema) {
            (*output)(json_response(std::move(schema)));
        });
    }, {drogon::Get});

    app.registerHandler("/plot/{1}/component/{2}/prop/{3}", [plots](
        const drogon::HttpRequestPtr& request,
        std::function<void(const drogon::HttpResponsePtr&)>&& callback,
        std::string plot_id,
        std::string component,
        std::string key) {
        auto plot = find_plot(*plots, plot_id);
        if (!plot) {
            callback(error_response(drogon::k404NotFound, "unknown plot"));
            return;
        }
        nlohmann::json value;
        try {
            value = nlohmann::json::parse(request->body());
        } catch (const nlohmann::json::exception&) {
            callback(error_response(drogon::k400BadRequest, "invalid JSON value"));
            return;
        }
        auto output = std::make_shared<std::function<void(const drogon::HttpResponsePtr&)>>(
            std::move(callback));
        plot->async_write_prop(std::move(component), std::move(key), std::move(value), [output](nlohmann::json result) {
            (*output)(json_response(std::move(result)));
        });
    }, {drogon::Put});

    app.registerController(websocket)
        .setDocumentRoot(asset_root.string())
        .setHomePage("index.html")
        .setStaticFileHeaders({{"Cache-Control", "no-store"}})
        .addListener("127.0.0.1", port)
        .setThreadNum(std::min(8U, hardware_threads))
        .setIdleConnectionTimeout(90)
        .run();
    graph_pool->stop();
    graph_pool->join();
    return 0;
}
}