ADOFAIRenderer/Util/Audio/OfflineAudioMixer.cs

using System;
using System.Collections.Generic;
using System.Runtime.CompilerServices;
using System.Threading;
using ADOFAIRenderer.Util.Audio.Clip;
using ADOFAIRenderer.Util.Audio.Filter;
using ADOFAIRenderer.Util.Audio.Mixer;
using MelonLoader;
using UnityEngine;

namespace ADOFAIRenderer.Util.Audio;

/// <summary>
/// Audio of one render: the song and effect sounds as <see cref="Voice"/>s of a <see cref="VoiceMixer"/>, clips from an
/// <see cref="AudioClipStore"/>, followed by the filter chain (RMS compressor → look-ahead limiter → hard clip → -1 dB gain).
///
/// Time: seconds relative to the render start (RenderManager.dspTimeBase); output sample n is at n / sample rate.
/// The limiter delays its output by its look-ahead (<see cref="IAudioFilter.Latency"/>), so output sample n is mixed from
/// voices at sample n + latency, and the filters are primed with samples [0, latency) before the first block.
///
/// Threads:
/// - Main thread: <see cref="InitializeBGM"/>, <see cref="CacheClip"/>, <see cref="RegisterSoundEvent"/>, <see cref="UpdateBGM"/>,
///   source/voice mapping, and <see cref="Commit"/> after each frame (all commands with an earlier time have been queued).
/// - Render thread: <see cref="GenerateAudioForFrame"/> waits until the main thread has committed the end of the samples it mixes,
///   so every command takes effect at its own sample. <see cref="Complete"/> ends the waiting when the main thread stops.
/// </summary>
internal sealed class OfflineAudioMixer : IDisposable
{
    /// <summary>Gain applied to every clip.</summary>
    private const double ClipVolume = 0.5;

    /// <summary>Additional gain of effect sounds (registered with <see cref="RegisterSoundEvent"/>).</summary>
    private const double EffectVolume = 2.0 / 3.0;

    /// <summary>
    /// Smallest number of tracked effect voices at which finished ones are removed from the list used by <see cref="StopAllEffects"/>.
    /// After each removal the next threshold is twice the remaining count (at least this value).
    /// </summary>
    private const int EffectVoicePruneCount = 1024;

    /// <summary>Longest wait for a commit before checking for cancellation again, in milliseconds.</summary>
    private const int CommitWaitMs = 50;

    private readonly int _channels;
    private readonly int _sampleRate;
    private readonly AudioClipStore _store;
    private readonly VoiceMixer _mixer;
    private readonly CancellationTokenSource _cancellation = new();
    private readonly List<IAudioFilter> _filters = new();
    private readonly int _latencySamples;

    // Main thread.
    private ConditionalWeakTable<AudioSource, Voice> _sourceVoices = new();
    private readonly List<Voice> _effectVoices = new();
    private int _effectVoicePruneAt = EffectVoicePruneCount;
    private AudioClipSource? _bgmSource;
    private string _bgmKey = "";
    private Voice? _bgmVoice;
    private double _bgmStart = double.NaN;
    private double _bgmPitch = double.NaN;
    private double _bgmVolume = double.NaN;

    // Shared: committed output sample (voices are final before it), guarded by _commitLock.
    private readonly object _commitLock = new();
    private long _committedSample = long.MinValue;

    // Render thread (and Dispose, under _renderLock).
    private readonly object _renderLock = new();
    private long _nextOutputSample;
    private bool _primed;
    private double[] _mixBuffer = Array.Empty<double>();
    private bool _disposed;

    public OfflineAudioMixer(int targetChannels, int targetSampleRate)
    {
        _channels = targetChannels;
        _sampleRate = targetSampleRate;
        _store = new AudioClipStore(targetSampleRate);
        _mixer = new VoiceMixer(_store, targetChannels, _cancellation.Token);

        var rmsCompressor = new RMSCompressor
        {
            ThresholdDb = -10, Ratio = 2,
            AttackTime = 0.1, ReleaseTime = 0.1,
            KneeDb = 10, RmsWindowTime = 0.02
        };
        var limiter = new RMSCompressor
        {
            ThresholdDb = 0, Ratio = double.PositiveInfinity,
            AttackTime = 0, ReleaseTime = 0.05,
            KneeDb = 0, RmsWindowTime = 0.02,
            LookAheadTime = 0.01
        };
        _filters.Add(rmsCompressor);
        _filters.Add(limiter);
        _filters.Add(new HardClip());
        _filters.Add(new Gain { Level = -1 });

        double latency = 0;
        foreach (var filter in _filters)
        {
            filter.Initialize(_sampleRate, _channels);
            latency += filter.Latency;
        }
        _latencySamples = (int)Math.Round(latency * _sampleRate);

        _store.Start();
    }

