// %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.Runtime.InteropServices;
using Unity.Collections.LowLevel.Unsafe;
using UnityEngine;
using System.Threading;
using UnityEngine.Rendering;

namespace MagicLeap.Spectator
{
    #region C++ API

    public class WinH264EncoderAPI
    {
        protected enum Result
        {
            NOT_INITIALIZED = -1,
            FAIL = 0,
            SUCCESS = 1,
        }

        protected delegate void Callback(IntPtr encoder, IntPtr data, int length);

#if UNITY_EDITOR_WIN || UNITY_STANDALONE_WIN

        private const string H264EncoderLib = "h264_encoder";

        [DllImport(H264EncoderLib, CallingConvention = CallingConvention.Cdecl)]
        protected static extern IntPtr H264Encoder_Create(int width, int height, int bitrate);

        [DllImport(H264EncoderLib, CallingConvention = CallingConvention.Cdecl)]
        protected static extern Result H264Encoder_Destroy(ref IntPtr encoder);

        [DllImport(H264EncoderLib, CallingConvention = CallingConvention.Cdecl)]
        protected static extern Result H264Encoder_Push(IntPtr encoder, IntPtr data, int length);

        [DllImport(H264EncoderLib, CallingConvention = CallingConvention.Cdecl)]
        protected static extern Result H264Encoder_Peek(IntPtr encoder, ref int size);

        [DllImport(H264EncoderLib, CallingConvention = CallingConvention.Cdecl)]
        protected static extern Result H264Encoder_Pull(IntPtr encoder, IntPtr data, int length);

#else

        protected static IntPtr H264Encoder_Create(int width, int height, int bitrate) => IntPtr.Zero;

        protected static Result H264Encoder_Destroy(ref IntPtr encoder) => Result.NOT_INITIALIZED;

        protected static Result H264Encoder_Push(IntPtr encoder, IntPtr data, int length) => Result.NOT_INITIALIZED;

        protected static Result H264Encoder_Peek(IntPtr encoder, ref int size) => Result.NOT_INITIALIZED;

        protected static Result H264Encoder_Pull(IntPtr encoder, IntPtr data, int length) => Result.NOT_INITIALIZED;

#endif
    }

    #endregion

    #region C# API

    public class WinH264Encoder : WinH264EncoderAPI, IVideoEncoder
    {
        #region Private variables

        IntPtr ptr = IntPtr.Zero;
        private RenderTexture input = null;
        private bool hasResources = false;

        private bool running
        {
            get { lock (runLock) { return _running; } }
            set { lock (runLock) { _running = value; } }
        }
        private bool _running = false;
        private object runLock = new();

        private Thread loop = null;

        private Camera camera = null;
        private CommandBuffer cmd = null;

        private object pushLock = new();
        private ulong frameId = 0, lastFrameId = 0;

        #endregion

        #region Public properties

        public int Width { get; private set; }
        public int Height { get; private set; }

        #endregion

        #region Public events

        public event Action<byte[], int> onFrameEncoded = null;

        #endregion

        #region Constructor / Destructor

        private WinH264Encoder(int width, int height, int bitrate)
        {
            Width = width;
            Height = height;

            ptr = H264Encoder_Create(width, height, bitrate);

            if (ptr == IntPtr.Zero)
                Debug.LogError("WinH264Encoder: WinH264Encoder() - Error encountered creating encoder!");
            else Debug.Log($"WinH264Encoder: WinH264Encoder() - {width}x{height} encoder created");
        }

        ~WinH264Encoder()
        {
            if (ptr == IntPtr.Zero) return;
            Debug.LogWarning("WinH264Encoder: ~WinH264Encoder() - Encoder was not disposed of properly");
            Dispose();
        }

        #endregion

        #region Public methods

        public static WinH264Encoder Create(int width, int height, int bitrate)
        {
            WinH264Encoder encoder = new WinH264Encoder(width, height, bitrate);
            if (encoder.ptr == IntPtr.Zero) return null;

            encoder.CreateResources();
            encoder.loop.Start();

            return encoder;
        }

