using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Net;
using System.Net.Sockets;
using System.Text;
using System.Text.Json;
using System.Threading;
using Microsoft.ML.OnnxRuntime;
using Microsoft.ML.OnnxRuntime.Tensors;

namespace LoopbackCapture
{
    // Immutable publication; never expose worker-owned/reused inference buffers.
    sealed class GenrePrediction
    {
        internal readonly string Label, Source;
        // Genere specifico dal catalogo ("Hard Rock"), vuoto per regole/ONNX.
        internal readonly string Detail;
        internal readonly float[] Probabilities;
        internal readonly long ExpiresMs;
        internal float Confidence { get { return Probabilities.Max(); } }
        internal GenrePrediction(float[] probabilities, string source, long expires, string detail = null)
        {
            Probabilities = probabilities; Source = source; ExpiresMs = expires; Detail = detail;
            int best = 0;
            for (int i = 1; i < probabilities.Length; i++) if (probabilities[i] > probabilities[best]) best = i;
            Label = probabilities[best] < 0.30f ? "unknown" : GenreOnnxProvider.Labels[best];
        }
    }

    sealed class GenreOnnxProvider
    {
        internal static readonly string[] Labels = { "latin", "electronic", "hiphop", "rock", "classical" };
        internal const int FftSize = 2048, MelBands = 128, FrameCount = 300;
        const int InferenceMs = 1000, ResultTtlMs = 3000;
        readonly int rate;
        readonly string dumpPath;
        readonly ManualResetEventSlim initialized = new ManualResetEventSlim(false);
        MelFrontend frontend;
        GenrePrediction prediction;
        int disabled, dumpWritten;
        string status = "warming";
        string failureReason = "";
        internal bool Failed { get { return Volatile.Read(ref disabled) != 0; } }
        internal bool DumpWritten { get { return Volatile.Read(ref dumpWritten) != 0; } }
        // Diagnostic status: off | no-model | warming | ok | failed. Read by the DIAG:
        // packet without locks (string reference assignment is atomic).
        internal string Status { get { return Volatile.Read(ref status); } }
        internal string FailureReason { get { return Volatile.Read(ref failureReason); } }

        void SetStatus(string s) { Volatile.Write(ref status, s); }
        void SetFailure(string s) { Volatile.Write(ref failureReason, s); }

        internal GenreOnnxProvider(int sampleRate, string dump)
        { rate = sampleRate; dumpPath = dump; }

        internal void Start()
        {
            try { new Thread(Worker) { IsBackground = true, Name = "Genre ONNX" }.Start(); }
            catch (Exception ex) { Disable(ex); initialized.Set(); }
        }

        void Disable(Exception ex)
        {
            Volatile.Write(ref prediction, null);
            if (Interlocked.Exchange(ref disabled, 1) == 0)
            {
                SetFailure(ex.Message);
                SetStatus("failed");
                Console.Error.WriteLine("ONNX disabled for this session: " + ex.Message +
                    (dumpPath != null && !DumpWritten ? " (tensor dump rejected: no valid tensor written)" : ""));
            }
        }

        internal bool WaitForInitialization(int milliseconds)
        { return initialized.Wait(milliseconds) && !Failed && Volatile.Read(ref frontend) != null; }

        // Called on audio thread: preallocated math only. No locks, I/O, or waits.
        internal void AddSample(float sample)
        {
            if (Failed) return;
            try { Volatile.Read(ref frontend)?.AddSample(sample); }
            catch (Exception ex) { Disable(ex); }
        }

        internal GenrePrediction GetPrediction()
        {
            if (Failed) return null;
            GenrePrediction value = Volatile.Read(ref prediction);
            return value != null && Environment.TickCount64 < value.ExpiresMs ? value : null;
        }