    // ── Main thread ──

    /// <summary>Registers the song clip. The song voice is created and kept in sync by <see cref="UpdateBGM"/>.</summary>
    public void InitializeBGM(AudioClip? clip)
    {
        if (!clip)
        {
            _bgmSource = null;
            return;
        }

        _bgmKey = clip!.name;
        _bgmSource = AudioClipResolver.GetOrAdd(_store, _bgmKey, clip);
        if (_bgmSource == null)
        {
            Melon<RenderMod>.Logger.Error($"[OfflineAudioMixer] Song clip '{_bgmKey}' could not be loaded, the song will be silent.");
            return;
        }
        Melon<RenderMod>.Logger.Msg($"[OfflineAudioMixer] Song '{_bgmKey}': {_bgmSource.Channels} ch, {_bgmSource.Duration:F2} s, {_bgmSource.Policy}");
    }

    /// <summary>
    /// Updates the song voice from the conductor (start time, pitch) and the level volume.
    /// A changed start time restarts the song voice at position (now - start) * pitch; pitch and volume changes take effect at <paramref name="now"/>.
    /// Start time: dspTimeSong + separate countdown (crotchetAtStart / pitch * adjustedCountdownTicks when separateCountdownTime),
    /// the time scrConductor.StartMusicCo schedules the song at.
    /// </summary>
    /// <param name="now">Current time, seconds relative to the render start.</param>
    /// <param name="dspTimeBase">AudioSettings.dspTime at the render start.</param>
    public void UpdateBGM(double now, double dspTimeBase)
    {
        if (_bgmSource == null) return;
        var conductor = scrConductor.instance;
        if (!conductor) return;

        double pitch = conductor.song.pitch;
        var countdown = conductor.separateCountdownTime ? conductor.crotchetAtStart / pitch * conductor.adjustedCountdownTicks : 0.0;
        var start = conductor.dspTimeSong + countdown - dspTimeBase;
        var levelVolume = scnGame.instance ? scnGame.instance.levelData.volume * 0.01 : 1.0;
        var volume = levelVolume * ClipVolume;

        if (_bgmVoice == null || start != _bgmStart)
        {
            // Song position at time t is (t - start) * pitch, also when the start is already past.
            _bgmVoice?.Stop(now);
            _bgmVoice = _mixer.Play(_bgmSource, _bgmKey, Math.Max(start, now), volume, pitch, false, now);
            if (start < now) _bgmVoice.SetPosition((now - start) * pitch, now);
            _bgmStart = start;
            _bgmPitch = pitch;
            _bgmVolume = volume;
            return;
        }

        if (pitch != _bgmPitch)
        {
            _bgmVoice.SetPitch(pitch, now);
            _bgmPitch = pitch;
        }
        if (volume != _bgmVolume)
        {
            _bgmVoice.SetVolume(volume, now);
            _bgmVolume = volume;
        }
    }

    /// <summary>Prepares the source of the game clip <paramref name="clipName"/> (audioLib key) so its first chunk starts loading.</summary>
    public void CacheClip(string clipName)
    {
        GetOrAddSource(clipName);
    }

    private AudioClipSource? GetOrAddSource(string clipName)
    {
        if (string.IsNullOrEmpty(clipName)) return null;
        if (_store.TryGet(clipName, out var existing)) return existing;
        var clip = AudioManager.Instance ? AudioManager.Instance.FindOrLoadAudioClip(clipName) : null;
        return AudioClipResolver.GetOrAdd(_store, clipName, clip);
    }

    /// <summary>
    /// Starts an effect sound at <paramref name="startTime"/>. An <paramref name="endTime"/> of 0 or more makes it loop until then.
    /// <paramref name="source"/> (when given) is mapped to the voice for later changes through the AudioSource.
    /// </summary>
    public Voice RegisterSoundEvent(AudioSource? source, string clipName, double startTime, double endTime, float volume, float pitch, double now)
    {
        var clipSource = GetOrAddSource(clipName);
        var loop = endTime >= 0;
        var voice = _mixer.Play(clipSource, clipName, startTime, volume * ClipVolume * EffectVolume, pitch, loop, now);
        if (loop) voice.SetEndTime(endTime, now);

        if (_effectVoices.Count >= _effectVoicePruneAt)
        {
            _effectVoices.RemoveAll(v => v.IsFinished);
            _effectVoicePruneAt = Math.Max(EffectVoicePruneCount, _effectVoices.Count * 2);
        }
        _effectVoices.Add(voice);

        if (source) MapSource(source!, voice);
        return voice;
    }

