using Microsoft.ML.OnnxRuntime;
using Microsoft.ML.OnnxRuntime.Tensors;

namespace WakewordProbe;

// Streaming wakeword inference port of openwakeword 0.6.0's predict pipeline.
// References (read alongside this file):
//   openwakeword/utils.py:AudioFeatures._streaming_features  (mel + embedding stages)
//   openwakeword/utils.py:AudioFeatures._streaming_melspectrogram
//   openwakeword/utils.py:AudioFeatures._get_embeddings
//   openwakeword/model.py:Model.predict                       (classifier stage)
internal sealed class WakewordModel : IDisposable
{
    // Audio I/O geometry — per-call input contract.
    public const int FrameSamples = 1280;       // 80 ms @ 16 kHz; each Predict() call.
    public const int SampleRate = 16_000;

    // Mel model: input is last (FrameSamples + MelContextSamples) raw samples
    // -- the +480 is openwakeword's `-n_samples-160*3:` slice (3 hops of context).
    public const int MelContextSamples = 480;          // 160 * 3
    public const int MelInputSamples = FrameSamples + MelContextSamples; // 1760
    public const int MelBins = 32;                     // melspectrogram model output dim
    public const int MelBufferMaxFrames = 970;         // 10 * 97 (openwakeword `melspectrogram_max_len`)

    // Embedding model: 76-mel-frame window in, 96-d embedding out.
    public const int EmbeddingWindowMelFrames = 76;
    public const int EmbeddingDim = 96;
    public const int EmbeddingBufferMax = 120;         // openwakeword `feature_buffer_max_len`
    public const string EmbeddingInputName = "input_1"; // openwakeword convention; assert at startup

    // Classifier: 16 embeddings in, scalar score out.
    public const int ClassifierEmbeddings = 16;

    // Skip the first N Predict() calls — buffer fill-up window.
    public const int WarmupFrames = ClassifierEmbeddings; // 16 frames ≈ 1.28 s

    private readonly InferenceSession _mel;
    private readonly InferenceSession _emb;
    private readonly InferenceSession _cls;

    private readonly string _melInputName;   // discovered at startup
    private readonly string _clsInputName;   // discovered at startup

    private readonly short[] _rawRing = new short[MelInputSamples];
    private int _rawRingFill = 0;                                      // samples buffered (≤ MelInputSamples)

    private readonly List<float[]> _melRing = new(MelBufferMaxFrames); // each entry is a length-32 mel frame
    private readonly List<float[]> _embRing = new(EmbeddingBufferMax); // each entry is a length-96 embedding

    private int _framesSeen = 0;

    public WakewordModel(string melPath, string embeddingPath, string classifierPath)
    {
        var opts = new SessionOptions
        {
            IntraOpNumThreads = 1,
            InterOpNumThreads = 1,
            LogSeverityLevel = OrtLoggingLevel.ORT_LOGGING_LEVEL_ERROR,
        };

        _mel = new InferenceSession(melPath, opts);
        _emb = new InferenceSession(embeddingPath, opts);
        _cls = new InferenceSession(classifierPath, opts);

        // Discover input names + assert shape geometry.
        _melInputName = _mel.InputMetadata.Keys.Single();

        AssertEmbeddingShape(_emb);   // input_1: [batch, 76, 32, 1], dtype float
        _clsInputName = _cls.InputMetadata.Keys.Single();
        AssertClassifierShape(_cls);  // [batch, 16, 96], dtype float
    }

    private static void AssertEmbeddingShape(InferenceSession sess)
    {
        if (!sess.InputMetadata.TryGetValue(EmbeddingInputName, out var meta))
            throw new InvalidOperationException(
                $"embedding_model.onnx: expected input named '{EmbeddingInputName}', got [{string.Join(",", sess.InputMetadata.Keys)}]");
        var d = meta.Dimensions;
        // Expected: [batch, 76, 32, 1] — batch may be -1 (dynamic).
        if (d.Length != 4 || d[1] != EmbeddingWindowMelFrames || d[2] != MelBins || d[3] != 1)
            throw new InvalidOperationException(
                $"embedding_model.onnx: expected input shape [batch,{EmbeddingWindowMelFrames},{MelBins},1], got [{string.Join(",", d)}]");
        if (meta.ElementType != typeof(float))
            throw new InvalidOperationException($"embedding_model.onnx: expected Single input, got {meta.ElementType.Name}");
    }

    private static void AssertClassifierShape(InferenceSession sess)
    {
        var inputName = sess.InputMetadata.Keys.Single();
        var meta = sess.InputMetadata[inputName];
        var d = meta.Dimensions;
        // Expected: [batch, 16, 96] — batch may be -1.
        if (d.Length != 3 || d[1] != ClassifierEmbeddings || d[2] != EmbeddingDim)
            throw new InvalidOperationException(
                $"alexa.onnx: expected input shape [batch,{ClassifierEmbeddings},{EmbeddingDim}], got [{string.Join(",", d)}]");
        if (meta.ElementType != typeof(float))
            throw new InvalidOperationException($"alexa.onnx: expected Single input, got {meta.ElementType.Name}");
    }

