dmf-studio-rnd/shared/DeckLinkReceiver.hpp

#pragma once

#include <algorithm>
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstring>
#include <mutex>
#include <numeric>
#include <stdexcept>
#include <string>
#include <vector>
#include <DeckLinkAPI.h>
#include "Signal.hpp"

namespace dmf {

class DeckLinkReceiver {
public:
    struct VideoInfo {
        int width   = 0;
        int height  = 0;
        int fps_num = 0;
        int fps_den = 0;
    };

    struct DeviceInfo {
        uint32_t    index;
        std::string name;
    };

    std::vector<DeviceInfo> devices;
    VideoInfo               video_info{};

    DeckLinkReceiver() { enumerate_devices(); }

    ~DeckLinkReceiver() {
        if (decklink_input) {
            decklink_input->StopStreams();
            decklink_input->DisableVideoInput();
            decklink_input->SetCallback(nullptr);
            decklink_input->Release();
        }
        delete input_callback;
        for (auto* d : raw_devices) d->Release();
        if (selected_device) selected_device->Release();
    }

    void start_capture(uint32_t device_index) {
        if (device_index >= raw_devices.size())
            throw std::runtime_error("Device index out of range");

        selected_device = raw_devices[device_index];
        selected_device->AddRef();

        HRESULT r = selected_device->QueryInterface(IID_IDeckLinkInput, (void**)&decklink_input);
        if (r != S_OK) throw std::runtime_error("Could not obtain IDeckLinkInput");

        input_callback = new InputCallback(*this);
        r = decklink_input->SetCallback(input_callback);
        if (r != S_OK) throw std::runtime_error("Could not set input callback");

        r = decklink_input->EnableVideoInput(bmdModeNTSC, bmdFormat10BitYUV,
                                              bmdVideoInputEnableFormatDetection);
        if (r != S_OK) throw std::runtime_error("Could not enable video input");

        r = decklink_input->StartStreams();
        if (r != S_OK) throw std::runtime_error("Could not start streams");
    }

    bool wait_for_format(uint64_t timeout_ms = 5000) {
        std::unique_lock<std::mutex> lk(mutex);
        return format_cv.wait_for(lk, std::chrono::milliseconds(timeout_ms),
            [this] { return format_detected; });
    }

    // Blocks until a fresh frame arrives or g_running goes false.
    // Copies frame data into dst with MXL stride. Returns false on shutdown.
    bool wait_for_frame(uint8_t* dst, uint32_t dst_stride, int width, int height) {
        std::unique_lock<std::mutex> lk(mutex);
        frame_cv.wait(lk, [this] {
            return frame_ready || !dmf::g_running.load(std::memory_order_relaxed);
        });
        if (!dmf::g_running.load(std::memory_order_relaxed)) return false;

        const uint32_t src_stride  = frame_row_bytes;
        const int      rows        = std::min(frame_height, height);
        const uint32_t copy_bytes  = std::min(src_stride, dst_stride);

        std::memset(dst, 0, static_cast<size_t>(dst_stride) * static_cast<size_t>(height));
        const uint8_t* src = frame_buffer.data();
        uint8_t*       d   = dst;
        for (int y = 0; y < rows; ++y, src += src_stride, d += dst_stride)
            std::memcpy(d, src, copy_bytes);

        frame_ready = false;
        return true;
    }

private:
    class InputCallback : public IDeckLinkInputCallback {
    public:
        explicit InputCallback(DeckLinkReceiver& owner) : owner(owner) {}

