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 <libswresample/swresample.h>
    #include <libavutil/opt.h>
}

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

namespace dmf {

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

    struct AudioInfo {
        int sample_rate = 0;
        int channels    = 0;
    };

    enum class FrameKind { None, Video, Audio };

    VideoInfo video_info{};
    AudioInfo audio_info{};
    bool has_audio  = false;
    bool has_video = false;

    explicit VideoReader(const std::string& filename) {
        if (!open_file(filename))
            return;
        if (!has_video && !has_audio) return;
        get_source_info();
        if (has_video) allocate_video_conversion_buffers();
        if (has_audio)  allocate_audio_conversion_buffers();
    }

    ~VideoReader() {
        avcodec_free_context(&video_codec_context);
        avcodec_free_context(&audio_codec_context);
        avformat_close_input(&format_context);
        sws_freeContext(sws_video_ctx);
        swr_free(&swr_audio_ctx);
        av_freep(&p10_data[0]);
        av_frame_free(&video_frame);
        av_frame_free(&audio_frame);
        av_packet_free(&packet);
    }

    // Returns Video or Audio when a frame/packet was decoded, None when g_running goes false.
    FrameKind read_next(uint8_t* video_buf, uint32_t mxl_stride, uint8_t* audio_buf, int max_audio_samples, int& out_samples_written) {
        while (dmf::g_running.load(std::memory_order_relaxed)) {
            // Drain any frames buffered in the decoder first
            if (has_video && avcodec_receive_frame(video_codec_context, video_frame) == 0) {
                if (!video_buf) {
                    av_frame_unref(video_frame);
                    continue;  // nowhere to write — discard frame
                }
                sws_scale(
                    sws_video_ctx,
                    video_frame->data,
                    video_frame->linesize,
                    0,
                    video_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,
                    video_info.width,
                    video_info.height,
                    p10_linesizes[0],
                    p10_linesizes[1],
                    p10_linesizes[2],
                    mxl_stride
                );
                av_frame_unref(video_frame);
                return FrameKind::Video;
            }

            if (has_audio && avcodec_receive_frame(audio_codec_context, audio_frame) == 0) {
                if (!audio_buf) {
                    av_frame_unref(audio_frame);
                    continue;  // nowhere to write — discard frame
                }

                int dst_nb_samples = av_rescale_rnd(
                    swr_get_delay(swr_audio_ctx, audio_codec_context->sample_rate) + audio_frame->nb_samples,
                    audio_codec_context->sample_rate, audio_codec_context->sample_rate, AV_ROUND_UP
                );

                // Guard against buffer overflows
                if (dst_nb_samples > max_audio_samples) {
                    dst_nb_samples = max_audio_samples;
                }

                std::vector<uint8_t*> dst(audio_info.channels);
                for (int ch = 0; ch < audio_info.channels; ch++) {
                    dst[ch] = audio_buf + ch * max_audio_samples * sizeof(float);
                }

                // Convert/Resample the audio layout and sample format
                int converted_samples = swr_convert(
                    swr_audio_ctx,
                    dst.data(), dst_nb_samples,
                    (const uint8_t**)audio_frame->data, audio_frame->nb_samples
                );
                out_samples_written = converted_samples;

                av_frame_unref(audio_frame);
                return FrameKind::Audio;
            }

            // 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);
                if (has_video) avcodec_flush_buffers(video_codec_context);
                if (has_audio) avcodec_flush_buffers(audio_codec_context);
                swr_close(swr_audio_ctx);
                swr_init(swr_audio_ctx);
                continue;
            }

            if (packet->stream_index == video_stream_index) {
                avcodec_send_packet(video_codec_context, packet);
            } else if (packet->stream_index == audio_stream_index) {
                avcodec_send_packet(audio_codec_context, packet);
            }
        }
        return FrameKind::None;
    }

private:
    int video_stream_index = -1;
    int audio_stream_index = -1;

    AVFormatContext* format_context        = nullptr;
    AVPacket*        packet                = av_packet_alloc();

    AVCodecContext*  video_codec_context   = nullptr;
    AVFrame*         video_frame           = av_frame_alloc();
    SwsContext* sws_video_ctx         = nullptr;
    int         p10_linesizes[4] = {0, 0, 0, 0};
    uint8_t*    p10_data[4]      = {nullptr, nullptr, nullptr, nullptr};

    AVCodecContext* audio_codec_context = nullptr;
    AVFrame*        audio_frame        = av_frame_alloc();
    SwrContext* swr_audio_ctx = 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);
                has_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() {
        if (video_stream_index != -1) {
            AVCodecParameters* video_codec_params = format_context->streams[video_stream_index]->codecpar;
            const AVCodec* video_codec = avcodec_find_decoder(video_codec_params->codec_id);
            if (!video_codec)
                throw std::runtime_error("Unsupported video codec");

            video_codec_context = avcodec_alloc_context3(video_codec);
            if (avcodec_parameters_to_context(video_codec_context, video_codec_params) < 0)
                throw std::runtime_error("Could not copy video codec parameters");

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

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

            video_info.width   = video_codec_context->width;
            video_info.height  = video_codec_context->height;
            video_info.fps_num = fps.num;
            video_info.fps_den = fps.den;
            video_info.pix_fmt = video_codec_context->pix_fmt;
        } else {
            has_video = false;
        }

        // audio part
        if (audio_stream_index == -1) return;
        AVCodecParameters* audio_codec_params = format_context->streams[audio_stream_index]->codecpar;
        const AVCodec* audio_codec = avcodec_find_decoder(audio_codec_params->codec_id);
        if (!audio_codec)
            throw std::runtime_error("Unsupported audio codec");

        audio_codec_context = avcodec_alloc_context3(audio_codec);
        if (avcodec_parameters_to_context(audio_codec_context, audio_codec_params) < 0)
            throw std::runtime_error("Could not copy audio codec parameters");

        if (avcodec_open2(audio_codec_context, audio_codec, nullptr) < 0) {
            avcodec_free_context(&audio_codec_context);
            throw std::runtime_error("Could not open audio codec");
        }
        audio_info.sample_rate = audio_codec_context->sample_rate;
        audio_info.channels = audio_codec_context->ch_layout.nb_channels;
    }

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

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

    void allocate_audio_conversion_buffers() {
        swr_audio_ctx = swr_alloc();

        // Set input options
        av_opt_set_chlayout(swr_audio_ctx, "in_chlayout",    &audio_codec_context->ch_layout, 0);
        av_opt_set_int(swr_audio_ctx,      "in_sample_rate", audio_info.sample_rate, 0);
        av_opt_set_sample_fmt(swr_audio_ctx, "in_sample_fmt", audio_codec_context->sample_fmt, 0);

        // Set output options
        av_opt_set_chlayout(swr_audio_ctx, "out_chlayout",   &audio_codec_context->ch_layout, 0);
        av_opt_set_int(swr_audio_ctx,      "out_sample_rate", audio_info.sample_rate, 0);
        av_opt_set_sample_fmt(swr_audio_ctx, "out_sample_fmt", AV_SAMPLE_FMT_FLTP, 0);

        // Initialize the context
        if (swr_init(swr_audio_ctx) < 0) {
            throw std::runtime_error("Failed to create SwrContext");
        }
    }
};

} // namespace dmf