ADOFAIRenderer/Util/Audio/Clip/FFmpegPCMReader.cs

using System;
using System.Runtime.InteropServices;
using FFmpeg.AutoGen;

namespace ADOFAIRenderer.Util.Audio.Clip;

/// <summary>
/// <see cref="IPCMReader"/> that decodes the best audio stream of a file with libavformat/libavcodec.
/// Output keeps the stream's sample rate and channel count; libswresample converts only the sample format to interleaved float.
///
/// Frame positions: frame 0 is the first sample the decoder outputs for the stream's start time.
/// The position of the first frame decoded after opening or seeking is taken from its timestamp
/// (best_effort_timestamp - stream start_time, in samples); following frames are counted sequentially.
/// Seek moves the demuxer to <see cref="SeekMarginSeconds"/> before the target and discards decoded samples before the target.
///
/// <see cref="TotalFrames"/> comes from the stream duration (or the container duration) and may differ from the decoded length;
/// reads stop at whichever ends first.
/// Reference: https://ffmpeg.org/doxygen/trunk/group__lavc__decoding.html (send/receive API)
/// Reference: https://ffmpeg.org/doxygen/trunk/group__lavf__decoding.html (av_read_frame, av_seek_frame)
/// </summary>
internal sealed unsafe class FFmpegPCMReader : IPCMReader
{
    /// <summary>
    /// Time the demuxer is moved before a seek target, in seconds. Decoders that depend on previous packets
    /// (for example the MP3 bit reservoir) produce settled output after this margin; the samples in it are discarded.
    /// </summary>
    private const double SeekMarginSeconds = 0.25;

    private AVFormatContext* _format;
    private AVCodecContext* _codec;
    private SwrContext* _swr;
    private AVPacket* _packet;
    private AVFrame* _frame;
    private readonly int _streamIndex;
    private readonly AVRational _timeBase;
    private readonly long _startPts;

    // Decoded samples not yet returned: frames [_pendingStart, _pendingStart + _pendingCount) at _pending[_pendingOffset * Channels].
    private float[] _pending = Array.Empty<float>();
    private int _pendingOffset;
    private int _pendingCount;
    private long _pendingStart;

    private long _nextDecodedFrame; // Frame index of the next decoded frame when counting sequentially
    private bool _syncToTimestamp;  // True until the first frame after open/seek sets _nextDecodedFrame from its timestamp
    private bool _flushSent;
    private bool _endOfStream;

    public string Path { get; }
    public int Channels { get; }
    public int SampleRate { get; }
    public long TotalFrames { get; }
    public bool IsForwardOnly => false;
    public long Position { get; private set; }

    /// <summary>Number of demuxer seeks performed (for diagnostics).</summary>
    public int SeekCount { get; private set; }

    public FFmpegPCMReader(string path)
    {
        Path = path;
        try
        {
            AVFormatContext* format = null;
            var ret = ffmpeg.avformat_open_input(&format, path, null, null);
            if (ret < 0) throw new InvalidOperationException($"Failed to open audio file '{path}': {ErrorString(ret)}");
            _format = format;

            ret = ffmpeg.avformat_find_stream_info(_format, null);
            if (ret < 0) throw new InvalidOperationException($"Failed to read stream info of '{path}': {ErrorString(ret)}");

            AVCodec* decoder = null;
            _streamIndex = ffmpeg.av_find_best_stream(_format, AVMediaType.AVMEDIA_TYPE_AUDIO, -1, -1, &decoder, 0);
            if (_streamIndex < 0 || decoder == null) throw new InvalidOperationException($"No decodable audio stream in '{path}': {ErrorString(_streamIndex)}");

            for (var i = 0; i < (int)_format->nb_streams; i++)
            {
                if (i != _streamIndex) _format->streams[i]->discard = AVDiscard.AVDISCARD_ALL;
            }

            var stream = _format->streams[_streamIndex];
            _timeBase = stream->time_base;
            _startPts = stream->start_time != ffmpeg.AV_NOPTS_VALUE ? stream->start_time : 0;

            _codec = ffmpeg.avcodec_alloc_context3(decoder);
            if (_codec == null) throw new InvalidOperationException("avcodec_alloc_context3 failed.");
            ret = ffmpeg.avcodec_parameters_to_context(_codec, stream->codecpar);
            if (ret < 0) throw new InvalidOperationException($"avcodec_parameters_to_context failed: {ErrorString(ret)}");
            _codec->pkt_timebase = _timeBase;
            ret = ffmpeg.avcodec_open2(_codec, decoder, null);
            if (ret < 0) throw new InvalidOperationException($"Failed to open audio decoder for '{path}': {ErrorString(ret)}");

            Channels = _codec->ch_layout.nb_channels;
            SampleRate = _codec->sample_rate;
            if (Channels <= 0 || SampleRate <= 0) throw new InvalidOperationException($"Invalid audio format in '{path}': {Channels}ch {SampleRate}Hz.");

            TotalFrames = EstimateTotalFrames(stream);

            SwrContext* swr = null;
            ret = ffmpeg.swr_alloc_set_opts2(&swr,
                &_codec->ch_layout, AVSampleFormat.AV_SAMPLE_FMT_FLT, SampleRate,
                &_codec->ch_layout, _codec->sample_fmt, SampleRate, 0, null);
            if (ret < 0 || swr == null) throw new InvalidOperationException($"swr_alloc_set_opts2 failed: {ErrorString(ret)}");
            _swr = swr;
            ret = ffmpeg.swr_init(_swr);
            if (ret < 0) throw new InvalidOperationException($"swr_init failed: {ErrorString(ret)}");

            _packet = ffmpeg.av_packet_alloc();
            _frame = ffmpeg.av_frame_alloc();
            if (_packet == null || _frame == null) throw new InvalidOperationException("Failed to allocate packet/frame.");

            _syncToTimestamp = true;
        }
        catch
        {
            Dispose();
            throw;
        }
    }

