dmf-studio-rnd/shared/VideoReader.hpp

#pragma once

extern "C" {
    #include <libavformat/avformat.h>
    #include <libavcodec/avcodec.h>
    #include <libavutil/avutil.h>
    #include <libavutil/imgutils.h>
    #include <libswscale/swscale.h>
}

#include <stdexcept>
#include <string>
#include "Signal.hpp"
#include "V210.hpp"

namespace dmf {

class VideoReader {
public:
    struct SourceInfo {
        int width   = 0;
        int height  = 0;
        int fps_num = 0;
        int fps_den = 0;
        AVPixelFormat pix_fmt{};
    };

    SourceInfo source_info{};
    bool has_audio  = false;
    bool have_video = false;

    explicit VideoReader(const std::string& filename) {
        if (!open_file(filename))
            return;
        if (have_video) {
            get_source_info();
            allocate_conversion_buffers();
        }
    }

    ~VideoReader() {
        avcodec_free_context(&codec_context);
        avformat_close_input(&format_context);
        sws_freeContext(sws_ctx);
        av_freep(&p10_data[0]);
        av_frame_free(&frame);
        av_packet_free(&packet);
    }

    // Returns true when a frame was decoded and written into video_buf.
    // Returns false when g_running goes false.
    bool get_next_frame(uint8_t* video_buf, uint32_t mxl_stride) {
        while (dmf::g_running.load(std::memory_order_relaxed)) {
            // Drain any frames buffered in the decoder first
            if (avcodec_receive_frame(codec_context, frame) == 0) {
                sws_scale(
                    sws_ctx,
                    frame->data,
                    frame->linesize,
                    0,
                    source_info.height,
                    p10_data,
                    p10_linesizes
                );
                dmf::v210::YUV422P10toV210(
                    reinterpret_cast<uint16_t*>(p10_data[0]),
                    reinterpret_cast<uint16_t*>(p10_data[1]),
                    reinterpret_cast<uint16_t*>(p10_data[2]),
                    video_buf,
                    source_info.width,
                    source_info.height,
                    p10_linesizes[0],
                    p10_linesizes[1],
                    p10_linesizes[2],
                    mxl_stride
                );
                av_frame_unref(frame);
                return true;
            }

            // No buffered frame — read next packet
            av_packet_unref(packet);
            if (av_read_frame(format_context, packet) < 0) {
                // EOF — loop back to start
                avformat_seek_file(format_context, -1, 0, 0, 0, AVSEEK_FLAG_BACKWARD);
                avcodec_flush_buffers(codec_context);
                continue;
            }
            if (packet->stream_index != video_stream_index) continue;
            avcodec_send_packet(codec_context, packet);
        }
        return false;
    }

private:
    AVFormatContext* format_context  = nullptr;
    AVCodecContext*  codec_context   = nullptr;
    AVPacket*        packet          = av_packet_alloc();
    AVFrame*         frame           = av_frame_alloc();
    int video_stream_index = -1;
    int audio_stream_index = -1;

    SwsContext* sws_ctx         = nullptr;
    int         p10_linesizes[4] = {0, 0, 0, 0};
    uint8_t*    p10_data[4]      = {nullptr, nullptr, nullptr, nullptr};

    bool open_file(const std::string& filename) {
        if (avformat_open_input(&format_context, filename.c_str(), nullptr, nullptr) != 0)
            throw std::runtime_error("Could not open file: " + filename);

        if (avformat_find_stream_info(format_context, nullptr) < 0) {
            avformat_close_input(&format_context);
            throw std::runtime_error("Could not find stream info");
        }

        for (unsigned int i = 0; i < format_context->nb_streams; ++i) {
            const AVMediaType type = format_context->streams[i]->codecpar->codec_type;
            if (type == AVMEDIA_TYPE_VIDEO && video_stream_index == -1) {
                video_stream_index = static_cast<int>(i);
                have_video = true;
            } else if (type == AVMEDIA_TYPE_AUDIO && audio_stream_index == -1) {
                audio_stream_index = static_cast<int>(i);
                has_audio = true;
            }
        }

        if (video_stream_index == -1 && audio_stream_index == -1) {
            avformat_close_input(&format_context);
            throw std::runtime_error("No audio/video stream found in: " + filename);
        }

        return true;
    }

    void get_source_info() {
        AVCodecParameters* codec_params = format_context->streams[video_stream_index]->codecpar;
        const AVCodec* codec = avcodec_find_decoder(codec_params->codec_id);
        if (!codec)
            throw std::runtime_error("Unsupported codec");

        codec_context = avcodec_alloc_context3(codec);
        if (avcodec_parameters_to_context(codec_context, codec_params) < 0)
            throw std::runtime_error("Could not copy codec parameters");

        if (avcodec_open2(codec_context, codec, nullptr) < 0) {
            avcodec_free_context(&codec_context);
            throw std::runtime_error("Could not open codec");
        }

        AVRational fps = codec_context->framerate;
        if (fps.num == 0 || fps.den == 0)
            fps = format_context->streams[video_stream_index]->avg_frame_rate;

        source_info.width   = codec_context->width;
        source_info.height  = codec_context->height;
        source_info.fps_num = fps.num;
        source_info.fps_den = fps.den;
        source_info.pix_fmt = codec_context->pix_fmt;
    }

    void allocate_conversion_buffers() {
        sws_ctx = sws_getContext(
            source_info.width, source_info.height, source_info.pix_fmt,
            source_info.width, source_info.height, AV_PIX_FMT_YUV422P10LE,
            SWS_BILINEAR, nullptr, nullptr, nullptr
        );
        if (!sws_ctx)
            throw std::runtime_error("Failed to create SwsContext");

        if (av_image_alloc(p10_data, p10_linesizes,
                source_info.width, source_info.height,
                AV_PIX_FMT_YUV422P10LE, 64) < 0)
            throw std::runtime_error("Failed to allocate YUV422P10 buffer");
    }
};

} // namespace dmf