using System;
using System.Collections.Generic;
using ADOFAIRenderer.Core;
using ADOFAIRenderer.Util.Audio.Clip;
using MelonLoader;
using UnityEngine;

namespace ADOFAIRenderer.Util.Audio;

/// <summary>
/// Replacement for AudioSource.GetSpectrumData during rendering: the spectrum of the song as the song AudioSource would output it.
///
/// Signal: the song clip from an <see cref="AudioClipStore"/> at the output sample rate (AudioSettings.outputSampleRate),
/// played at conductor.song.pitch (linear interpolation between frames) and scaled by conductor.song.volume.
/// Unity analyzes the output signal of the source: bin spacing follows the output rate, pitch moves the frequencies,
/// and volume scales the magnitudes (measured with a Unity 6000.3.10f1 test project).
/// Channel: stereo clips use the requested channel (0 = left, 1 = right); a mono clip is scaled by <see cref="MonoGain"/>.
/// Window: twice the requested spectrum length in output samples, centered on the current render time.
///
/// The song source is shared with the <see cref="OfflineAudioMixer"/> store when its output rate matches; otherwise a separate
/// store is created. A cursor at the analysis position keeps the needed chunks loaded. Main thread only.
/// </summary>
internal static class AudioSpectrumGenerator
{
    /// <summary>
    /// Gain of a mono clip relative to one channel of a stereo clip with the same samples.
    /// Measured with a Unity 6000.3.10f1 test project: the spectrum of a mono clip is about 0.70 times that of each stereo channel (1/√2).
    /// </summary>
    private const double MonoGain = 0.70710678118654752;

    /// <summary>Store created by this class when the mixer store cannot be shared (disposed by <see cref="Clear"/>); otherwise null.</summary>
    private static AudioClipStore? _ownStore;

    /// <summary>Store the song source belongs to: the mixer store or <see cref="_ownStore"/>.</summary>
    private static AudioClipStore? _store;

    /// <summary>Song source at the output sample rate, or null when not initialized or the song could not be loaded.</summary>
    private static AudioClipSource? _source;

    /// <summary>Cursor on <see cref="_source"/> moved to each analysis window, so the decoder thread loads the chunks there.</summary>
    private static AudioCursor? _cursor;

    /// <summary>Chunk currently held for <see cref="ReadSample"/> (invalid when none).</summary>
    private static ChunkView _view;

    /// <summary>True after a chunk loading error was logged (logged once per render).</summary>
    private static bool _loadErrorLogged;

    /// <summary>RenderManager.frameCount of the spectra in <see cref="CachedSpectrums"/>; -1 when empty.</summary>
    private static long _cachedFrame = -1;

    /// <summary>
    /// Spectra computed in the current render frame, by spectrum length and source channel. Several game objects read the
    /// spectrum every frame; within one frame the render time does not change, so they receive the same values.
    /// </summary>
    private static readonly Dictionary<(int Size, int Channel), float[]> CachedSpectrums = new();

    /// <summary>Registers the song clip. Call after the <see cref="OfflineAudioMixer"/> is created (null when audio is not rendered).</summary>
    public static void Initialize(AudioClip? clip, OfflineAudioMixer? mixer)
    {
        Clear();
        if (!clip) return;

        try
        {
            var rate = AudioSettings.outputSampleRate;
            if (mixer != null && mixer.Store.OutputRate == rate)
            {
                _store = mixer.Store;
            }
            else
            {
                _ownStore = new AudioClipStore(rate);
                _ownStore.Start();
                _store = _ownStore;
            }

            _source = AudioClipResolver.GetOrAdd(_store, clip!.name, clip);
            if (_source == null)
            {
                Melon<RenderMod>.Logger.Error($"[AudioSpectrumGenerator] Song clip '{clip.name}' could not be loaded, the spectrum will be empty.");
                return;
            }
            _cursor = _source.AddCursor(0, false);
            Melon<RenderMod>.Logger.Msg($"[AudioSpectrumGenerator] Song '{clip.name}': {_source.Channels} ch at {rate} Hz ({(_ownStore != null ? "own store" : "shared with the mixer")})");
        }
        catch (Exception e)
        {
            Melon<RenderMod>.Logger.Error($"[AudioSpectrumGenerator] Failed to initialize: {e}");
            Clear();
        }
    }