    public float Predict(short[] frame1280)
    {
        if (frame1280.Length != FrameSamples)
            throw new ArgumentException($"Expected {FrameSamples} samples, got {frame1280.Length}");

        // 1. Append 1280 new samples to the raw ring (shift older samples down if full).
        if (_rawRingFill < MelInputSamples)
        {
            int copyToFront = Math.Min(MelInputSamples - _rawRingFill, FrameSamples);
            Array.Copy(frame1280, 0, _rawRing, _rawRingFill, copyToFront);
            _rawRingFill += copyToFront;
            if (copyToFront < FrameSamples)
            {
                // Boundary case: ring was partially full and the new frame overshoots
                // remaining capacity. Fires exactly once during warm-up (typically call 2,
                // when _rawRingFill = 1280 and the incoming 1280 samples overshoot the
                // remaining 480 capacity). Shift the older samples left to make room,
                // then write the leftover at the tail.
                int leftover = FrameSamples - copyToFront;
                Array.Copy(_rawRing, leftover, _rawRing, 0, MelInputSamples - leftover);
                Array.Copy(frame1280, copyToFront, _rawRing, MelInputSamples - leftover, leftover);
            }
        }
        else
        {
            // Shift older samples left by FrameSamples, then append new at the tail.
            Array.Copy(_rawRing, FrameSamples, _rawRing, 0, MelInputSamples - FrameSamples);
            Array.Copy(frame1280, 0, _rawRing, MelInputSamples - FrameSamples, FrameSamples);
        }

        // Skip everything until we have the full mel-context window primed.
        if (_rawRingFill < MelInputSamples)
        {
            _framesSeen++;
            return 0f;
        }

        // 2. Mel stage: feed the full _rawRing as float32 (1, MelInputSamples) into mel model.
        var melInputData = new float[MelInputSamples];
        for (int i = 0; i < MelInputSamples; i++) melInputData[i] = _rawRing[i]; // int16 → float32, NO normalisation
        var melInputTensor = new DenseTensor<float>(melInputData, new[] { 1, MelInputSamples });
        using var melResults = _mel.Run(new[] {
            NamedOnnxValue.CreateFromTensor(_melInputName, melInputTensor)
        });
        var melOut = melResults.First().AsTensor<float>(); // shape (1, 1, n_frames, 32)

        // Apply openwakeword's `x / 10 + 2` transform and append each new frame to _melRing.
        // FIXED: actual mel output shape is (1, 1, n_frames, 32) — n_frames is at Dimensions[2], NOT [1].
        int nFrames = melOut.Dimensions[2];
        for (int f = 0; f < nFrames; f++)
        {
            var bin = new float[MelBins];
            for (int b = 0; b < MelBins; b++)
                bin[b] = melOut[0, 0, f, b] / 10f + 2f;
            _melRing.Add(bin);
        }
        if (_melRing.Count > MelBufferMaxFrames)
            _melRing.RemoveRange(0, _melRing.Count - MelBufferMaxFrames);

        // 3. Embedding stage: need ≥ 76 mel frames; slice the last 76 → (1, 76, 32, 1).
        if (_melRing.Count < EmbeddingWindowMelFrames)
        {
            _framesSeen++;
            return 0f;
        }
        var embInputData = new float[EmbeddingWindowMelFrames * MelBins];
        int startMel = _melRing.Count - EmbeddingWindowMelFrames;
        for (int f = 0; f < EmbeddingWindowMelFrames; f++)
            Array.Copy(_melRing[startMel + f], 0, embInputData, f * MelBins, MelBins);
        var embInputTensor = new DenseTensor<float>(embInputData, new[] { 1, EmbeddingWindowMelFrames, MelBins, 1 });
        using var embResults = _emb.Run(new[] {
            NamedOnnxValue.CreateFromTensor(EmbeddingInputName, embInputTensor)
        });
        var embOut = embResults.First().AsTensor<float>(); // shape (1, 1, 1, 96)
        var newEmb = new float[EmbeddingDim];
        for (int i = 0; i < EmbeddingDim; i++) newEmb[i] = embOut[0, 0, 0, i];
        _embRing.Add(newEmb);
        if (_embRing.Count > EmbeddingBufferMax)
            _embRing.RemoveRange(0, _embRing.Count - EmbeddingBufferMax);

        _framesSeen++;

        // 4. Warm-up + classifier stage.
        if (_embRing.Count < ClassifierEmbeddings || _framesSeen <= WarmupFrames)
            return 0f;

        var clsInputData = new float[ClassifierEmbeddings * EmbeddingDim];
        int startEmb = _embRing.Count - ClassifierEmbeddings;
        for (int i = 0; i < ClassifierEmbeddings; i++)
            Array.Copy(_embRing[startEmb + i], 0, clsInputData, i * EmbeddingDim, EmbeddingDim);
        var clsInputTensor = new DenseTensor<float>(clsInputData, new[] { 1, ClassifierEmbeddings, EmbeddingDim });
        using var clsResults = _cls.Run(new[] {
            NamedOnnxValue.CreateFromTensor(_clsInputName, clsInputTensor)
        });
        var clsOut = clsResults.First().AsTensor<float>(); // shape (1, 1)
        return clsOut[0, 0];
    }

    public void Dispose()
    {
        _mel.Dispose();
        _emb.Dispose();
        _cls.Dispose();
    }
}