    private long EstimateTotalFrames(AVStream* stream)
    {
        if (stream->duration != ffmpeg.AV_NOPTS_VALUE && stream->duration > 0)
            return ffmpeg.av_rescale_q(stream->duration, _timeBase, new AVRational { num = 1, den = SampleRate });
        if (_format->duration != ffmpeg.AV_NOPTS_VALUE && _format->duration > 0)
            return ffmpeg.av_rescale(_format->duration, SampleRate, ffmpeg.AV_TIME_BASE);
        throw new InvalidOperationException($"Audio file '{Path}' has no duration.");
    }

    public void Seek(long frame)
    {
        if (frame < 0 || frame > TotalFrames) throw new ArgumentOutOfRangeException(nameof(frame));
        if (frame == Position) return;

        var marginFrames = (long)(SeekMarginSeconds * SampleRate);

        // Target inside the pending block or less than the seek margin after it: decode forward and discard.
        if (!_syncToTimestamp && frame >= _pendingStart + _pendingOffset
            && frame <= _pendingStart + _pendingOffset + _pendingCount + marginFrames)
        {
            Position = frame;
            return;
        }

        var targetTs = _startPts + ffmpeg.av_rescale_q(Math.Max(0, frame - marginFrames), new AVRational { num = 1, den = SampleRate }, _timeBase);
        var ret = ffmpeg.av_seek_frame(_format, _streamIndex, targetTs, ffmpeg.AVSEEK_FLAG_BACKWARD);
        if (ret < 0)
            ret = ffmpeg.av_seek_frame(_format, _streamIndex, _startPts, ffmpeg.AVSEEK_FLAG_BACKWARD);
        if (ret < 0)
            throw new InvalidOperationException($"Failed to seek '{Path}' to frame {frame}: {ErrorString(ret)}");

        ffmpeg.avcodec_flush_buffers(_codec);
        SeekCount++;
        _pendingCount = 0;
        _pendingOffset = 0;
        _syncToTimestamp = true;
        _flushSent = false;
        _endOfStream = false;
        Position = frame;
    }

    public int Read(float[] buffer, int frameCount)
    {
        var limit = (int)Math.Min(frameCount, TotalFrames - Position);
        var written = 0;
        while (written < limit)
        {
            if (_pendingCount == 0 || Position >= _pendingStart + _pendingOffset + _pendingCount)
            {
                if (!DecodeNextFrame()) break;
                continue;
            }

            var pendingPosition = _pendingStart + _pendingOffset;
            if (Position > pendingPosition)
            {
                // Discard decoded samples before the read position (after a seek).
                var skip = (int)Math.Min(Position - pendingPosition, _pendingCount);
                _pendingOffset += skip;
                _pendingCount -= skip;
                continue;
            }
            if (Position < pendingPosition)
            {
                // A gap in the decoded timestamps: fill with silence up to the next decoded sample.
                var gap = (int)Math.Min(pendingPosition - Position, limit - written);
                Array.Clear(buffer, written * Channels, gap * Channels);
                written += gap;
                Position += gap;
                continue;
            }

            var count = Math.Min(_pendingCount, limit - written);
            Array.Copy(_pending, _pendingOffset * Channels, buffer, written * Channels, count * Channels);
            _pendingOffset += count;
            _pendingCount -= count;
            written += count;
            Position += count;
        }
        return written;
    }

