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 sound effects 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 Unity 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 Unity thread stops.
/// </summary>
internal sealed class OfflineAudioMixer : IDisposable
{
    /// <summary>Linear gain applied to every clip (song and effects), before the filter chain (0.5 ≈ -6 dB headroom).</summary>
    private const double ClipVolume = 0.5;

    /// <summary>Additional gain of sound effects (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;

    /// <summary>Output channel count (interleaved).</summary>
    private readonly int _channels;

    /// <summary>Output sample rate in Hz.</summary>
    private readonly int _sampleRate;

    private readonly AudioClipStore _store;
    private readonly VoiceMixer _mixer;

    /// <summary>Canceled by <see cref="Cancel"/>/<see cref="Dispose"/>; ends commit waits and chunk waits on the render thread.</summary>
    private readonly CancellationTokenSource _cancellation = new();

    /// <summary>Filter chain applied in list order to every mixed block.</summary>
    private readonly List<IAudioFilter> _filters = new();

    /// <summary>Sum of the filter latencies in output samples: round(sum of <see cref="IAudioFilter.Latency"/> * sample rate).</summary>
    private readonly int _latencySamples;

    // Main thread.

    /// <summary>
    /// Mapping of game AudioSources to the state of the sound effect last started on them. Weak keys: a destroyed source
    /// does not keep its state alive. Replaced by a new table to drop every mapping at once.
    /// </summary>
    private ConditionalWeakTable<AudioSource, SourceState> _sourceStates = new();

    /// <summary>Effect voices started since the last <see cref="StopAllEffects"/> (finished ones are pruned periodically).</summary>
    private readonly List<Voice> _effectVoices = new();

    /// <summary>Count of <see cref="_effectVoices"/> at which the next prune runs.</summary>
    private int _effectVoicePruneAt = EffectVoicePruneCount;

    /// <summary>Song source, or null when there is no song or it could not be loaded.</summary>
    private AudioClipSource? _bgmSource;

    /// <summary>audioLib key (AudioClip.name) of the song.</summary>
    private string _bgmKey = "";

    /// <summary>Current song voice, created by the first <see cref="UpdateBGM"/> and replaced when the start time changes.</summary>
    private Voice? _bgmVoice;

    // Values last sent to the song voice (start in seconds relative to the render start); NaN before the first update.
    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();

    /// <summary>Index of the next output sample <see cref="GenerateAudioForFrame"/> writes (frames per channel since the render start).</summary>
    private long _nextOutputSample;

    /// <summary>True once the filters have been fed samples [0, latency).</summary>
    private bool _primed;

    /// <summary>Reused interleaved mix buffer, grown to the largest block requested.</summary>
    private double[] _mixBuffer = Array.Empty<double>();

    private bool _disposed;

