dmf-studio-rnd/shared/VideoReader.hpp

#pragma once

extern "C" {
    #include <libavformat/avformat.h>
    #include <libavcodec/avcodec.h>
    #include <libavutil/avutil.h>
    #include <libswscale/swscale.h>
    #include <libavutil/imgutils.h>
    #include <libavutil/pixdesc.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;
            int stride  = 0;
            AVPixelFormat pix_fmt{};
        };

        struct AudioInfo {
            int sample_rate    = 0;
            int channels       = 0;
            int samples        = 0;
            int channel_stride = 0;  // floats between channel planes
        };

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

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

        ~VideoReader() {
            avcodec_free_context(&codec_context);
            avformat_close_input(&format_context);
            sws_freeContext(sws_ctx);
            av_frame_unref(frame);
            av_packet_unref(packet);
        }

        bool get_next_frame(uint8_t* video_buf, uint8_t* audiobuf) {
            while (dmf::g_running.load(std::memory_order_relaxed)) {
                // Try to get a buffered frame from previous packet 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],
                        v210_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 — seek back to start and keep going
                    avformat_seek_file(format_context, -1, 0, 0, 0, AVSEEK_FLAG_BACKWARD);
                    avcodec_flush_buffers(codec_context);
                    av_packet_unref(packet);
                    continue;
                }
                if (packet->stream_index != video_stream_index) continue;
                avcodec_send_packet(codec_context, packet);
            }
        }

    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;

        // conversion data
        struct SwsContext *sws_ctx{};
        uint8_t* p10_buffer = nullptr;
        int p10_linesizes[4] = {0, 0, 0, 0};
        uint8_t* p10_data[4] = {nullptr, nullptr, nullptr, nullptr};
        uint8_t* v210_buffer = nullptr;
        int v210_stride = 0;

        bool open_file(std::string filename) {
            if (avformat_open_input(&format_context, filename.c_str(), nullptr, nullptr) != 0) {
                throw std::runtime_error("Could not open file: " + filename);
                return false;
            }
            
            // Find stream info
            if (avformat_find_stream_info(format_context, nullptr) < 0) {
                avformat_close_input(&format_context);
                throw std::runtime_error("Could not find stream info");
                return false;
            }
            
            // Find streams
            for (unsigned int i = 0; i < format_context->nb_streams; i++) {
                AVMediaType data_type = format_context->streams[i]->codecpar->codec_type;
                if (data_type == AVMEDIA_TYPE_VIDEO) {
                    video_stream_index = i;
                    have_video = true;
                } else if (data_type == AVMEDIA_TYPE_AUDIO && audio_stream_index != -1) {
                    // TODO: show list of available audio tracks and allow user to pick
                    // or handle multiple audio streams
                    audio_stream_index = 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");
                return false;
            }

            return true;
        }

        void get_source_info() {
            // Get codec parameters
            AVCodecParameters* codec_params = format_context->streams[video_stream_index]->codecpar;
            const AVCodec* codec = avcodec_find_decoder(codec_params->codec_id);
            
            if (!codec) {
                avformat_close_input(&format_context);
                throw std::runtime_error("Unsupported codec");
            }
            
            // Open codec
            codec_context = avcodec_alloc_context3(codec);
            avcodec_parameters_to_context(codec_context, codec_params);
            
            if (avcodec_open2(codec_context, codec, nullptr) < 0) {
                avcodec_free_context(&codec_context);
                avformat_close_input(&format_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()
        {
            // Create Sws context to convert to planar YUV 4:2:2, 20bpp, (1 Cr & Cb sample per 2x1 Y samples), LE
            sws_ctx = sws_getContext(
                source_info.width, source_info.height, source_info.pix_fmt,
                source_info.width, source_info.height, AV_PIX_FMT_YUV422P10LE,
                0, NULL, NULL, NULL
            );

            if (!sws_ctx) {
                throw std::runtime_error("Failed to create SwsContext");
                return;
            }

            // Allocate P10 image
            int p10_buf_size = av_image_get_buffer_size(
                AV_PIX_FMT_YUV422P10LE,
                source_info.width, source_info.height,
                64
            );
            p10_buffer = (uint8_t*)av_malloc(p10_buf_size);
            av_image_alloc(
                p10_data, p10_linesizes,
                source_info.width, source_info.height,
                AV_PIX_FMT_YUV422P10LE, 64
            ); 

            // Allocate v210 buffer
            int blocks_per_row = (codec_context->width + 5) / 6;
            int v210_bytes_per_row = blocks_per_row * 16;
            v210_stride = ((v210_bytes_per_row + 63) / 64) * 64;
            int v210_buffer_size = v210_stride * source_info.height;

            v210_buffer = (uint8_t*)av_malloc(v210_buffer_size);
            if (!v210_buffer) {
                throw std::runtime_error("Failed to allocate V210 buffer");
                return;
            }
        }
        
};
}