// OG Sound Remaster - RIFF/WAVE decoder
// Author: StixsworldHD (StixsmasterHD4k)
//
// Supports: PCM u8 / s16 / s24 / s32, IEEE float32 / float64,
//           IMA (DVI) ADPCM 4-bit, and WAVE_FORMAT_EXTENSIBLE wrappers.
// Output is always interleaved float32 in [-1, 1], which is exactly what
// UnityEngine.AudioClip.SetData expects.
//
// The decoder is allocation-lean: one managed float[] for the result and
// nothing else on the hot path. All parsing is done over a byte[] the caller
// owns, so the file buffer can be pooled and reused.

using System;

namespace OGSoundRemaster.Audio
{
    /// <summary>Fully decoded linear PCM, interleaved, float32.</summary>
    public sealed class PcmData
    {
        public float[] Samples;   // interleaved, length == Frames * Channels
        public int Frames;        // sample frames (per channel)
        public int Channels;
        public int SampleRate;

        public long ByteSize => (long)Frames * Channels * sizeof(float);
        public float Duration => SampleRate > 0 ? (float)Frames / SampleRate : 0f;
    }

    public struct WavFormat
    {
        public int FormatTag;
        public int Channels;
        public int SampleRate;
        public int BlockAlign;
        public int BitsPerSample;
        public int SamplesPerBlock;   // IMA ADPCM only
        public int DataOffset;
        public int DataLength;
    }

    public static class WavDecoder
    {
        public const int FMT_PCM = 0x0001;
        public const int FMT_ADPCM_MS = 0x0002;
        public const int FMT_IEEE_FLOAT = 0x0003;
        public const int FMT_ADPCM_IMA = 0x0011;
        public const int FMT_EXTENSIBLE = 0xFFFE;

        // ---- IMA ADPCM tables (ITU / IMA reference) -------------------------

        private static readonly sbyte[] IndexTable =
        {
            -1, -1, -1, -1, 2, 4, 6, 8,
            -1, -1, -1, -1, 2, 4, 6, 8
        };

        private static readonly short[] StepTable =
        {
            7,     8,     9,     10,    11,    12,    13,    14,    16,    17,
            19,    21,    23,    25,    28,    31,    34,    37,    41,    45,
            50,    55,    60,    66,    73,    80,    88,    97,    107,   118,
            130,   143,   157,   173,   190,   209,   230,   253,   279,   307,
            337,   371,   408,   449,   494,   544,   598,   658,   724,   796,
            876,   963,   1060,  1166,  1282,  1411,  1552,  1707,  1878,  2066,
            2272,  2499,  2749,  3024,  3327,  3660,  4026,  4428,  4871,  5358,
            5894,  6484,  7132,  7845,  8630,  9493,  10442, 11487, 12635, 13899,
            15289, 16818, 18500, 20350, 22385, 24623, 27086, 29794, 32767
        };

        // ---- Header parsing -------------------------------------------------

        private static uint RdU32(byte[] b, int o) =>
            (uint)(b[o] | (b[o + 1] << 8) | (b[o + 2] << 16) | (b[o + 3] << 24));

        private static int RdU16(byte[] b, int o) => b[o] | (b[o + 1] << 8);

        private static short RdS16(byte[] b, int o) => unchecked((short)(b[o] | (b[o + 1] << 8)));

        private static bool Tag(byte[] b, int o, char a, char c, char d, char e) =>
            o + 3 < b.Length && b[o] == (byte)a && b[o + 1] == (byte)c && b[o + 2] == (byte)d && b[o + 3] == (byte)e;

        /// <summary>Parse the RIFF header only. Cheap enough to call on every file at scan time.</summary>
        public static bool TryReadFormat(byte[] file, int fileLength, out WavFormat fmt, out string error)
            => TryReadFormat(file, fileLength, fileLength, out fmt, out error);

