// %BANNER_BEGIN%
// ---------------------------------------------------------------------
// %COPYRIGHT_BEGIN%
// Copyright (c) 2022-2023 Magic Leap, Inc. All Rights Reserved.
// Use of this file is governed by the Magic Leap 2 Software License Agreement, located here: https://www.magicleap.com/software-license-agreement-ml2  
// Terms and conditions applicable to third-party materials accompanying this distribution may also be found in the top-level NOTICE file appearing herein.
// %COPYRIGHT_END%
// ---------------------------------------------------------------------
// %BANNER_END%

using System;
using System.Buffers;
using System.Collections.Concurrent;
using System.Runtime.InteropServices;
using System.Threading;
using System.Threading.Tasks;
using UnityEngine;
using UnityEngine.SceneManagement;
using static MagicLeap.Spectator.SpectatorMessages;

namespace MagicLeap.Spectator
{
    public class MLSpectatorAudio : MonoBehaviour
    {
        #region Private classes

        // Audio capture interface
        private interface IAudioCapture : IDisposable
        {
            public int SampleRate { get; }
            public int Channels { get; }
            public bool Active { get; }

            public void AddBufferAction(Action<ulong, float[], int, bool> action);
            public void RemoveBufferAction(Action<ulong, float[], int, bool> action);

            public void StartCapture();
            public void StopCapture();
        }

        // ML2 audio capture - not a monobehaviour
        private class DeviceCapture : IAudioCapture
        {
            private MLVirtualMic virtualMic = null;

            public int SampleRate => virtualMic.SampleRate;
            public int Channels => virtualMic.Channels;
            public bool Active => virtualMic.Active;

            public DeviceCapture()
            {
                virtualMic = new MLVirtualMic();
            }

            ~DeviceCapture()
            {
                Dispose();
            }

            public void AddBufferAction(Action<ulong, float[], int, bool> action)
            {
                virtualMic.OnBufferReceived += action;
            }

            public void RemoveBufferAction(Action<ulong, float[], int, bool> action)
            {
                virtualMic.OnBufferReceived -= action;
            }

            public void StartCapture() => virtualMic.Start();
            public void StopCapture() => virtualMic.Stop();

            public void Dispose()
            {
                if (virtualMic == null) return;
                virtualMic.Dispose();
                virtualMic = null;
            }
        }

        // Editor audio capture - place on a root object
        private class EditorManager : MonoBehaviour, IAudioCapture
        {
            [RequireComponent(typeof(AudioListener))]
            private class EditorCapture : MonoBehaviour
            {
                private class CopyBuffer<T> : IDisposable
                {
                    private static ArrayPool<T> bufferPool = ArrayPool<T>.Create();

                    public T[] data { get; private set; }
                    public int length { get; private set; }
                    public ulong timestamp { get; private set; }

                    public CopyBuffer(T[] array, ulong time)
                    {
                        data = bufferPool.Rent(array.Length);
                        Array.Copy(array, data, array.Length);
                        length = array.Length;
                        timestamp = time;
                    }

                    ~CopyBuffer()
                    {
                        Dispose();
                    }

                    public void Dispose()
                    {
                        if (data != null)
                        {
                            bufferPool.Return(data);
                            data = null;
                        }
                    }
                }

                public event Action<ulong, float[], int, bool> OnBufferReceived;

                private Task process = null;
                private SafeBool active = new();

                
                private ConcurrentQueue<CopyBuffer<float>> audioBuffers = new();

                private void Awake()
                {
                    process = new Task(ProcessLoop, TaskCreationOptions.LongRunning);
                    process.Start();
                }

                private void OnDestroy()
                {
                    active.Set(false);
                }

                private void OnAudioFilterRead(float[] data, int channels)
                {
                    ulong utcNow = (ulong)DateTime.UtcNow.Ticks * 100;
                    audioBuffers.Enqueue(new CopyBuffer<float>(data, utcNow));
                }

                private void ProcessLoop()
                {
                    active.Set(true);
                    while (active)
                    {
                        if (audioBuffers.TryDequeue(out var buffer))
                        {
                            bool silence = true;
                            for (int i = 0; silence && i < buffer.length; ++i)
                                silence &= Mathf.Abs(buffer.data[i]) < 1e-4;
                            
                            OnBufferReceived?.Invoke(buffer.timestamp, buffer.data, buffer.length, silence);
                            buffer.Dispose();
                        }
                        else Thread.Sleep(1);
                    }
                }
            }

            private EditorCapture virtualMic = null;
            private event Action<ulong, float[], int, bool> OnBufferReceived = null;
            private bool active = false;
            private bool destroyed = false;

            public int SampleRate => AudioSettings.outputSampleRate;
            public int Channels => 2;
            public bool Active => active;

