#include <algorithm>
#include <cstring>
#include <string>
#include <vector>
#include <mxl/flow.h>
#include <mxl/time.h>
#include "NodeBase.hpp"
#include "FlowDef.hpp"
#include "NDIReceiver.hpp"

class NDIInNode : public dmf::NodeBase {
    void run() override {
        const auto source_num = static_cast<uint32_t>(config().value("source_num", 0));

        dmf::NDIReceiver ndi;
        dmf::NDIReceiver::SourceInfo src;
        try {
            auto sources = ndi.find_sources(5000);
            log("Available NDI sources:");
            for (const auto& name : sources)
                log("  %s", name.c_str());
            ndi.connect(source_num);
            src = ndi.probe();
        } catch (const std::runtime_error& e) {
            log("Error: %s", e.what());
            return;
        }

        // --- video flow ---
        const auto video_flow_info = config().at("video_flow_id");
        const auto video_flow_id   = video_flow_info.at("id").get<std::string>();
        const int  width           = video_flow_info.value("width",   src.width);
        const int  height          = video_flow_info.value("height",  src.height);
        const int  fps_num         = video_flow_info.value("fps_num", src.fps_num);
        const int  fps_den         = video_flow_info.value("fps_den", src.fps_den);

        log("video flow=%s  %dx%d @ %d/%d fps", video_flow_id.c_str(), width, height, fps_num, fps_den);

        mxlFlowWriter     video_writer{};
        mxlFlowConfigInfo video_cfg{};
        bool created = false;
        mxlStatus st = mxlCreateFlowWriter(
            instance(),
            dmf::make_video_flow_def(video_flow_id, node_id(), width, height, fps_num, fps_den).c_str(),
            nullptr, &video_writer, &video_cfg, &created);
        if (st != MXL_STATUS_OK) { log("video mxlCreateFlowWriter failed (status=%d)", st); return; }

        const uint32_t video_stride = video_cfg.discrete.sliceSizes[0];
        log("video stride=%u B/line  grain=%u B  ring=%u grains",
            video_stride, video_stride * static_cast<uint32_t>(height), video_cfg.discrete.grainCount);

        // --- audio flow (optional — only created when graph wires audio_flow_id) ---
        mxlFlowWriter audio_writer{};
        mxlFlowConfigInfo audio_cfg{};
        int  sample_rate  = 0;
        int  channels     = 0;
        int  bit_depth    = 32;
        bool has_audio    = config().contains("audio_flow_id");

        if (has_audio) {
            const auto audio_flow_info = config().at("audio_flow_id");
            const auto audio_flow_id   = audio_flow_info.at("id").get<std::string>();
            sample_rate = audio_flow_info.value("sample_rate", 48000);
            channels    = audio_flow_info.value("channels", 2);
            bit_depth   = audio_flow_info.value("bit_depth", 32);

            log("audio flow=%s  %d Hz  %dch  %d-bit", audio_flow_id.c_str(), sample_rate, channels, bit_depth);

            mxlStatus ast = mxlCreateFlowWriter(
                instance(),
                dmf::make_audio_flow_def(audio_flow_id, node_id(), sample_rate, channels, bit_depth,
                                         fps_num, fps_den).c_str(),
                nullptr, &audio_writer, &audio_cfg, &created);
            if (ast != MXL_STATUS_OK) {
                log("audio mxlCreateFlowWriter failed (status=%d) — continuing without audio", ast);
                has_audio = false;
            } else {
                log("audio grain=%u B", audio_cfg.discrete.sliceSizes[0]);
            }
        }

        // --- main loop ---
        const mxlRational video_rate = {fps_num, fps_den};
        const mxlRational audio_rate = {sample_rate, 1};

        uint64_t video_index = mxlGetCurrentIndex(&video_rate);
        uint64_t audio_index = has_audio ? mxlGetCurrentIndex(&audio_rate) : 0;
        log("start video_index=%llu", video_index);

        std::vector<uint8_t> latest_video(video_stride * height);
        bool have_video = false;
        std::vector<float> audio_buf;
        dmf::NDIReceiver::AudioInfo audio_info;

        while (dmf::g_running.load(std::memory_order_relaxed)) {
            dmf::NDIReceiver::FrameKind kind;
            try {
                kind = ndi.capture(latest_video.data(), video_stride, audio_buf, audio_info);
            } catch (const std::runtime_error& e) {
                log("NDI error: %s — stopping", e.what());
                break;
            }

            if (kind == dmf::NDIReceiver::FrameKind::Video) {
                have_video = true;
            } else if (kind == dmf::NDIReceiver::FrameKind::Audio && has_audio) {
                mxlGrainInfo grain{};
                uint8_t* buf = nullptr;
                if (mxlFlowWriterOpenGrain(audio_writer, audio_index, &grain, &buf) == MXL_STATUS_OK) {
                    // Convert float32 planar → int32 interleaved PCM
                    const int audio_grain_size = audio_cfg.discrete.sliceSizes[0];
                    const int capacity = audio_grain_size / (channels * static_cast<int>(sizeof(int32_t)));
                    const int to_write = std::min(audio_info.samples, capacity);
                    auto* out = reinterpret_cast<int32_t*>(buf);
                    for (int s = 0; s < to_write; ++s)
                        for (int c = 0; c < audio_info.channels; ++c) {
                            float f = audio_buf[c * audio_info.channel_stride + s];
                            f = std::max(-1.0f, std::min(1.0f, f));
                            *out++ = static_cast<int32_t>(f * 2147483647.0f);
                        }
                    grain.validSlices = grain.totalSlices;
                    mxlFlowWriterCommitGrain(audio_writer, &grain);
                }
                audio_index += audio_info.samples;
            }

            // Write video grain whenever the MXL clock has reached video_index
            const uint64_t current = mxlGetCurrentIndex(&video_rate);
            if (current >= video_index) {
                mxlGrainInfo grain{};
                uint8_t* buf = nullptr;
                st = mxlFlowWriterOpenGrain(video_writer, video_index, &grain, &buf);
                if (st == MXL_STATUS_OK) {
                    if (have_video) {
                        std::memcpy(buf, latest_video.data(), latest_video.size());
                        grain.flags = 0;
                    } else {
                        grain.flags = MXL_GRAIN_FLAG_INVALID;
                    }
                    grain.validSlices = grain.totalSlices;
                    mxlFlowWriterCommitGrain(video_writer, &grain);
                } else {
                    log("video OpenGrain failed (status=%d) at index=%llu", st, video_index);
                }
                video_index = current + 1;
            }
        }

        log("stopped at video_index=%llu", video_index);
        mxlReleaseFlowWriter(instance(), video_writer);
        if (has_audio) mxlReleaseFlowWriter(instance(), audio_writer);
    }
};

int main() {
    NDIInNode node;
    return node.execute();
}