        /// <summary>
        /// Header-only parse where the buffer may hold just the first slice of the file.
        /// <paramref name="available"/> is how many bytes are actually in the buffer;
        /// <paramref name="totalLength"/> is the real size of the file on disk. Keeping these
        /// separate is what lets the bank index 500+ files by reading only their headers while
        /// still recording correct durations for multi-megabyte music tracks.
        /// </summary>
        public static bool TryReadFormat(byte[] file, int available, long totalLength, out WavFormat fmt, out string error)
        {
            fmt = default;
            error = null;

            int fileLength = available;
            if (totalLength < available) totalLength = available;

            if (file == null || fileLength < 44) { error = "file shorter than a minimal RIFF header"; return false; }
            if (!Tag(file, 0, 'R', 'I', 'F', 'F')) { error = "missing RIFF magic"; return false; }
            if (!Tag(file, 8, 'W', 'A', 'V', 'E')) { error = "missing WAVE magic"; return false; }

            bool haveFmt = false, haveData = false;
            int pos = 12;

            while (pos + 8 <= fileLength)
            {
                int id = pos;
                long size = RdU32(file, pos + 4);
                int body = pos + 8;

                // A data chunk routinely extends past the scan window; only clamp against the real
                // file size, never against how much happens to be buffered.
                if (size < 0 || body + size > totalLength)
                    size = totalLength - body;             // tolerate truncated / mis-sized final chunk
                if (size < 0) size = 0;

                if (Tag(file, id, 'f', 'm', 't', ' ') && size >= 16)
                {
                    fmt.FormatTag = RdU16(file, body + 0);
                    fmt.Channels = RdU16(file, body + 2);
                    fmt.SampleRate = (int)RdU32(file, body + 4);
                    fmt.BlockAlign = RdU16(file, body + 12);
                    fmt.BitsPerSample = RdU16(file, body + 14);

                    int cbSize = size >= 18 ? RdU16(file, body + 16) : 0;

                    if (fmt.FormatTag == FMT_EXTENSIBLE && size >= 40 && cbSize >= 22)
                    {
                        // SubFormat GUID's first two bytes carry the real format tag.
                        fmt.FormatTag = RdU16(file, body + 24);
                        int valid = RdU16(file, body + 18);
                        if (valid > 0) fmt.BitsPerSample = valid;
                    }
                    else if (fmt.FormatTag == FMT_ADPCM_IMA && size >= 20 && cbSize >= 2)
                    {
                        fmt.SamplesPerBlock = RdU16(file, body + 18);
                    }

                    haveFmt = true;
                }
                else if (Tag(file, id, 'd', 'a', 't', 'a'))
                {
                    fmt.DataOffset = body;
                    fmt.DataLength = (int)size;
                    haveData = true;
                }

                long next = body + size;
                if ((size & 1) != 0) next++;               // RIFF chunks are word aligned
                if (next <= id || next > int.MaxValue) break;   // malformed; refuse to spin
                pos = (int)next;
            }

            if (!haveFmt) { error = "no fmt chunk"; return false; }
            if (!haveData) { error = "no data chunk"; return false; }
            if (fmt.Channels < 1 || fmt.Channels > 8) { error = "unsupported channel count " + fmt.Channels; return false; }
            if (fmt.SampleRate < 1000 || fmt.SampleRate > 384000) { error = "unsupported sample rate " + fmt.SampleRate; return false; }

            if (fmt.FormatTag == FMT_ADPCM_IMA)
            {
                if (fmt.BlockAlign < 4 * fmt.Channels) { error = "IMA block align too small"; return false; }
                if (fmt.SamplesPerBlock <= 0)
                    fmt.SamplesPerBlock = ((fmt.BlockAlign - 4 * fmt.Channels) * 8) / (4 * fmt.Channels) + 1;
            }

            return true;
        }

        /// <summary>How many sample frames the file will decode to. Header-only, no decode.</summary>
        public static int FrameCount(in WavFormat f)
        {
            switch (f.FormatTag)
            {
                case FMT_ADPCM_IMA:
                {
                    int full = f.DataLength / f.BlockAlign;
                    int tail = f.DataLength - full * f.BlockAlign;
                    int frames = full * f.SamplesPerBlock;
                    if (tail >= 4 * f.Channels)
                        frames += 1 + ((tail - 4 * f.Channels) * 8) / (4 * f.Channels);
                    return frames;
                }
                default:
                {
                    int bytesPerFrame = (f.BitsPerSample / 8) * f.Channels;
                    return bytesPerFrame > 0 ? f.DataLength / bytesPerFrame : 0;
                }
            }
        }

        // ---- Decode ---------------------------------------------------------

        public static bool TryDecode(byte[] file, int fileLength, out PcmData pcm, out string error)
        {
            pcm = null;
            if (!TryReadFormat(file, fileLength, out WavFormat f, out error))
                return false;

            int frames = FrameCount(in f);
            if (frames <= 0) { error = "decodes to zero frames"; return false; }

            float[] outBuf;
            try { outBuf = new float[(long)frames * f.Channels]; }
            catch (OutOfMemoryException) { error = "out of memory allocating " + ((long)frames * f.Channels * 4) + " bytes"; return false; }

            bool ok;
            switch (f.FormatTag)
            {
                case FMT_PCM: ok = DecodePcm(file, in f, outBuf, frames, out error); break;
                case FMT_IEEE_FLOAT: ok = DecodeFloat(file, in f, outBuf, frames, out error); break;
                case FMT_ADPCM_IMA: ok = DecodeImaAdpcm(file, in f, outBuf, frames, out error); break;
                default:
                    error = "unsupported WAVE format tag 0x" + f.FormatTag.ToString("X4");
                    return false;
            }

            if (!ok) return false;

            pcm = new PcmData
            {
                Samples = outBuf,
                Frames = frames,
                Channels = f.Channels,
                SampleRate = f.SampleRate
            };
            return true;
        }