        void Worker()
        {
            try
            {
                string modelPath = Path.Combine(AppContext.BaseDirectory, "genre_model.onnx");
                string jsonPath = Path.Combine(AppContext.BaseDirectory, "genre_model.json");
                if (!File.Exists(modelPath) || !File.Exists(jsonPath))
                {
                    // Permanent: the model is absent, not a transient load failure.
                    SetFailure("genre_model.onnx / genre_model.json missing");
                    SetStatus("no-model");
                    Volatile.Write(ref disabled, 1);
                    return;
                }
                using var json = JsonDocument.Parse(File.ReadAllText(jsonPath));
                var root = json.RootElement;
                if (root.GetProperty("version").GetInt32() != 1 ||
                    !root.GetProperty("labels").EnumerateArray().Select(x => x.GetString()).SequenceEqual(Labels))
                    throw new InvalidDataException("Unsupported model version or label order");
                double mean = root.GetProperty("mean").GetDouble(), std = root.GetProperty("std").GetDouble();
                if (!double.IsFinite(mean) || !double.IsFinite(std) || std <= 0) throw new InvalidDataException("Invalid model mean/std");
                using var options = new SessionOptions { IntraOpNumThreads = 1, InterOpNumThreads = 1, ExecutionMode = ExecutionMode.ORT_SEQUENTIAL };
                using var session = new InferenceSession(modelPath, options);
                if (session.InputMetadata.Count != 1 || session.OutputMetadata.Count != 1) throw new InvalidDataException("Expected one model input and output");
                var input = session.InputMetadata.Single(); var output = session.OutputMetadata.Single();
                if (input.Value.ElementType != typeof(float) || !input.Value.Dimensions.SequenceEqual(new[] { 1, 1, MelBands, FrameCount }) ||
                    output.Value.ElementType != typeof(float) || !output.Value.Dimensions.SequenceEqual(new[] { 1, 5 }))
                    throw new InvalidDataException("Expected float32 [1,1,128,300] -> [1,5]");
                var localFrontend = new MelFrontend(rate, mean, std);
                var snapshot = new float[MelBands * FrameCount];
                var tensor = new DenseTensor<float>(snapshot, new[] { 1, 1, MelBands, FrameCount });
                var inputs = new[] { NamedOnnxValue.CreateFromTensor(input.Key, tensor) };
                Volatile.Write(ref frontend, localFrontend);
                initialized.Set();
                SetStatus("warming");
                long next = Environment.TickCount64, lastFrame = -1;
                while (!Failed)
                {
                    long wait = next - Environment.TickCount64;
                    if (wait > 0) Thread.Sleep((int)Math.Min(wait, InferenceMs)); // Worker only.
                    if (Failed) break;
                    next += InferenceMs;
                    if (localFrontend.TrySnapshot(snapshot, out long frame) && frame != lastFrame)
                    {
                        lastFrame = frame;
                        if (dumpPath != null && !DumpWritten)
                        {
                            // BinaryWriter writes IEEE float32 little-endian; exactly 153600 bytes.
                            // Write atomically: no empty/partial target if initialization or I/O fails.
                            string temporary = dumpPath + "." + Guid.NewGuid().ToString("N") + ".tmp";
                            try
                            {
                                using (var writer = new BinaryWriter(File.Open(temporary, FileMode.CreateNew)))
                                    foreach (float value in snapshot) writer.Write(value);
                                File.Move(temporary, dumpPath, true);
                                Volatile.Write(ref dumpWritten, 1);
                            }
                            finally { if (File.Exists(temporary)) File.Delete(temporary); }
                        }
                        using var result = session.Run(inputs);
                        float[] probabilities = result.Single().AsTensor<float>().ToArray();
                        double sum = probabilities.Sum(x => (double)x);
                        if (probabilities.Length != 5 || probabilities.Any(x => !float.IsFinite(x) || x < 0 || x > 1) || Math.Abs(sum - 1) > 0.001)
                            throw new InvalidDataException("Model output is not five probabilities summing to one");
                        if (!Failed) Volatile.Write(ref prediction, new GenrePrediction(probabilities, "onnx", Environment.TickCount64 + ResultTtlMs));
                        if (!Failed) SetStatus("ok");
                    }
                    // Skip missed slots. Never accumulate work when Run exceeds one second.
                    long now = Environment.TickCount64;
                    if (next <= now) next += ((now - next) / InferenceMs + 1) * InferenceMs;
                }
            }
            catch (Exception ex) { Disable(ex); }
            finally { initialized.Set(); }
        }

        sealed class MelFrontend
        {
            readonly int hop;
            readonly double mean, std;
            readonly double[] pcm = new double[FftSize], window = new double[FftSize];
            readonly double[] real = new double[FftSize], imag = new double[FftSize], power = new double[FftSize / 2 + 1];
            readonly double[] cos = new double[FftSize / 2], sin = new double[FftSize / 2];
            readonly int[] reverse = new int[FftSize], starts = new int[MelBands];
            readonly double[][] weights = new double[MelBands][];
            readonly float[] ring = new float[MelBands * FrameCount];
            int write, pending, sequence;
            long frames;

