ADOFAIRenderer/Util/Audio/OfflineAudioMixer.cs

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

namespace ADOFAIRenderer.Util.Audio
{
    public class AudioClipCache
    {
        private readonly Dictionary<string, CachedClipData> _cache = new();

        public class CachedClipData
        {
            public float[] Samples = null!; // Pre-resampled to target channels and sample rate
            public double Duration;
        }

        public CachedClipData? GetOrCache(string clipName, int targetChannels, int targetSampleRate, bool allowLoad = true)
        {
            if (string.IsNullOrEmpty(clipName)) return null;

            lock (_cache)
            {
                if (_cache.TryGetValue(clipName, out var cached))
                    return cached;

                if (!allowLoad) return null;

                if (!AudioManager.Instance) return null;

                var clip = AudioManager.Instance.FindOrLoadAudioClip(clipName);
                if (!clip) return null;

                try
                {
                    var totalSamples = clip.samples;
                    var channels = clip.channels;
                    var rawSamples = new float[totalSamples * channels];
                    clip.GetData(rawSamples, 0);

                    // Auto-resample during cache generation
                    var resampled = AudioResampler.Resample(rawSamples, channels, clip.frequency, targetChannels, targetSampleRate);

                    var data = new CachedClipData
                    {
                        Samples = resampled,
                        Duration = (double)resampled.Length / (targetChannels * targetSampleRate)
                    };

                    _cache[clipName] = data;
                    return data;
                }
                catch (Exception e)
                {
                    MelonLogger.Error($"[AudioClipCache] Failed to cache and resample clip {clipName}: {e}");
                    return null;
                }
            }
        }

        public void Clear()
        {
            lock (_cache)
            {
                _cache.Clear();
            }
        }
    }

    public class OfflineAudioMixer
    {
        public class SoundEvent
        {
            public string ClipName = null!;
            public double StartTime; // In seconds relative to dspTimeBase
            public double EndTime;   // In seconds relative to dspTimeBase (-1 if one-shot)
            public float Volume;
            public float Pitch;
        }

        private readonly List<SoundEvent> _activeEvents = new();
        private ConditionalWeakTable<AudioSource, SoundEvent> _sourceMappings = new();
        private readonly object _lock = new();

        // Target Audio Settings
        private readonly int _targetChannels;
        private readonly int _targetSampleRate;
        private const float ClipVolume = .5f;
        
        // Audio Filter Settings
        private readonly LinkedList<IAudioFilter> _filters = new();

        // BGM Cached Samples (Pre-resampled)
        private float[]? _bgmSamples;
        private double _bgmDuration;

        // Clip Cache
        private readonly AudioClipCache _clipCache = new();
        private bool _isPrimed;

        public OfflineAudioMixer(int targetChannels, int targetSampleRate)
        {
            _targetChannels = targetChannels;
            _targetSampleRate = targetSampleRate;

            var rmsCompressor = new RMSCompressor
            {
                ThresholdDb = -10, Ratio = 2,
                AttackTime = 0.1f, ReleaseTime = 0.1f,
                KneeDb = 10, RmsWindowTime = 0.02f
            };
            rmsCompressor.Initialize(_targetSampleRate, _targetChannels);
            _filters.AddLast(rmsCompressor);
            
            var limiter = new RMSCompressor
            {
                ThresholdDb = 0, Ratio = float.PositiveInfinity,
                AttackTime = 0f, ReleaseTime = 0.05f,
                KneeDb = 0, RmsWindowTime = 0.02f,
                LookAheadTime = 0.01f
            };
            limiter.Initialize(_targetSampleRate, _targetChannels);
            _filters.AddLast(limiter);
            
            var hardClip = new HardClip();
            hardClip.Initialize(_targetSampleRate, _targetChannels);
            _filters.AddLast(hardClip);

            var truePeakSafetyGain = new Gain { Level = -1f };
            truePeakSafetyGain.Initialize(_targetSampleRate, _targetChannels);
            _filters.AddLast(truePeakSafetyGain);
        }

        public void InitializeBGM(AudioClip? clip)
        {
            lock (_lock)
            {
                if (!clip)
                {
                    _bgmSamples = null;
                    _bgmDuration = 0;
                    return;
                }

                try
                {
                    int originalSamples = clip.samples;
                    int originalChannels = clip.channels;
                    float[] rawSamples = new float[originalSamples * originalChannels];
                    clip.GetData(rawSamples, 0);

                    // Pre-resample the BGM to the target format
                    _bgmSamples = AudioResampler.Resample(rawSamples, originalChannels, clip.frequency, _targetChannels, _targetSampleRate);
                    _bgmDuration = (double)_bgmSamples.Length / (_targetChannels * _targetSampleRate);

                    Melon<RenderMod>.Logger.Msg($"[OfflineAudioMixer] Initialized resampled BGM: {clip.name}, Samples: {_bgmSamples.Length / _targetChannels}, TargetFreq: {_targetSampleRate}Hz, TargetChannels: {_targetChannels}");
                }
                catch (Exception e)
                {
                    Melon<RenderMod>.Logger.Error($"[OfflineAudioMixer] Failed to initialize and resample BGM: {e}");
                    _bgmSamples = null;
                    _bgmDuration = 0;
                }
            }
        }