        private static bool DecodePcm(byte[] src, in WavFormat f, float[] dst, int frames, out string error)
        {
            error = null;
            int ch = f.Channels;
            int n = frames * ch;
            int o = f.DataOffset;

            switch (f.BitsPerSample)
            {
                case 8:
                    // 8-bit WAV PCM is unsigned, biased at 128.
                    for (int i = 0; i < n; i++)
                        dst[i] = (src[o + i] - 128) * (1f / 128f);
                    return true;

                case 16:
                    for (int i = 0; i < n; i++)
                        dst[i] = RdS16(src, o + (i << 1)) * (1f / 32768f);
                    return true;

                case 24:
                    for (int i = 0; i < n; i++)
                    {
                        int p = o + i * 3;
                        int v = (src[p] << 8) | (src[p + 1] << 16) | (src[p + 2] << 24); // sign-extend via top byte
                        dst[i] = (v >> 8) * (1f / 8388608f);
                    }
                    return true;

                case 32:
                    for (int i = 0; i < n; i++)
                        dst[i] = unchecked((int)RdU32(src, o + (i << 2))) * (1f / 2147483648f);
                    return true;

                default:
                    error = "unsupported PCM bit depth " + f.BitsPerSample;
                    return false;
            }
        }

        private static bool DecodeFloat(byte[] src, in WavFormat f, float[] dst, int frames, out string error)
        {
            error = null;
            int n = frames * f.Channels;
            int o = f.DataOffset;

            if (f.BitsPerSample == 32)
            {
                Buffer.BlockCopy(src, o, dst, 0, n * 4);
                return true;
            }
            if (f.BitsPerSample == 64)
            {
                for (int i = 0; i < n; i++)
                    dst[i] = (float)BitConverter.Int64BitsToDouble(
                        (long)RdU32(src, o + (i << 3)) | ((long)RdU32(src, o + (i << 3) + 4) << 32));
                return true;
            }
            error = "unsupported IEEE float bit depth " + f.BitsPerSample;
            return false;
        }

        /// <summary>
        /// Expand one 4-bit IMA nibble. Uses the multiply form
        /// diff = ((2*delta + 1) * step) >> 3, which is what FFmpeg, FMOD and
        /// every modern decoder use, so output is bit-identical to a reference decode.
        /// </summary>
        private static int Expand(int nibble, ref int predictor, ref int index)
        {
            int step = StepTable[index];

            index += IndexTable[nibble];
            if (index < 0) index = 0;
            else if (index > 88) index = 88;

            int diff = (((nibble & 7) * 2 + 1) * step) >> 3;

            if ((nibble & 8) != 0) predictor -= diff;
            else predictor += diff;

            if (predictor > 32767) predictor = 32767;
            else if (predictor < -32768) predictor = -32768;

            return predictor;
        }

        private static bool DecodeImaAdpcm(byte[] src, in WavFormat f, float[] dst, int frames, out string error)
        {
            error = null;

            int ch = f.Channels;
            int blockAlign = f.BlockAlign;
            int spb = f.SamplesPerBlock;
            int dataEnd = f.DataOffset + f.DataLength;

            var predictor = new int[ch];
            var index = new int[ch];

            const float SCALE = 1f / 32768f;

            int outFrame = 0;

            for (int blockStart = f.DataOffset; blockStart < dataEnd && outFrame < frames; blockStart += blockAlign)
            {
                int blockEnd = blockStart + blockAlign;
                if (blockEnd > dataEnd) blockEnd = dataEnd;
                if (blockEnd - blockStart < 4 * ch) break;         // runt block, nothing usable

                // --- per-channel block preamble: int16 predictor, uint8 step index, uint8 reserved
                int p = blockStart;
                for (int c = 0; c < ch; c++)
                {
                    predictor[c] = RdS16(src, p);
                    int ix = src[p + 2];
                    index[c] = ix > 88 ? 88 : ix;
                    p += 4;

                    // The preamble predictor IS the block's first output sample.
                    dst[outFrame * ch + c] = predictor[c] * SCALE;
                }

                int produced = 1;

                // --- nibble payload: groups of 4 bytes (8 samples) per channel, round robin
                while (p + 4 * ch <= blockEnd && produced < spb && outFrame + produced < frames)
                {
                    for (int c = 0; c < ch; c++)
                    {
                        int frame = produced;
                        int pred = predictor[c];
                        int idx = index[c];

                        for (int k = 0; k < 4; k++)
                        {
                            byte b = src[p++];

                            if (frame < spb && outFrame + frame < frames)
                                dst[(outFrame + frame) * ch + c] = Expand(b & 0x0F, ref pred, ref idx) * SCALE;
                            frame++;

                            if (frame < spb && outFrame + frame < frames)
                                dst[(outFrame + frame) * ch + c] = Expand(b >> 4, ref pred, ref idx) * SCALE;
                            frame++;
                        }

                        predictor[c] = pred;
                        index[c] = idx;
                    }

                    produced += 8;
                }

                outFrame += produced < spb ? produced : spb;
            }

            return true;
        }
    }
}