            private void Awake()
            {
                // If we are a root gameobject
                if (transform.parent == null)
                {
                    // Make sure we aren't destroyed between scenes
                    DontDestroyOnLoad(this);
                }
                else Debug.LogWarning("MLSpectatorAudio/EditorManager: Awake() - Must be root object to persist between scenes");
            }

            private void OnDestroy()
            {
                destroyed = true;
                Destroy(virtualMic);
            }

            private void OnEnable()
            {
                // Subscribe to scene changes
                SceneManager.sceneLoaded += SceneLoaded;

                // Kick off Awake like a scene load
                SceneLoaded(SceneManager.GetActiveScene(), LoadSceneMode.Single);
            }

            private void OnDisable()
            {
                SceneManager.sceneLoaded -= SceneLoaded;
            }

            private void SceneLoaded(Scene scene, LoadSceneMode mode)
            {
                // Get a reference to the main camera object in our new scene
                var mainCam = Camera.main?.gameObject;
                if (mainCam == null)
                {
                    Debug.LogWarning("MLSpectatorAudio_EditorManager: SceneLoaded() - Main camera not found");
                    return;
                }

                // Get our audio capture component or add one if necessary
                if (!(virtualMic = mainCam.GetComponent<EditorCapture>()))
                    virtualMic = mainCam.AddComponent<EditorCapture>();
            }

            public void AddBufferAction(Action<ulong, float[], int, bool> action)
            {
                OnBufferReceived += action;
            }

            public void RemoveBufferAction(Action<ulong, float[], int, bool> action)
            {
                OnBufferReceived -= action;
            }

            public void StartCapture()
            {
                if (active)
                {
                    Debug.LogWarning("MLSpectatorAudio/EditorManager: StartCapture() - Capture is already running");
                    return;
                }

                virtualMic.OnBufferReceived += OnBufferReceived;
                active = true;

                Debug.Log("MLSpectatorAudio/EditorManager: StartCapture()");
            }

            public void StopCapture()
            {
                if (!active)
                {
                    Debug.LogWarning("MLSpectatorAudio/EditorManager: StopCapture() - Capture is not running");
                    return;
                }

                virtualMic.OnBufferReceived -= OnBufferReceived;
                active = false;

                Debug.Log("MLSpectatorAudio/EditorManager: StopCapture()");
            }

            public void Dispose()
            {
                if (destroyed) return;

                if (gameObject != null)
                    Destroy(gameObject);
            }
        }

        #endregion

        #region Private variables

        IAudioCapture virtualMic = null;

        AACEncoder encoder
        {
            get { lock (audioLock) { return _encoder; } }
            set { lock (audioLock) { _encoder = value; } }
        }
        AACEncoder _encoder;
        object audioLock = new();

        private bool connected = false;

        private long maxPacketDiff = 1024;
        private long totalPackets = 0;
        private ulong startTime = 0;

        private int SampleRate = 48000;
        private int Channels = 2;

        #endregion

        #region Public properties

        public bool Paused;

        #endregion

        #region MonoBehaviour methods

        private void Awake()
        {
#if !UNITY_EDITOR && UNITY_ANDROID
            virtualMic = new DeviceCapture();
            Debug.Log("MLSpectatorAudio: Awake() - Created virtual mic");
#else
            var gameObject = new GameObject("MLSpectatorAudio");
            virtualMic = gameObject.AddComponent<EditorManager>();
#endif
            virtualMic.AddBufferAction(OnBufferReceived);
            maxPacketDiff = virtualMic.SampleRate / 30;

            SampleRate = virtualMic.SampleRate;
            Channels = virtualMic.Channels;
        }

        private void OnDestroy()
        {
            if (virtualMic == null) return;
            if (virtualMic.Active) virtualMic.StopCapture();
            virtualMic.RemoveBufferAction(OnBufferReceived);
            virtualMic.Dispose();
            virtualMic = null;
        }

        private void OnEnable()
        {
            //Debug.Log("MLSpectatorAudio: OnEnable()");

            // Subscribe to relevant messages
            Host.Instance.RegisterCallback(AudioMode, AudioStateChange);

            // Subscribe to host connection
            Host.Instance.OnConnected += OnConnected;

            // Check if we are already connected
            if (Host.Instance.Connected) OnConnected();
        }

        private void OnDisable()
        {
            //Debug.Log("MLSpectatorAudio: OnDisable()");

            // Unsubscribe from host connection
            Host.Instance.OnConnected -= OnConnected;

            // Disconnect if we are connected
            if (connected) OnDisconnect();

            // Unsubscribe from relevant messages
            Host.Instance.UnregisterCallback(AudioMode, AudioStateChange);
        }

        #endregion

        #region Host callbacks

        private void OnConnected()
        {
            //Debug.Log("MLSpectatorAudio: OnConnected()");

            // Subscribe to host disconnect
            Host.Instance.OnDisconnect += OnDisconnect;

            // Mark that we are connected
            connected = true;
        }