        public void CacheClip(string clipName)
        {
            _clipCache.GetOrCache(clipName, _targetChannels, _targetSampleRate, allowLoad: true);
        }

        public void RegisterSoundEvent(AudioSource? source, string clipName, double startTime, double endTime, float volume, float pitch)
        {
            lock (_lock)
            {
                var ev = new SoundEvent
                {
                    ClipName = clipName,
                    StartTime = startTime,
                    EndTime = endTime,
                    Volume = volume,
                    Pitch = pitch
                };
                _activeEvents.Add(ev);
                
                if (!source) return;
                _sourceMappings.Remove(source);
                _sourceMappings.Add(source, ev);
            }
        }

        public void RegisterSoundEvent(string clipName, double startTime, double endTime, float volume, float pitch)
        {
            RegisterSoundEvent(null, clipName, startTime, endTime, volume, pitch);
        }

        public bool TryGetSoundEvent(AudioSource? source, out SoundEvent ev)
        {
            if (source == null)
            {
                ev = null!;
                return false;
            }
            lock (_lock)
            {
                return _sourceMappings.TryGetValue(source, out ev);
            }
        }

        public bool RemoveSourceMapping(AudioSource source)
        {
            lock (_lock)
            {
                return _sourceMappings.Remove(source);
            }
        }

        public void Reset()
        {
            lock (_lock)
            {
                _activeEvents.Clear();
                _sourceMappings = new ConditionalWeakTable<AudioSource, SoundEvent>();
                _bgmSamples = null;
                _bgmDuration = 0;
                _clipCache.Clear();
                _isPrimed = false;
                foreach (var filter in _filters)
                {
                    filter.Reset();
                }
            }
        }

        private void PrimeFilters(double totalLatency, double dspTimeBase)
        {
            int latencySamples = (int)Math.Round(totalLatency * _targetSampleRate);
            if (latencySamples <= 0) return;

            // Rent temporary buffer to prime filters with initial audio
            float[] tempBuffer = System.Buffers.ArrayPool<float>.Shared.Rent(latencySamples * _targetChannels);
            try
            {
                Array.Clear(tempBuffer, 0, latencySamples * _targetChannels);

                // Mix starting from -totalLatency to 0
                MixBGM(tempBuffer, latencySamples, -totalLatency, dspTimeBase, totalLatency);
                MixEffectSounds(tempBuffer, latencySamples, -totalLatency, totalLatency);

                // Feed the buffer into filters to prime circular delay buffers
                foreach (var filter in _filters)
                {
                    filter.ProcessBlock(tempBuffer.AsSpan(0, latencySamples * _targetChannels));
                }
            }
            finally
            {
                System.Buffers.ArrayPool<float>.Shared.Return(tempBuffer);
            }
        }

        public void GenerateAudioForFrame(float[] mixedBuffer, int samplesNeeded, double frameStartTime, double dspTimeBase)
        {
            // mixedBuffer size must be at least samplesNeeded * _targetChannels
            Array.Clear(mixedBuffer, 0, samplesNeeded * _targetChannels);

            lock (_lock)
            {
                double totalLatency = 0;
                foreach (var filter in _filters)
                {
                    totalLatency += filter.Latency;
                }

                // Prime filters on first frame execution if latency exists
                if (!_isPrimed && totalLatency > 0)
                {
                    PrimeFilters(totalLatency, dspTimeBase);
                    _isPrimed = true;
                }

                // 1. Mix BGM
                MixBGM(mixedBuffer, samplesNeeded, frameStartTime, dspTimeBase, totalLatency);

                // 2. Mix Hitsounds/Effect Sounds
                MixEffectSounds(mixedBuffer, samplesNeeded, frameStartTime, totalLatency);

                foreach (var filter in _filters)
                    filter.ProcessBlock(mixedBuffer.AsSpan(0, samplesNeeded * _targetChannels));

                for (var i = 0; i < samplesNeeded * _targetChannels; i++)
                {
                    if (Math.Abs(mixedBuffer[i]) > 1) Melon<RenderMod>.Logger.Msg("Warning Why Audio Sample Is Larger Than 1");
                }
            }
        }