    /// <summary>Releases the chunk, the cursor and the own store. Call before the mixer store is disposed.</summary>
    public static void Clear()
    {
        if (_source != null)
        {
            _source.Release(_view);
            if (_cursor != null) _source.RemoveCursor(_cursor);
        }
        _view = default;
        _cursor = null;
        _source = null;
        _store = null;
        _ownStore?.Dispose();
        _ownStore = null;
        _loadErrorLogged = false;
        _cachedFrame = -1;
        CachedSpectrums.Clear();
    }

    /// <summary>
    /// Fills <paramref name="samples"/> with the magnitude spectrum of the song at the current render time
    /// (called by the AudioSource.GetSpectrumData patch while capturing with ProcessAudio on, for any AudioSource).
    /// Bin k covers frequency k * outputRate / (2 * samples.Length) Hz in output time; values are 0 ~ 1.
    /// Zeros are written when the song is not available, the conductor has not started the song, or the length is not a power of two.
    /// <paramref name="window"/> is not read: the Blackman-Harris window is always used (every GetSpectrumData call in the game
    /// passes FFTWindow.BlackmanHarris). Main thread only.
    /// </summary>
    /// <param name="samples">Destination; its length is the spectrum length (power of two).</param>
    /// <param name="channel">Requested channel (0 = left, 1 = right), clamped to the song's channels; ignored for mono songs.</param>
    /// <param name="window">Requested window (not used).</param>
    public static void GetSpectrumData(float[] samples, int channel, FFTWindow window)
    {
        var n = samples.Length;
        var conductor = scrConductor.instance;
        if (_source == null || _store == null || !conductor || !conductor!.hasSongStarted || !conductor.song
            || n == 0 || (n & (n - 1)) != 0)
        {
            Array.Clear(samples, 0, n);
            return;
        }

        var render = RenderManager.Instance;
        var currentFrame = render.frameCount;
        if (currentFrame != _cachedFrame)
        {
            _cachedFrame = currentFrame;
            CachedSpectrums.Clear();
        }

        var sourceChannel = _source.Channels == 1 ? 0 : Math.Min(Math.Max(channel, 0), _source.Channels - 1);
        if (CachedSpectrums.TryGetValue((n, sourceChannel), out var cached))
        {
            Array.Copy(cached, samples, n);
            return;
        }

        // Song position (seconds of song time) at the current render time, as in scrConductor.StartMusicCo:
        // the song starts at dspTimeSong + separate countdown (crotchetAtStart * adjustedCountdownTicks in song time).
        double pitch = conductor.song.pitch;
        double volume = conductor.song.volume;
        var countdown = conductor.separateCountdownTime ? conductor.crotchetAtStart * conductor.adjustedCountdownTicks : 0.0;
        var now = render.dspTimeBase + render.AudioTime;
        var songSeconds = (now - conductor.dspTimeSong) * pitch - countdown;

        // Window of 2n output samples centered on the current time: output sample i reads source frame first + i * pitch,
        // so the window starts n output samples (n * pitch source frames) before the song position.
        var fftSize = 2 * n;
        var gain = volume * (_source.Channels == 1 ? MonoGain : 1.0);
        var first = songSeconds * _store.OutputRate - n * pitch;

        var windowSamples = new float[fftSize];
        try
        {
            // Keep the cursor at the window start so the decoder thread loads (and keeps) the chunks the window reads.
            _cursor!.Position = Math.Min(Math.Max((long)first, 0), _source.TotalFrames - 1);
            for (var i = 0; i < fftSize; i++)
            {
                var position = first + i * pitch;
                var index = (long)Math.Floor(position);
                var fraction = position - index;
                var s0 = ReadSample(index, sourceChannel);
                var s1 = fraction > 0 ? ReadSample(index + 1, sourceChannel) : 0.0;
                windowSamples[i] = (float)((s0 + (s1 - s0) * fraction) * gain);
            }
        }
        catch (InvalidOperationException e)
        {
            if (!_loadErrorLogged)
            {
                _loadErrorLogged = true;
                Melon<RenderMod>.Logger.Error($"[AudioSpectrumGenerator] Song samples are unavailable, the spectrum will be empty: {e.Message}");
            }
            Array.Clear(samples, 0, n);
            return;
        }

        var spectrum = new float[n];
        CalculateFFT(windowSamples, spectrum);
        CachedSpectrums[(n, sourceChannel)] = spectrum;
        Array.Copy(spectrum, samples, n);
    }

