dmf-studio-rnd/shared/NDIReceiver.hpp

#pragma once

#include <stdexcept>
#include <string>
#include <vector>
#include <Processing.NDI.Lib.h>
#include "Signal.hpp"
#include "V210.hpp"

namespace dmf {

class NDIReceiver {
public:
    struct SourceInfo {
        int width   = 0;
        int height  = 0;
        int fps_num = 0;
        int fps_den = 0;
        int stride  = 0;
        NDIlib_FourCC_type_e fourcc{};
    };

    NDIReceiver() {
        if (!NDIlib_is_supported_CPU())
            throw std::runtime_error("CPU is not sufficient for NDI");
        if (!NDIlib_initialize())
            throw std::runtime_error("NDI lib init failed");
    }

    ~NDIReceiver() {
        if (recv_) NDIlib_recv_destroy(recv_);
        if (find_) NDIlib_find_destroy(find_);
        NDIlib_destroy();
    }

    // Discovers NDI sources on the network. Polls up to max_attempts times,
    // each waiting timeout_ms milliseconds. Respects g_running.
    std::vector<std::string> find_sources(uint32_t timeout_ms, int max_attempts = 10) {
        find_ = NDIlib_find_create_v2();
        if (!find_)
            throw std::runtime_error("Cannot create NDI finder");

        uint32_t count = 0;
        const NDIlib_source_t* p_sources = nullptr;
        for (int i = 0; !count && i < max_attempts; ++i) {
            if (!dmf::g_running.load(std::memory_order_relaxed)) {
                NDIlib_find_destroy(find_);
                find_ = nullptr;
                throw std::runtime_error("Interrupted while searching for NDI sources");
            }
            NDIlib_find_wait_for_sources(find_, timeout_ms);
            p_sources = NDIlib_find_get_current_sources(find_, &count);
        }
        if (!count) {
            NDIlib_find_destroy(find_);
            find_ = nullptr;
            throw std::runtime_error("No NDI sources found after timeout");
        }

        // Copy while finder is alive: p_sources points into finder-owned memory
        std::vector<std::string> names;
        for (uint32_t i = 0; i < count; ++i) {
            names.emplace_back(p_sources[i].p_ndi_name);
            source_names_.emplace_back(p_sources[i].p_ndi_name);
            source_urls_.emplace_back(p_sources[i].p_url_address);
        }
        sources_.reserve(source_names_.size());
        for (size_t i = 0; i < source_names_.size(); ++i)
            sources_.push_back({source_names_[i].c_str(), source_urls_[i].c_str()});

        NDIlib_find_destroy(find_);
        find_ = nullptr;
        return names;
    }

    // Connects to a discovered source by index.
    void connect(uint32_t source_num) {
        if (sources_.empty())
            throw std::runtime_error("No sources available — call find_sources() first");
        if (source_num >= static_cast<uint32_t>(sources_.size()))
            throw std::runtime_error("source_num exceeds available source count");

        recv_ = NDIlib_recv_create_v3();
        if (!recv_)
            throw std::runtime_error("Cannot create NDI receive instance");
        NDIlib_recv_connect(recv_, &sources_[source_num]);
    }

    // Captures the first video frame to determine resolution, frame rate, and format.
    // Stores the result internally for use by capture_v210(). Respects g_running.
    SourceInfo probe() {
        while (dmf::g_running.load(std::memory_order_relaxed)) {
            NDIlib_video_frame_v2_t frame;
            auto type = NDIlib_recv_capture_v3(recv_, &frame, nullptr, nullptr, 1000);
            if (type == NDIlib_frame_type_error)
                throw std::runtime_error("NDI source lost during probe");
            if (type != NDIlib_frame_type_video)
                continue;

            info_.width   = frame.xres;
            info_.height  = frame.yres;
            info_.fps_num = frame.frame_rate_N;
            info_.fps_den = frame.frame_rate_D;
            info_.fourcc  = static_cast<NDIlib_FourCC_type_e>(frame.FourCC);
            info_.stride  = frame.line_stride_in_bytes;
            if (info_.stride == 0)
                info_.stride = info_.width * bytes_per_pixel(info_.fourcc);

            NDIlib_recv_free_video_v2(recv_, &frame);
            return info_;
        }
        throw std::runtime_error("Interrupted during probe");
    }

    enum class FrameKind { None, Video, Audio };

    struct AudioInfo {
        int sample_rate    = 0;
        int channels       = 0;
        int samples        = 0;
        int channel_stride = 0;  // floats between channel planes (NDI planar layout)
    };

    // Waits for the first audio packet and returns its format.
    // Call after probe(). Audio packets are typically already queued at that point.
    // Discards any video frames encountered while searching. Respects g_running.
    AudioInfo probe_audio(uint32_t timeout_ms = 5000) {
        const auto deadline = std::chrono::steady_clock::now()
                              + std::chrono::milliseconds(timeout_ms);
        while (dmf::g_running.load(std::memory_order_relaxed) &&
               std::chrono::steady_clock::now() < deadline) {
            NDIlib_video_frame_v2_t video_frame{};
            NDIlib_audio_frame_v3_t audio_frame{};
            auto type = NDIlib_recv_capture_v3(recv_, &video_frame, &audio_frame, nullptr, 100);
            if (type == NDIlib_frame_type_error)
                throw std::runtime_error("NDI source lost during audio probe");
            if (type == NDIlib_frame_type_video) {
                NDIlib_recv_free_video_v2(recv_, &video_frame);
                continue;
            }
            if (type == NDIlib_frame_type_audio) {
                const bool is_fltp = (audio_frame.FourCC == NDIlib_FourCC_audio_type_FLTP);
                AudioInfo info;
                if (is_fltp) {
                    info.sample_rate    = audio_frame.sample_rate;
                    info.channels       = audio_frame.no_channels;
                    info.samples        = audio_frame.no_samples;
                    info.channel_stride = audio_frame.channel_stride_in_bytes / sizeof(float);
                }
                NDIlib_recv_free_audio_v3(recv_, &audio_frame);
                if (is_fltp) return info;
            }
        }
        throw std::runtime_error("Timeout waiting for first audio packet");
    }