            internal MelFrontend(int rate, double mean, double std)
            {
                if (rate != 44100 && rate != 48000) throw new NotSupportedException("ONNX front-end supports 44100/48000 Hz");
                hop = rate / 100; this.mean = mean; this.std = std;
                for (int i = 0; i < FftSize; i++)
                {
                    window[i] = 0.5 - 0.5 * Math.Cos(2 * Math.PI * i / FftSize); // PERIODIC, independent of rules.
                    int x = i, r = 0;
                    for (int b = 0; b < 11; b++) { r = (r << 1) | (x & 1); x >>= 1; }
                    reverse[i] = r;
                }
                for (int i = 0; i < cos.Length; i++) { cos[i] = Math.Cos(-2 * Math.PI * i / FftSize); sin[i] = Math.Sin(-2 * Math.PI * i / FftSize); }
                double low = 2595 * Math.Log10(1 + 30.0 / 700), high = 2595 * Math.Log10(1 + 11000.0 / 700);
                var edges = new double[MelBands + 2];
                for (int i = 0; i < edges.Length; i++) edges[i] = 700 * (Math.Pow(10, (low + (high - low) * i / (MelBands + 1)) / 2595) - 1) * FftSize / rate;
                for (int m = 0; m < MelBands; m++)
                {
                    double left = edges[m], center = edges[m + 1], right = edges[m + 2];
                    starts[m] = Math.Max(0, (int)Math.Ceiling(left));
                    int end = Math.Min(FftSize / 2, (int)Math.Floor(right));
                    weights[m] = new double[Math.Max(0, end - starts[m] + 1)];
                    for (int j = 0; j < weights[m].Length; j++)
                    {
                        int k = starts[m] + j;
                        // Fractional edges; divide by FULL support width in bins (not half-area).
                        weights[m][j] = Math.Max(0, Math.Min((k - left) / (center - left), (right - k) / (right - center))) / (right - left);
                    }
                }
            }

            internal void AddSample(float sample)
            {
                if (!float.IsFinite(sample)) throw new InvalidDataException("Non-finite PCM");
                pcm[write] = sample; write = (write + 1) % FftSize;
                if (++pending < hop) return;
                pending = 0;
                // The initial zeroed ring supplies exactly N-hop left padding.
                for (int i = 0; i < FftSize; i++) { real[reverse[i]] = pcm[(write + i) % FftSize] * window[i]; imag[reverse[i]] = 0; }
                for (int n = 2; n <= FftSize; n <<= 1)
                    for (int start = 0; start < FftSize; start += n)
                        for (int j = 0; j < n / 2; j++)
                        {
                            int a = start + j, b = a + n / 2, t = j * (FftSize / n);
                            double tr = cos[t] * real[b] - sin[t] * imag[b], ti = sin[t] * real[b] + cos[t] * imag[b];
                            real[b] = real[a] - tr; imag[b] = imag[a] - ti; real[a] += tr; imag[a] += ti;
                        }
                for (int k = 0; k < power.Length; k++) power[k] = real[k] * real[k] + imag[k] * imag[k];
                Interlocked.Increment(ref sequence); // odd: producer updating frame. Never waits.
                int column = (int)(frames % FrameCount);
                for (int m = 0; m < MelBands; m++)
                {
                    double energy = 0;
                    for (int j = 0; j < weights[m].Length; j++) energy += weights[m][j] * power[starts[m] + j];
                    ring[m * FrameCount + column] = (float)((Math.Log10(Math.Max(energy, 1e-10)) - mean) / std);
                }
                frames++;
                Interlocked.Increment(ref sequence); // even: snapshot available.
            }

            internal bool TrySnapshot(float[] target, out long frame)
            {
                // Seqlock: only the worker retries; audio producer is never locked or blocked.
                frame = 0;
                for (int attempt = 0; attempt < 3; attempt++)
                {
                    int before = Volatile.Read(ref sequence);
                    if ((before & 1) != 0) continue;
                    long count = Interlocked.Read(ref frames);
                    if (count < FrameCount) return false;
                    int start = (int)(count % FrameCount);
                    for (int m = 0; m < MelBands; m++)
                    {
                        Array.Copy(ring, m * FrameCount + start, target, m * FrameCount, FrameCount - start);
                        Array.Copy(ring, m * FrameCount, target, m * FrameCount + FrameCount - start, start);
                    }
                    Thread.MemoryBarrier();
                    if (before == Volatile.Read(ref sequence)) { frame = count; return true; }
                }
                return false;
            }
        }