    /// <summary>Sample of source channel <paramref name="channel"/> at <paramref name="frame"/>; 0 outside the clip or when the source failed.</summary>
    /// <exception cref="InvalidOperationException">Chunk loading failed fatally (see <see cref="AudioClipStore.Acquire"/>).</exception>
    private static double ReadSample(long frame, int channel)
    {
        var source = _source!;
        if (frame < 0 || frame >= source.TotalFrames) return 0.0;
        if (!_view.IsValid || !_view.Contains(frame))
        {
            source.Release(_view);
            _view = default;
            _view = _store!.Acquire(source, frame);
            if (!_view.IsValid) return 0.0;
        }
        return _view.Data[(int)(frame - _view.StartFrame) * _view.Channels + channel];
    }

    /// <summary>
    /// Magnitude spectrum of <paramref name="samples"/> (length n, a power of two) after the Blackman-Harris window:
    /// writes min(n / 2, spectrum.Length) bins, bin k = min(|X[k]| * 2.5 / n, 1).
    /// </summary>
    private static void CalculateFFT(float[] samples, float[] spectrum)
    {
        var n = samples.Length;

        // Blackman-Harris window, periodic form: w(i) = a0 - a1 cos(2πi/n) + a2 cos(4πi/n) - a3 cos(6πi/n), i = 0 ~ n - 1
        const float a0 = 0.35875f;
        const float a1 = 0.48829f;
        const float a2 = 0.14128f;
        const float a3 = 0.01168f;

        var re = new float[n];
        var im = new float[n];

        for (var i = 0; i < n; i++)
        {
            var angle = 2.0 * Math.PI * i / n;
            var w = (float)(a0 - a1 * Math.Cos(angle) + a2 * Math.Cos(2.0 * angle) - a3 * Math.Cos(3.0 * angle));
            re[i] = samples[i] * w;
            im[i] = 0f;
        }

        // Cooley-Tukey Radix-2 FFT (iterative, in place, decimation in time); m = log2(n) stages.
        var m = (int)Math.Round(Math.Log(n, 2));

        // Bit-reversal: reorder the input so element i moves to the index with the bits of i reversed (j tracks that index).
        var j = 0;
        for (var i = 0; i < n - 1; i++)
        {
            if (i < j)
            {
                var temp = re[i];
                re[i] = re[j];
                re[j] = temp;

                temp = im[i];
                im[i] = im[j];
                im[j] = temp;
            }
            var k = n / 2;
            while (k <= j)
            {
                j -= k;
                k /= 2;
            }
            j += k;
        }

        // FFT stages: stage s combines pairs of transforms of length le/2 into length le = 2^s with butterflies.
        // (ur, ui) is the twiddle factor exp(-i * pi * sub / le2), advanced by multiplying with (wr, wi) = exp(-i * pi / le2).
        for (var stage = 1; stage <= m; stage++)
        {
            var le = 1 << stage;
            var le2 = le / 2;
            var ur = 1f;
            var ui = 0f;
            var wr = (float)Math.Cos(Math.PI / le2);
            var wi = -(float)Math.Sin(Math.PI / le2);

            for (var sub = 0; sub < le2; sub++)
            {
                for (var item = sub; item < n; item += le)
                {
                    var ip = item + le2;
                    var tr = re[ip] * ur - im[ip] * ui;
                    var ti = re[ip] * ui + im[ip] * ur;

                    re[ip] = re[item] - tr;
                    im[ip] = im[item] - ti;
                    re[item] += tr;
                    im[item] += ti;
                }
                var tempUr = ur;
                ur = tempUr * wr - ui * wi;
                ui = tempUr * wi + ui * wr;
            }
        }

        // Magnitude scale 2.5 / n (n = FFT size, twice the spectrum length), clamped to 1.0: the scale of Unity GetSpectrumData.
        // A sine of amplitude A centered on a bin gives A * a0 * 1.25 with the Blackman-Harris window.
        var factor = 2.5f / n;
        var outLen = Math.Min(spectrum.Length, n / 2);
        for (var i = 0; i < outLen; i++)
        {
            spectrum[i] = Math.Min((float)Math.Sqrt(re[i] * re[i] + im[i] * im[i]) * factor, 1f);
        }
    }
}