    /// <summary>Maps <paramref name="source"/> to <paramref name="voice"/> for later changes through the AudioSource (replaces an earlier mapping).</summary>
    public void MapSource(AudioSource source, Voice voice)
    {
        _sourceVoices.Remove(source);
        _sourceVoices.Add(source, voice);
    }

    /// <summary>Voice started for <paramref name="source"/> by <see cref="RegisterSoundEvent"/>.</summary>
    public bool TryGetVoice(AudioSource? source, out Voice voice)
    {
        voice = null!;
        return source != null && _sourceVoices.TryGetValue(source, out voice);
    }

    public bool RemoveSourceMapping(AudioSource source) => _sourceVoices.Remove(source);

    /// <summary>
    /// Stops every effect sound at <paramref name="now"/>, including ones scheduled to start later
    /// (AudioManager.StopAllSounds destroys every AudioSource in liveSources, which holds the sources of
    /// AudioManager._Play and scrConductor.PlayWithEndTime). The song voice keeps playing.
    /// </summary>
    public void StopAllEffects(double now)
    {
        foreach (var voice in _effectVoices)
            voice.Stop(now);
        _effectVoices.Clear();
        _sourceVoices = new ConditionalWeakTable<AudioSource, Voice>();
    }

    /// <summary>
    /// Marks every command with a time before <paramref name="time"/> as queued
    /// (called after each frame with the time of the next frame).
    /// </summary>
    public void Commit(double time)
    {
        var sample = VoiceMixer.ToSample(time, _sampleRate);
        lock (_commitLock)
        {
            if (sample <= _committedSample) return;
            _committedSample = sample;
            Monitor.PulseAll(_commitLock);
        }
    }

    /// <summary>No more commands will be queued (render ending): mixing no longer waits for commits.</summary>
    public void Complete()
    {
        lock (_commitLock)
        {
            _committedSample = long.MaxValue;
            Monitor.PulseAll(_commitLock);
        }
    }

    /// <summary>Stops waiting for commits and chunks; the render thread's current or next mix throws OperationCanceledException.</summary>
    public void Cancel()
    {
        _cancellation.Cancel();
        Complete();
    }

    // ── Render thread ──

    /// <summary>
    /// Writes the next <paramref name="samplesNeeded"/> output frames (interleaved float, continuing from the previous call).
    /// </summary>
    /// <exception cref="InvalidOperationException">Chunk loading failed fatally (out of memory or the decoder thread stopped).</exception>
    /// <exception cref="OperationCanceledException">The render was canceled (<see cref="Cancel"/>).</exception>
    public void GenerateAudioForFrame(float[] mixedBuffer, int samplesNeeded)
    {
        lock (_renderLock)
        {
            if (_disposed) throw new ObjectDisposedException(nameof(OfflineAudioMixer));

            if (!_primed)
            {
                _primed = true;
                if (_latencySamples > 0)
                {
                    var prime = MixBlock(0, _latencySamples);
                    foreach (var filter in _filters)
                        filter.ProcessBlock(prime);
                }
            }

            var block = MixBlock(_nextOutputSample + _latencySamples, samplesNeeded);
            foreach (var filter in _filters)
                filter.ProcessBlock(block);
            _nextOutputSample += samplesNeeded;

            for (var i = 0; i < block.Length; i++)
                mixedBuffer[i] = (float)block[i];
        }
    }

    /// <summary>Mixes voices at samples [first, first + frames) into the reused double buffer, after the main thread committed them.</summary>
    private Span<double> MixBlock(long first, int frames)
    {
        WaitForCommit(first + frames);

        var length = frames * _channels;
        if (_mixBuffer.Length < length) _mixBuffer = new double[length];
        var block = _mixBuffer.AsSpan(0, length);
        block.Clear();
        _mixer.Mix(first, block);
        return block;
    }

    private void WaitForCommit(long endSample)
    {
        lock (_commitLock)
        {
            while (_committedSample < endSample)
            {
                _cancellation.Token.ThrowIfCancellationRequested();
                Monitor.Wait(_commitLock, CommitWaitMs);
            }
        }
        _cancellation.Token.ThrowIfCancellationRequested();
    }

    /// <summary>
    /// Cancels waits, then releases the voices, the clip store (decoder thread) and the filter state.
    /// Waits for a running <see cref="GenerateAudioForFrame"/> to return.
    /// </summary>
    public void Dispose()
    {
        Cancel();
        lock (_renderLock)
        {
            if (_disposed) return;
            _disposed = true;
            _mixer.Dispose();
            _store.Dispose();
            foreach (var filter in _filters)
                filter.Reset();
        }
        _sourceVoices = new ConditionalWeakTable<AudioSource, Voice>();
        _effectVoices.Clear();
    }
}