        // Diagnostic input bypasses WASAPI, but uses the SAME frontend/session worker.
        // --rate reinterprets samples (no resampling), taking the first 300*hop samples.
        internal static int VerifyWav(string path, int forcedRate, string dump, int port)
        {
            try
            {
                float[] pcm = ReadWav(path, out int rate);
                if (forcedRate != 0) rate = forcedRate;
                int count = (rate / 100) * FrameCount;
                if (pcm.Length < count) throw new InvalidDataException("Fixture too short for 300 frames at requested rate");
                var provider = new GenreOnnxProvider(rate, dump);
                provider.Start();
                if (!provider.WaitForInitialization(30000))
                {
                    if (!provider.Failed) provider.Disable(new TimeoutException("ONNX initialization timed out"));
                    return 5;
                }
                for (int i = 0; i < count; i++) provider.AddSample(pcm[i]);
                var clock = Stopwatch.StartNew();
                while (!provider.Failed && (provider.GetPrediction() == null || (dump != null && !provider.DumpWritten)) && clock.ElapsedMilliseconds < 30000) Thread.Sleep(10);
                if (provider.Failed) return 5;
                var result = provider.GetPrediction();
                if (result == null) throw new TimeoutException("ONNX verification inference timed out");
                var message = new StringBuilder();
                message.AppendFormat(System.Globalization.CultureInfo.InvariantCulture, "GENR:2;{0};{1:F3};onnx", result.Label, result.Confidence);
                for (int i = 0; i < 5; i++) message.AppendFormat(System.Globalization.CultureInfo.InvariantCulture, ";{0}={1:F3}", Labels[i], result.Probabilities[i]);
                using var udp = new UdpClient();
                byte[] bytes = Encoding.UTF8.GetBytes(message.ToString());
                udp.Send(bytes, bytes.Length, new IPEndPoint(IPAddress.Loopback, port));
                return 0;
            }
            catch (Exception ex) { Console.Error.WriteLine("ONNX verification rejected: " + ex.Message); return 5; }
        }

        static float[] ReadWav(string path, out int rate)
        {
            using var reader = new BinaryReader(File.OpenRead(path));
            if (new string(reader.ReadChars(4)) != "RIFF") throw new InvalidDataException("Expected RIFF WAV");
            reader.ReadUInt32();
            if (new string(reader.ReadChars(4)) != "WAVE") throw new InvalidDataException("Expected WAVE");
            int format = 0, channels = 0, bits = 0, align = 0; rate = 0; byte[] data = null;
            while (reader.BaseStream.Position + 8 <= reader.BaseStream.Length)
            {
                string id = new string(reader.ReadChars(4)); uint size = reader.ReadUInt32();
                long next = reader.BaseStream.Position + size + (size & 1);
                if (next > reader.BaseStream.Length + 1) throw new InvalidDataException("Truncated WAV");
                if (id == "fmt ")
                {
                    if (size < 16) throw new InvalidDataException("Invalid WAV fmt");
                    format = reader.ReadUInt16(); channels = reader.ReadUInt16(); rate = reader.ReadInt32();
                    reader.ReadUInt32(); align = reader.ReadUInt16(); bits = reader.ReadUInt16();
                    if (format == 65534 && size >= 40) { reader.ReadUInt16(); reader.ReadUInt16(); reader.ReadUInt32(); format = reader.ReadUInt16(); }
                }
                else if (id == "data") data = reader.ReadBytes(checked((int)size));
                reader.BaseStream.Position = next;
            }
            if (data == null || channels < 1 || (format != 1 && format != 3) || align != channels * (bits / 8) || align == 0 || data.Length % align != 0)
                throw new InvalidDataException("Unsupported WAV format");
            var samples = new float[data.Length / align];
            for (int i = 0; i < samples.Length; i++)
            {
                double sum = 0;
                for (int c = 0; c < channels; c++)
                {
                    int offset = i * align + c * (bits / 8);
                    if (format == 3 && bits == 32) sum += BitConverter.ToSingle(data, offset);
                    else if (format == 3 && bits == 64) sum += BitConverter.ToDouble(data, offset);
                    else if (format == 1 && bits == 8) sum += (data[offset] - 128) / 128.0;
                    else if (format == 1 && bits == 16) sum += BitConverter.ToInt16(data, offset) / 32768.0;
                    else if (format == 1 && bits == 24)
                    { int v = data[offset] | data[offset + 1] << 8 | data[offset + 2] << 16; if ((v & 0x800000) != 0) v |= unchecked((int)0xff000000); sum += v / 8388608.0; }
                    else if (format == 1 && bits == 32) sum += BitConverter.ToInt32(data, offset) / 2147483648.0;
                    else throw new InvalidDataException("Unsupported WAV sample encoding");
                }
                samples[i] = (float)(sum / channels);
            }
            return samples;
        }
    }
}