        private void OnDisconnect()
        {
            //Debug.Log("MLSpectatorAudio: OnDisconnect()");

            // Unsubscribe from host disconnect
            Host.Instance.OnDisconnect -= OnDisconnect;

            // Reset our state
            Reset();

            // Mark that we are no longer connected
            connected = false;
        }

        private void AudioStateChange(Message message)
        {
            bool active = BitConverter.ToBoolean(message.data);

            if (virtualMic == null)
            {
                Debug.LogError("MLSpectatorAudio: AudioStateChange() - Missing virtual mic");
                return;
            }    

            if (active)
            {
                if (virtualMic.Active)
                {
                    Debug.Log("MLSpectatorAudio: AudioStateChange() - Already enabled");
                    return;
                }

                Debug.Log("MLSpectatorAudio: AudioStateChange() - Enabled");

                if (MLSpectator.Instance.EncodeAudio)
                {
                    if (encoder != null)
                    {
                        Debug.LogWarning("MLSpectatorAudio: AudioStateChange() - Encoder already exists");
                    }
                    else
                    {
                        encoder = AACEncoder.Create(SampleRate, Channels, 128000);
                        encoder.onOutput += SendEncodedAudio;
                    }
                }
                
                virtualMic.StartCapture();
            }

            else
            {
                if (!virtualMic.Active)
                {
                    Debug.LogWarning("MLSpectatorAudio: AudioStateChange() - Already disabled");
                    return;
                }

                Debug.Log("MLSpectatorAudio: AudioStateChange() - Disabled");
                virtualMic.StopCapture();
                encoder?.Dispose();
            }
        }

        #endregion

        #region Private methods

        private void OnBufferReceived(ulong time, float[] samples, int numSamples, bool isSilence)
        {
            if (!Host.Instance.Connected) return;

            // Calculate time from total packet count
            ulong sampleTime = (ulong)Math.Round((1e9 * totalPackets) / SampleRate);

            if (!Paused)
            {
                // If we don't have an encoder, send raw audio
                if (encoder == null) SendRawAudio(samples, numSamples, sampleTime);
                // Otherwise, push audio to our encoder
                else encoder.Push(new(samples, 0, numSamples), sampleTime);
            }

            // Make corrections to fix audio timing, as necessary
            CorrectAudioTiming(time, numSamples, isSilence);
        }

        private void CorrectAudioTiming(ulong time, int numSamples, bool isSilence)
        {
            // Make sure we have a start time
            if (startTime == 0) startTime = time;

            // Calculate our number of packets
            int packets = numSamples / Channels;

            // Calculate our expected total packet length
            float duration = (time - startTime) / 1e9f;
            long expectedPackets = Mathf.RoundToInt(duration * SampleRate);

            bool behind = (expectedPackets - totalPackets) > maxPacketDiff;

            // If we're silent or too far behind, correct our packet count
            if (isSilence || behind)
            {
                if (behind)
                {
                    Debug.Log($"SpectatorAudioCapture: CorrectAudioTiming() - Audio fell behind by {expectedPackets - totalPackets} samples");
                }

                totalPackets = expectedPackets;
            }

            // Otherwise, add our current packet count
            else totalPackets += packets;
        }

        private void Reset()
        {
            startTime = 0;
            totalPackets = 0;

            if (virtualMic.Active)
                virtualMic.StopCapture();
        }

        private void SendEncodedAudio(IntPtr data, int length, ulong timestamp)
        {
            // Make a new audio message with required size
            Message audioMessage = new Message(AudioSamples, sizeof(ulong) + sizeof(int) + length);

            // Copy time
            BitConverter.GetBytes(timestamp).CopyTo(audioMessage.data, 0);
            // Copy channels
            BitConverter.GetBytes(Channels).CopyTo(audioMessage.data, sizeof(ulong));
            // Copy encoded audio data
            Marshal.Copy(data, audioMessage.data, sizeof(ulong) + sizeof(int), length);

            // Send our audio message
            Host.Instance.Send(audioMessage);
            audioMessage.Dispose();
        }

        private void SendRawAudio(float[] samples, int length, ulong timestamp)
        {
            // Make a new audio message with required size
            Message audioMessage = new Message(AudioSamples, sizeof(ulong) + sizeof(int) + length * sizeof(float));

            // Copy time
            BitConverter.GetBytes(timestamp).CopyTo(audioMessage.data, 0);
            // Copy channels
            BitConverter.GetBytes(Channels).CopyTo(audioMessage.data, sizeof(ulong));
            // Copy raw audio data
            Buffer.BlockCopy(samples, 0, audioMessage.data, sizeof(ulong) + sizeof(int), length * sizeof(float));

            // Send our audio message
            Host.Instance.Send(audioMessage);
            audioMessage.Dispose();
        }

        #endregion
    }
}