        private void MixBGM(float[] mixedBuffer, int samplesNeeded, double frameStartTime, double dspTimeBase, double totalLatency)
        {
            if (_bgmSamples == null || _bgmSamples.Length == 0) return;

            var conductor = scrConductor.instance;
            if (!conductor) return;

            double pitch = conductor.song.pitch;
            var dspTimeSong = conductor.dspTimeSong;
            var separateCountdownOffset = conductor.separateCountdownTime
                ? conductor.crotchetAtStart / pitch * conductor.adjustedCountdownTicks
                : 0.0;

            // BGM Playback start time relative to dspTimeBase
            var time1 = dspTimeSong + separateCountdownOffset;
            var relativeTime1 = time1 - dspTimeBase;
            var adjustedRelativeTime1 = relativeTime1 - totalLatency;

            var totalSamples = _bgmSamples.Length / _targetChannels;

            for (int s = 0; s < samplesNeeded; s++)
            {
                double sampleTime = frameStartTime + (double)s / _targetSampleRate;
                double clipTime = (sampleTime - adjustedRelativeTime1) * pitch;

                if (clipTime >= 0 && clipTime < _bgmDuration)
                {
                    double sampleIdxDouble = clipTime * _targetSampleRate;
                    long idxLow = (long)Math.Floor(sampleIdxDouble);
                    long idxHigh = idxLow + 1;
                    float t = (float)(sampleIdxDouble - idxLow);

                    if (idxHigh >= totalSamples) idxHigh = totalSamples - 1;
                    if (idxLow >= totalSamples) idxLow = totalSamples - 1;

                    for (var c = 0; c < _targetChannels; c++)
                    {
                        float valLow = _bgmSamples[idxLow * _targetChannels + c];
                        float valHigh = _bgmSamples[idxHigh * _targetChannels + c];
                        mixedBuffer[s * _targetChannels + c] += (valLow + t * (valHigh - valLow)) * (scnGame.instance.levelData.volume * 0.01f)
                            * ClipVolume;
                    }
                }
            }
        }

        private void MixEffectSounds(float[] mixedBuffer, int samplesNeeded, double frameStartTime, double totalLatency)
        {
            for (var i = _activeEvents.Count - 1; i >= 0; i--)
            {
                var ev = _activeEvents[i];
                var clipData = _clipCache.GetOrCache(ev.ClipName, _targetChannels, _targetSampleRate, allowLoad: false);
                if (clipData == null)
                {
                    _activeEvents.RemoveAt(i);
                    continue;
                }

                var isSoundFinished = false;
                var totalSamples = clipData.Samples.Length / _targetChannels;
                double adjustedStartTime = ev.StartTime - totalLatency;
                double adjustedEndTime = ev.EndTime >= 0 ? ev.EndTime - totalLatency : -1.0;

                for (var s = 0; s < samplesNeeded; s++)
                {
                    var sampleTime = frameStartTime + (double)s / _targetSampleRate;

                    // If it is a looping sound and has reached or exceeded its end time, stop playing.
                    if (adjustedEndTime >= 0 && sampleTime >= adjustedEndTime)
                    {
                        isSoundFinished = true;
                        break;
                    }

                    var elapsed = sampleTime - adjustedStartTime;
                    if (elapsed < 0) continue; // Sound hasn't started yet

                    var soundTime = elapsed * ev.Pitch;

                    if (soundTime >= clipData.Duration)
                    {
                        if (ev.EndTime < 0) // One-shot
                        {
                            isSoundFinished = true;
                            break;
                        }

                        // Looping
                        soundTime %= clipData.Duration;
                    }

                    // Linear interpolation for pitch shifting in the mixer loop
                    double sampleIdxDouble = soundTime * _targetSampleRate;
                    long idxLow = (long)Math.Floor(sampleIdxDouble);
                    long idxHigh = idxLow + 1;
                    float t = (float)(sampleIdxDouble - idxLow);

                    if (idxHigh >= totalSamples) idxHigh = totalSamples - 1;
                    if (idxLow >= totalSamples) idxLow = totalSamples - 1;

                    for (var c = 0; c < _targetChannels; c++)
                    {
                        float valLow = clipData.Samples[idxLow * _targetChannels + c];
                        float valHigh = clipData.Samples[idxHigh * _targetChannels + c];
                        mixedBuffer[s * _targetChannels + c] += (valLow + t * (valHigh - valLow)) * ev.Volume 
                            * ClipVolume * 2 / 3;
                    }
                }

                // Clean up finished sound events (both one-shot and looping)
                if (isSoundFinished)
                {
                    _activeEvents.RemoveAt(i);
                }
            }
        }
    }
}