    /// <summary>
    /// Creates the clip store (and starts its decoder thread), the voice mixer and the filter chain.
    /// Filter settings: RMS compressor (threshold -10 dB, ratio 2, attack/release 0.1 s, knee 10 dB, RMS window 0.02 s),
    /// limiter (RMS compressor with threshold 0 dB, infinite ratio, attack 0 raised to the 0.0001 s minimum, release 0.05 s,
    /// hard knee, RMS window 0.02 s, look-ahead 0.01 s), hard clip at ±1.0, then -1 dB gain.
    /// </summary>
    /// <param name="targetChannels">Output channel count (interleaved).</param>
    /// <param name="targetSampleRate">Output sample rate in Hz; every clip is resampled to it.</param>
    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 });

        // Total delay of the chain; only the limiter's look-ahead contributes (0.01 s).
        double latency = 0;
        foreach (var filter in _filters)
        {
            filter.Initialize(_sampleRate, _channels);
            latency += filter.Latency;
        }
        _latencySamples = (int)Math.Round(latency * _sampleRate);

        _store.Start();
    }

    /// <summary>Clip store of this render (output rate = mixer sample rate), also read by <see cref="AudioSpectrumGenerator"/>.</summary>
    internal AudioClipStore Store => _store;

    // ── 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;
        // levelData.volume is a percentage (100 = unchanged; scrConductor sets song.volume = volume * 0.01 the same way).
        // Without a scnGame instance the level volume is 1.0.
        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);
    }

    /// <summary>
    /// Source of the game clip <paramref name="clipName"/> (audioLib key): an existing source, or one opened through
    /// <see cref="AudioClipResolver"/> with the clip from AudioManager.FindOrLoadAudioClip. Null for an empty name or an
    /// unreadable clip (the key is then remembered as missing by the store). Main thread only.
    /// </summary>
    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 a sound effect 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 returned state, which reports and applies later changes through the AudioSource.
    /// </summary>
    /// <param name="source">Game AudioSource playing the sound, or null to map it later with <see cref="MapSource"/>.</param>
    /// <param name="clipName">audioLib key of the clip. A clip that cannot be loaded gives a silent voice.</param>
    /// <param name="startTime">Scheduled start, seconds relative to the render start (a start before <paramref name="now"/> starts at now).</param>
    /// <param name="endTime">Loop end, seconds relative to the render start; negative for a sound that plays once.</param>
    /// <param name="volume">Game volume (linear); multiplied by the clip and effect gains.</param>
    /// <param name="pitch">Playback rate at the start.</param>
    /// <param name="now">Current render time in seconds.</param>
    /// <returns>State of the new sound, also used by the AudioSource patches.</returns>
    public SourceState 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);

        // Amortized pruning: the list is scanned only when it has doubled since the last prune.
        if (_effectVoices.Count >= _effectVoicePruneAt)
        {
            _effectVoices.RemoveAll(v => v.IsFinished);
            _effectVoicePruneAt = Math.Max(EffectVoicePruneCount, _effectVoices.Count * 2);
        }
        _effectVoices.Add(voice);

        var state = new SourceState(voice, Math.Max(startTime, now), pitch, loop, loop ? endTime : double.PositiveInfinity,
            clipSource?.Duration ?? 0.0, _sampleRate);
        if (source) MapSource(source!, state);
        return state;
    }

    /// <summary>
    /// Maps <paramref name="source"/> to <paramref name="state"/> (replaces an earlier mapping) and takes the clip length and
    /// sample rate from the clip of the source.
    /// </summary>
    public void MapSource(AudioSource source, SourceState state)
    {
        state.BindClip(source.clip);
        _sourceStates.Remove(source);
        _sourceStates.Add(source, state);
    }

    /// <summary>State of the sound effect last started for <paramref name="source"/>. The mapping stays after the sound ends or stops.</summary>
    public bool TryGetSourceState(AudioSource? source, out SourceState state)
    {
        state = null!;
        return source != null && _sourceStates.TryGetValue(source, out state);
    }

    /// <summary>
    /// Stops every sound effect 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();
        _sourceStates = new ConditionalWeakTable<AudioSource, SourceState>();
    }

    /// <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.
    /// May be called from any thread, more than once.
    /// </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).
    /// Blocks until the main Unity thread has committed the voice samples it needs (<see cref="Commit"/>, <see cref="Complete"/>).
    /// </summary>
    /// <param name="mixedBuffer">Destination; must hold at least samplesNeeded * channels floats. Values beyond that are untouched.</param>
    /// <param name="samplesNeeded">Frames (samples per channel) to write.</param>
    /// <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));

            // First call: run voice samples [0, latency) through the filters and discard the output, which is the
            // zero-filled delay lines. Afterwards output sample n is the filtered voice sample n.
            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 Unity 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;
    }

    /// <summary>
    /// Waits until the committed sample reaches <paramref name="endSample"/> (exclusive end of the block to mix), checking for
    /// cancellation every <see cref="CommitWaitMs"/> ms.
    /// </summary>
    /// <exception cref="OperationCanceledException">The render was canceled.</exception>
    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();
        }
        _sourceStates = new ConditionalWeakTable<AudioSource, SourceState>();
        _effectVoices.Clear();
    }
}