    /// <summary>Decodes the next audio frame into the pending buffer. Returns false at the end of the stream.</summary>
    private bool DecodeNextFrame()
    {
        while (!_endOfStream)
        {
            var ret = ffmpeg.avcodec_receive_frame(_codec, _frame);
            if (ret == 0)
            {
                StoreFrame();
                ffmpeg.av_frame_unref(_frame);
                return true;
            }
            if (ret == ffmpeg.AVERROR_EOF)
            {
                _endOfStream = true;
                break;
            }
            if (ret != ffmpeg.AVERROR(ffmpeg.EAGAIN))
                throw new InvalidOperationException($"Audio decoding failed in '{Path}': {ErrorString(ret)}");

            if (_flushSent)
            {
                _endOfStream = true;
                break;
            }

            ret = ffmpeg.av_read_frame(_format, _packet);
            if (ret < 0)
            {
                if (ret != ffmpeg.AVERROR_EOF)
                    throw new InvalidOperationException($"Failed to read '{Path}': {ErrorString(ret)}");
                ffmpeg.avcodec_send_packet(_codec, null);
                _flushSent = true;
                continue;
            }

            try
            {
                if (_packet->stream_index != _streamIndex) continue;
                ret = ffmpeg.avcodec_send_packet(_codec, _packet);
                // A corrupt packet is skipped; decoding continues with the next one.
                if (ret < 0 && ret != ffmpeg.AVERROR_INVALIDDATA && ret != ffmpeg.AVERROR(ffmpeg.EAGAIN))
                    throw new InvalidOperationException($"Failed to send audio packet of '{Path}': {ErrorString(ret)}");
            }
            finally
            {
                ffmpeg.av_packet_unref(_packet);
            }
        }
        return false;
    }

    private void StoreFrame()
    {
        var samples = _frame->nb_samples;
        if (samples <= 0) return;

        if (_syncToTimestamp)
        {
            var pts = _frame->best_effort_timestamp;
            _nextDecodedFrame = pts != ffmpeg.AV_NOPTS_VALUE
                ? ffmpeg.av_rescale_q(pts - _startPts, _timeBase, new AVRational { num = 1, den = SampleRate })
                : Position;
            _syncToTimestamp = false;
        }

        if (_pending.Length < samples * Channels)
            _pending = new float[samples * Channels];

        fixed (float* pPending = _pending)
        {
            var outPlane = (byte*)pPending;
            var converted = ffmpeg.swr_convert(_swr, &outPlane, samples, _frame->extended_data, samples);
            if (converted < 0)
                throw new InvalidOperationException($"Sample format conversion failed in '{Path}': {ErrorString(converted)}");
            samples = converted;
        }

        _pendingStart = _nextDecodedFrame;
        _pendingOffset = 0;
        _pendingCount = samples;
        _nextDecodedFrame += samples;
    }

    public void Dispose()
    {
        if (_frame != null)
        {
            var frame = _frame;
            ffmpeg.av_frame_free(&frame);
            _frame = null;
        }
        if (_packet != null)
        {
            var packet = _packet;
            ffmpeg.av_packet_free(&packet);
            _packet = null;
        }
        if (_swr != null)
        {
            var swr = _swr;
            ffmpeg.swr_free(&swr);
            _swr = null;
        }
        if (_codec != null)
        {
            var codec = _codec;
            ffmpeg.avcodec_free_context(&codec);
            _codec = null;
        }
        if (_format != null)
        {
            var format = _format;
            ffmpeg.avformat_close_input(&format);
            _format = null;
        }
    }

    private static string ErrorString(int error)
    {
        var buffer = stackalloc byte[1024];
        ffmpeg.av_strerror(error, buffer, 1024);
        return $"{Marshal.PtrToStringAnsi((IntPtr)buffer)} ({error})";
    }
}