        HRESULT STDMETHODCALLTYPE VideoInputFormatChanged(
            BMDVideoInputFormatChangedEvents events,
            IDeckLinkDisplayMode*            mode,
            BMDDetectedVideoInputFormatFlags /*flags*/) override
        {
            if (!(events & bmdVideoInputDisplayModeChanged)) return S_OK;

            {
                std::lock_guard<std::mutex> lk(owner.mutex);
                owner.video_info.width  = static_cast<int>(mode->GetWidth());
                owner.video_info.height = static_cast<int>(mode->GetHeight());
                BMDTimeValue dur = 0; BMDTimeScale ts = 0;
                mode->GetFrameRate(&dur, &ts);
                if (dur > 0 && ts > 0) {
                    owner.video_info.fps_num = static_cast<int>(ts);
                    owner.video_info.fps_den = static_cast<int>(dur);
                    const int g = std::gcd(owner.video_info.fps_num, owner.video_info.fps_den);
                    if (g > 1) { owner.video_info.fps_num /= g; owner.video_info.fps_den /= g; }
                }
                owner.frame_ready    = false;
                owner.frame_buffer.clear();
                owner.format_detected = true;
            }
            owner.format_cv.notify_one();

            HRESULT r = owner.decklink_input->StopStreams();
            if (r == S_OK) {
                r = owner.decklink_input->EnableVideoInput(
                    mode->GetDisplayMode(), bmdFormat10BitYUV,
                    bmdVideoInputEnableFormatDetection);
                if (r == S_OK) owner.decklink_input->StartStreams();
            }
            return S_OK;
        }

        HRESULT STDMETHODCALLTYPE VideoInputFrameArrived(
            IDeckLinkVideoInputFrame*  video_frame,
            IDeckLinkAudioInputPacket* /*audio_packet*/) override
        {
            if (!video_frame) return S_OK;
            void* src = nullptr;
            video_frame->GetBytes(&src);
            if (!src) return S_OK;

            const uint32_t row_bytes = static_cast<uint32_t>(video_frame->GetRowBytes());
            const int      fw        = static_cast<int>(video_frame->GetWidth());
            const int      fh        = static_cast<int>(video_frame->GetHeight());
            const size_t   sz        = static_cast<size_t>(row_bytes) * static_cast<size_t>(fh);

            {
                std::lock_guard<std::mutex> lk(owner.mutex);
                if (owner.frame_buffer.size() < sz) owner.frame_buffer.resize(sz);
                std::memcpy(owner.frame_buffer.data(), src, sz);
                owner.frame_row_bytes = row_bytes;
                owner.frame_width     = fw;
                owner.frame_height    = fh;
                owner.frame_ready     = true;
            }
            owner.frame_cv.notify_one();
            return S_OK;
        }

        HRESULT STDMETHODCALLTYPE QueryInterface(REFIID, LPVOID*) override { return E_NOINTERFACE; }
        ULONG   STDMETHODCALLTYPE AddRef()  override { return ++ref_count; }
        ULONG   STDMETHODCALLTYPE Release() override { return --ref_count; }

    private:
        DeckLinkReceiver&    owner;
        std::atomic<int32_t> ref_count{1};
    };

    std::vector<IDeckLink*>  raw_devices;
    IDeckLink*               selected_device = nullptr;
    IDeckLinkInput*          decklink_input  = nullptr;
    InputCallback*           input_callback  = nullptr;

    std::mutex              mutex;
    std::condition_variable format_cv;
    std::condition_variable frame_cv;
    bool                    format_detected = false;

    std::vector<uint8_t>    frame_buffer;
    uint32_t                frame_row_bytes = 0;
    int                     frame_width     = 0;
    int                     frame_height    = 0;
    bool                    frame_ready     = false;

    void enumerate_devices() {
        IDeckLinkIterator* it = CreateDeckLinkIteratorInstance();
        if (!it) throw std::runtime_error("DeckLink drivers not installed");

        IDeckLink* device = nullptr;
        uint32_t   index  = 0;
        while (it->Next(&device) == S_OK) {
            IDeckLinkInput* inp = nullptr;
            if (device->QueryInterface(IID_IDeckLinkInput, (void**)&inp) == S_OK) {
                inp->Release();
                const char* name = nullptr;
                device->GetDisplayName(&name);
                devices.push_back({index, name ? name : "?"});
                raw_devices.push_back(device);
                index++;
            } else {
                device->Release();
            }
        }
        it->Release();
        if (raw_devices.empty())
            throw std::runtime_error("No DeckLink input devices found");
    }
};

} // namespace dmf