        public void Dispose()
        {
            if (ptr == IntPtr.Zero) return;

            // Must do this first to join our loop before destroying the encoder
            ReleaseResources();

            if (H264Encoder_Destroy(ref ptr) != Result.SUCCESS)
                Debug.LogError("WinH264Encoder: Dispose() - Failed to destroy encoder");

            Debug.Log("WinH264Encoder: Dispose() - Success");
        }

        public void AttachToCamera(Camera camera, Material material)
        {
            if (camera.targetTexture == null)
            {
                Debug.LogError("WinH264Encoder: AttachToCamera() - Camera must have target texture");
                return;
            }

            this.camera = camera;
            cmd = new CommandBuffer();

            if (!material) cmd.Blit(camera.targetTexture, input);
            else cmd.Blit(camera.targetTexture, input, material);
            cmd.RequestAsyncReadback(input, 0,  r => AsyncPush(r, ++frameId));

            camera.AddCommandBuffer(CameraEvent.AfterEverything, cmd);
        }

        public void Push(Texture frame, Material material)
        {
            if (camera != null)
            {
                Debug.LogWarning("WinH264Encoder: Push() - Encoder is attached to a camera, cannot push");
                return;
            }

            if (!material) Graphics.Blit(frame, input);
            else Graphics.Blit(frame, input, material);
            GL.Flush();

            AsyncGPUReadback.Request(input, 0, r => AsyncPush(r, ++frameId));
        }

        public bool TryRequestSyncFrame() => false;

        #endregion

        #region Private methods

        private void CreateResources()
        {
            if (hasResources) return;

            int start = Environment.TickCount;

            // Create input texture
            input = new RenderTexture(Width, Height, 0, RenderTextureFormat.BGRA32, 0);
            input.Create();
            Graphics.Blit(Texture2D.blackTexture, input);

            loop = new Thread(OutputLoop);

            Debug.Log($"WinH264Encoder: CreateResources() - Took {Environment.TickCount - start} ms");
            hasResources = true;
        }

        private void ReleaseResources()
        {
            running = false;
            loop.Join();

            if (input != null)
            {
                input.Release();
                UnityEngine.Object.Destroy(input);
                input = null;
            }

            if (cmd != null)
            {
                if (camera != null)
                {
                    camera.RemoveCommandBuffer(CameraEvent.AfterEverything, cmd);
                    camera = null;
                }
                cmd.Release();
                cmd = null;
            }

            hasResources = false;
        }

        private bool TryPull(ref byte[] data, ref int length)
        {
            if (ptr == IntPtr.Zero) return false;

            if (H264Encoder_Peek(ptr, ref length) != Result.SUCCESS)
                return false;

            if (data == null || data.Length < length)
                data = new byte[length];

            GCHandle gcHandle = default;
            try
            {
                gcHandle = GCHandle.Alloc(data, GCHandleType.Pinned);
                IntPtr dataPtr = gcHandle.AddrOfPinnedObject();
                return H264Encoder_Pull(ptr, dataPtr, length) == Result.SUCCESS;
            }
            catch (Exception ex)
            {
                Debug.LogError($"WinH264Encoder: TryPull() - {ex.Message}");
            }
            finally
            {
                if (gcHandle != default)
                    gcHandle.Free();
            }
            return false;
        }

        private void OutputLoop()
        {
            int length = 0;
            byte[] buffer = null;

            running = true;
            while (running)
            {
                if (TryPull(ref buffer, ref length))
                    onFrameEncoded?.Invoke(buffer, length);
                else Thread.Sleep(1);
            }
        }

        private unsafe void AsyncPush(AsyncGPUReadbackRequest request, ulong id)
        {
            if (!running) return;

            if (!request.hasError)
            {
                var data = request.GetData<byte>();
                var dataPtr = (IntPtr)NativeArrayUnsafeUtility.GetUnsafePtr(data);

                lock (pushLock)
                {
                    if (id < lastFrameId) Debug.LogError("WinH264Encoder: TryPush() - Pushing data out of order");
                    if (H264Encoder_Push(ptr, dataPtr, data.Length) != Result.SUCCESS)
                        Debug.LogError("WinH264Encoder: AsyncPush() - Failed to push data to encoder");
                    lastFrameId = id;
                }

                data.Dispose();
            }
        }

        #endregion
    }

    #endregion
}