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

namespace ADOFAIRenderer.Util.Audio;

internal 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();
        }
    }
}

internal 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);
            }
        }
    }
}