    // Receives one NDI frame. On video: converts to V210 in frame_buffer/frame_stride.
    // On audio: copies float32 planar samples into audio_out and fills audio_info.
    // Returns FrameKind::None on timeout or non-A/V frames.
    // Throws on source lost or unsupported video format.
    // timeout_ms=0 is non-blocking (useful for draining queued frames after an initial wait).
    FrameKind capture(uint8_t* frame_buffer, uint32_t frame_stride,
                      std::vector<float>& audio_out, AudioInfo& audio_info,
                      bool want_audio = false, int timeout_ms = 5) {
        NDIlib_video_frame_v2_t video_frame{};
        NDIlib_audio_frame_v3_t audio_frame{};
        auto type = NDIlib_recv_capture_v3(recv_, &video_frame,
                                           want_audio ? &audio_frame : nullptr,
                                           nullptr, timeout_ms);

        if (type == NDIlib_frame_type_error)
            throw std::runtime_error("NDI source lost");

        if (type == NDIlib_frame_type_status_change) {
            // Don't call probe() here — it would consume a video frame from the queue.
            // Format changes are detected on the next actual video frame below.
            return FrameKind::None;
        }

        if (type == NDIlib_frame_type_video) {
            // Use the actual frame's FourCC and stride, not the stale values from probe().
            const auto fourcc  = static_cast<NDIlib_FourCC_type_e>(video_frame.FourCC);
            const int   stride = video_frame.line_stride_in_bytes > 0
                                     ? video_frame.line_stride_in_bytes
                                     : video_frame.xres * bytes_per_pixel(fourcc);
            switch (fourcc) {
                case NDIlib_FourCC_type_UYVY:
                    v210::UYVYtoV210(video_frame.p_data, frame_buffer,
                                     video_frame.xres, video_frame.yres, stride, frame_stride);
                    break;
                case NDIlib_FourCC_type_P216: {
                    NDIlib_video_frame_v2_t dst{};
                    dst.p_data = frame_buffer;
                    dst.line_stride_in_bytes = frame_stride;
                    NDIlib_util_P216_to_V210(&video_frame, &dst);
                    break;
                }
                default:
                    NDIlib_recv_free_video_v2(recv_, &video_frame);
                    throw std::runtime_error("Unsupported NDI color format: " + fourcc_str(fourcc));
            }
            NDIlib_recv_free_video_v2(recv_, &video_frame);
            return FrameKind::Video;
        }

        if (type == NDIlib_frame_type_audio) {
            if (audio_frame.FourCC != NDIlib_FourCC_audio_type_FLTP) {
                NDIlib_recv_free_audio_v3(recv_, &audio_frame);
                return FrameKind::None;  // compressed or unknown format — skip
            }
            audio_info.sample_rate    = audio_frame.sample_rate;
            audio_info.channels       = audio_frame.no_channels;
            audio_info.samples        = audio_frame.no_samples;
            audio_info.channel_stride = audio_frame.channel_stride_in_bytes / sizeof(float);
            const int total = audio_frame.no_channels * audio_info.channel_stride;
            const auto* fdata = reinterpret_cast<const float*>(audio_frame.p_data);
            audio_out.assign(fdata, fdata + total);
            NDIlib_recv_free_audio_v3(recv_, &audio_frame);
            return FrameKind::Audio;
        }

        return FrameKind::None;
    }

private:
    NDIlib_find_instance_t       find_ = nullptr;
    NDIlib_recv_instance_t       recv_ = nullptr;
    SourceInfo                   info_;
    std::vector<std::string>     source_names_;
    std::vector<std::string>     source_urls_;
    std::vector<NDIlib_source_t> sources_;

    static int bytes_per_pixel(NDIlib_FourCC_type_e fc) {
        switch (fc) {
            case NDIlib_FourCC_video_type_UYVY:
            case NDIlib_FourCC_video_type_YV12:
            case NDIlib_FourCC_video_type_I420:
            case NDIlib_FourCC_video_type_NV12:  return 2;
            case NDIlib_FourCC_video_type_BGRA:
            case NDIlib_FourCC_video_type_RGBA:
            case NDIlib_FourCC_video_type_BGRX:
            case NDIlib_FourCC_video_type_RGBX:  return 4;
            case NDIlib_FourCC_video_type_UYVA:  return 3;
            case NDIlib_FourCC_video_type_P216:  return 4;
            case NDIlib_FourCC_video_type_PA16:  return 6;
            default:                              return 2;
        }
    }

    static std::string fourcc_str(NDIlib_FourCC_type_e fc) {
        uint32_t v = static_cast<uint32_t>(fc);
        char s[5] = {
            static_cast<char>((v >>  0) & 0xFF),
            static_cast<char>((v >>  8) & 0xFF),
            static_cast<char>((v >> 16) & 0xFF),
            static_cast<char>((v >> 24) & 0xFF),
            '\0'
        };
        return s;
    }
};

} // namespace dmf