#include <ROOT/RNTuple.hxx>
#include <ROOT/RNTupleModel.hxx>
#include <ROOT/RNTupleWriter.hxx>
#include <ROOT/RVec.hxx>
#include <TFile.h>
#include <algorithm>
#include <filesystem>
#include <regex>
#include <spdlog/spdlog.h>
#include <vector>

#include "internal/mana_lib.hpp"
#include "internal/mana_root.hpp"
#include "internal/rxi/log.hpp"

using namespace mesytec::mnode;
using namespace mesytec::mnode::mana;

#if 0
using RNTuple = ROOT::Experimental::RNTuple;
#else
using RNTuple = ROOT::RNTuple;
#endif
using RNTupleModel = ROOT::Experimental::RNTupleModel;
using RNTupleWriter = ROOT::Experimental::RNTupleWriter;

struct Context
{
    std::vector<std::shared_ptr<std::vector<float>>> rvecs;
    std::unique_ptr<RNTupleWriter> writer;
};

static Context *g_ctx = nullptr;

MANA_DEFINE_PLUGIN_INIT(init)
{
    if (g_ctx)
    {
        log_warn("init() called multiple times. This plugin is a singleton!");
        return g_ctx;
    }
    log_set_level(LOG_INFO);
    log_debug("init");
    return g_ctx = new Context;
}

MANA_DEFINE_PLUGIN_SHUTDOWN(shutdown)
{
    log_debug("shutdown");
    if (context != g_ctx)
    {
        log_warn("shutdown() called with invalid context");
        return;
    }
    delete g_ctx;
    g_ctx = nullptr;
}

MANA_DEFINE_PLUGIN_BEGIN_RUN(begin_run)
{
    assert(context);
    assert(context == g_ctx);
    log_debug("begin_run: context=%p, descriptor_json=%s", context, descriptor_json);
    auto jRun = nlohmann::json::parse(descriptor_json);

    if (jRun["events"].size() > 1)
    {
        log_error("root-rntuple-writer: currently only supports a single event/trigger");
        mana_status_t result;
        result.code = mana_status_invalid_argument;
        strncpy(result.message,
                "root-rntuple-writer currently only supports a single event/trigger",
                sizeof(result.message));
        return result;
    }

    auto ctx = reinterpret_cast<Context *>(context);
    ctx->rvecs.clear();
    auto jEvent = jRun["events"][0];
    auto model = RNTupleModel::Create();
    for (const auto &jOutput: jEvent["outputs"])
    {
        auto name = jOutput["name"].get<std::string>();
        std::replace(std::begin(name), std::end(name), '.', '_');
        auto rvec = model->MakeField<std::vector<float>>(name);
        rvec->resize(jOutput["size"].get<size_t>());
        ctx->rvecs.emplace_back(std::move(rvec));
    }

    std::filesystem::path rp(jRun["name"].get<std::string>());
    auto filename =
        fmt::format("{}_rntuple_events.root", rp.filename().replace_extension().string());

    log_info("writing rntuple into: %s", filename.c_str());

    ctx->writer =
        RNTupleWriter::Recreate(std::move(model), jEvent["name"].get<std::string>(), filename);
    mana_status_t result;
    result.code = mana_status_ok;
    strncpy(result.message, "ok", sizeof(result.message));
    return result;
}

MANA_DEFINE_PLUGIN_END_RUN(end_run)
{
    assert(context);
    assert(context == g_ctx);
    log_debug("end: context=%p, descriptor_json=%s", context, descriptor_json);
    auto ctx = reinterpret_cast<Context *>(context);
    *ctx = {};
    mana_status_t result;
    result.code = mana_status_ok;
    strncpy(result.message, "ok", sizeof(result.message));
    return result;
}

MANA_DEFINE_PLUGIN_EVENT_DATA(process_event)
{
    log_trace("event: ctx=%p, eventIndex=%d, arrayCount=%zu, totalBytes=%zu", context, eventIndex,
              arrayCount, totalBytes);
    if (eventIndex != 0)
        return;

    auto ctx = reinterpret_cast<Context *>(context);

    for (size_t ai = 0; ai < arrayCount; ++ai)
    {
        auto &rvec = ctx->rvecs.at(ai);
        auto input = mana::get_span<float>(arrays[ai]);

        rvec->resize(input.size());
        std::copy(std::begin(input), std::end(input), std::begin(*rvec));
    }
    ctx->writer->Fill();
}

MANA_DEFINE_PLUGIN_SYSTEM_EVENT(process_system_event)
{
    log_trace("system_event: ctx=%p, size=%zu", context, size);
}

extern "C" MANA_C_SINK_PLUGIN()
{
    mana_sink_plugin_t plugin;
    plugin.init = init;
    plugin.shutdown = shutdown;
    plugin.begin_run = begin_run;
    plugin.end_run = end_run;
    plugin.process_event = process_event;
    plugin.process_system_event = process_system_event;
    return plugin;
}
