#include <mesytec-mvlc/cpp_compat.h>
#include <pybind11/embed.h>
#include <pybind11/stl.h>

#include "internal/argh.h"
#include "internal/mana_lib.hpp"
#include "internal/rxi/log.hpp"

using namespace mesytec;
using namespace mesytec::mnode;
namespace py = pybind11;

PYBIND11_EMBEDDED_MODULE(py_mana, m) {}

namespace
{
struct Context
{
    py::scoped_interpreter interp;
    py::module usercode;
    py::object userobject;
    py::object py_begin_run;
    py::object py_process_event;
    py::object py_process_system_event;
    py::object py_end_run;
    std::vector<std::vector<py::object>> eventBuffers;
};

void reload_py(Context *ctx)
{
    ctx->usercode.reload();
    // TODO: check if the retrieved attributes are callable
    ctx->py_begin_run = ctx->usercode.attr("begin_run");
    ctx->py_process_event = ctx->usercode.attr("process_event");
    ctx->py_process_system_event = ctx->usercode.attr("process_system_event");
    ctx->py_end_run = ctx->usercode.attr("end_run");
}

} // namespace

MANA_DEFINE_PLUGIN_INIT(init)
{
    log_set_level(LOG_INFO);
    static Context g_ctx;
    auto ctx = &g_ctx;
    log_debug("init: ctx=%p", ctx);
    try
    {
        argh::parser parser;
        parser.parse(plugin_argc, plugin_argv);
        if (parser.size() != 1)
        {
            log_error("missing <script.py> argument");
            return nullptr;
        }
        auto filename = parser[0];
        ctx->usercode = py::module_::import(filename.c_str());
        try
        {
            py::exec(R"(import debugpy; debugpy.listen(5678))");
            // py::exec(R"(import numpy)");
            // py::module::import("numpy");
        }
        catch (const std::exception &e)
        {
            log_warn("warning: startup code failed: %s", e.what());
        }
        reload_py(ctx);
    }
    catch (const std::exception &e)
    {
        log_error("error: %s", e.what());
        return nullptr;
    }
    return ctx;
}

MANA_DEFINE_PLUGIN_SHUTDOWN(shutdown) { log_info("shutdown: ctx=%p", context); }

MANA_DEFINE_PLUGIN_BEGIN_RUN(begin_run)
{
    log_info("begin_run: ctx=%p", context);
    auto jRun = nlohmann::json::parse(descriptor_json);
    auto ctx = reinterpret_cast<Context *>(context);
    reload_py(ctx);
    ctx->py_begin_run(descriptor_json);
    ctx->eventBuffers.clear();
    ctx->eventBuffers.resize(jRun["events"].size());

    for (size_t ei = 0; ei < jRun["events"].size(); ++ei)
    {
        auto &jEvent = jRun["events"][ei];
        ctx->eventBuffers[ei].resize(jEvent["outputs"].size());
    }

    // TODO: check the above for errors and return a proper status
    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)
{
    log_info("end: ctx=%p", context);
    auto ctx = reinterpret_cast<Context *>(context);
    ctx->py_end_run(descriptor_json);

    // TODO: check the above for errors and return a proper status
    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);

    auto ctx = reinterpret_cast<Context *>(context);
    auto &buffers = ctx->eventBuffers.at(eventIndex);
    assert(buffers.size() == arrayCount);
    for (size_t ai = 0; ai < arrayCount; ++ai)
    {
        auto arraySpan = mana::get_span<float>(arrays[ai]);
        ctx->eventBuffers[eventIndex][ai] =
            py::memoryview::from_memory(arraySpan.data(), arraySpan.size() * sizeof(float));
    }
    ctx->py_process_event(eventIndex, ctx->eventBuffers);
}

MANA_DEFINE_PLUGIN_SYSTEM_EVENT(process_system_event)
{
    log_debug("system_event: ctx=%p, size=%zu", context, size);
    auto ctx = reinterpret_cast<Context *>(context);
    auto view = py::memoryview::from_memory(data, size * sizeof(uint32_t));
    ctx->py_process_system_event(view);
}

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;
}
