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

namespace MagicLeap.Spectator
{
    #region C++ API

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

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

#if UNITY_EDITOR_OSX || UNITY_STANDALONE_OSX

        private const string H264EncoderLib = "macVideoEncoder";

        [DllImport(H264EncoderLib, CallingConvention = CallingConvention.Cdecl)]
        protected static extern int CreateVideoEncoder(int width, int height, int bitrate, Callback callback);

        [DllImport(H264EncoderLib, CallingConvention = CallingConvention.Cdecl)]
        protected static extern Result DestroyVideoEncoder(int id);

        [DllImport(H264EncoderLib, CallingConvention = CallingConvention.Cdecl)]
        protected static extern Result PushImageToVideoEncoder(int id, IntPtr imagePtr);
        
        [DllImport(H264EncoderLib, CallingConvention = CallingConvention.Cdecl)]
        protected static extern Result PushDataToVideoEncoder(int id, int width, int height, IntPtr data);

#else

        protected static int CreateVideoEncoder(int width, int height, int bitrate, Callback callback) => 0;

        protected static Result DestroyVideoEncoder(int id) => Result.NOT_INITIALIZED;

        protected static Result PushImageToVideoEncoder(int id, IntPtr imagePtr) => Result.NOT_INITIALIZED;
        
        protected static Result PushDataToVideoEncoder(int id, int width, int height, IntPtr data) => Result.NOT_INITIALIZED;

#endif
    }

    #endregion

    #region C# API

    public class MacH264Encoder : MacH264EncoderAPI, IVideoEncoder
    {
        #region Static

        private static ConcurrentDictionary<int, MacH264Encoder> encoderMap = new();

        [MonoPInvokeCallback(typeof(Callback))]
        private static void StaticCallback(int id, IntPtr dataPtr, int length)
        {
            if (encoderMap.TryGetValue(id, out MacH264Encoder encoder))
                encoder.OnFrameReceived(dataPtr, length);
        }

        #endregion

        #region Private variables

        private int encoderId = 0;
        private byte[] buffer = null;
        private RenderTexture input = null;
        private IntPtr inputPtr = IntPtr.Zero;
        private bool hasResources = false;

        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 MacH264Encoder(int width, int height, int bitrate)
        {
            Width = width;
            Height = height;

            encoderId = CreateVideoEncoder(width, height, bitrate, StaticCallback);

            if (encoderId == 0)
                Debug.LogError("MacH264Encoder: H264Encoder() - Error encountered creating encoder");
            else Debug.Log($"MacH264Encoder: H264Encoder() - {width}x{height} encoder created");
        }

        ~MacH264Encoder()
        {
            if (encoderId == 0) return;
            Debug.LogWarning("MacH264Encoder: ~MacH264Encoder() - Encoder was not disposed of properly");
            Dispose();
        }

        #endregion

        #region Public methods

        public static MacH264Encoder Create(int width, int height, int bitrate)
        {
            MacH264Encoder encoder = new MacH264Encoder(width, height, bitrate);
            if (encoder.encoderId == 0) return null;

            encoder.CreateResources();
            if (!encoderMap.TryAdd(encoder.encoderId, encoder))
                Debug.Log("MacH264Encoder: Create() - Failed to add encoder to map");

            return encoder;
        }

        public void Dispose()
        {
            if (encoderId == 0) return;

            if (!encoderMap.TryRemove(encoderId, out var _))
                Debug.LogError("MacH264Encoder: Dispose() - Failed to remove encoder from map");

            if (DestroyVideoEncoder(encoderId) == Result.SUCCESS) encoderId = 0;
            else Debug.LogError("MacH264Encoder: Dispose() - Failed to destroy encoder");
           
            ReleaseResources();

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

        public void AttachToCamera(Camera camera, Material material)
        {
            if (camera.targetTexture == null)
            {
                Debug.LogError("MacH264Encoder: 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("MacH264Encoder: Push() - Encoder is attached to a camera, cannot push");
                return;
            }
            
            if (!material) Graphics.Blit(frame, input);
            else Graphics.Blit(frame, input, material);

            if (PushImageToVideoEncoder(encoderId, inputPtr) != Result.SUCCESS)
                Debug.LogError("MacH264Encoder: Push() - Failed to push frame to encoder");
        }

        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);

            // Assign our input texture to our encoder
            inputPtr = input.GetNativeTexturePtr();

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

        private void ReleaseResources()
        {
            lock (pushLock)
            {
                if (input != null)
                {
                    inputPtr = IntPtr.Zero;

                    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 void OnFrameReceived(IntPtr data, int size)
        {
            // Invalid frames will have a size of 0
            if (size == 0) return;

            if (buffer == null || buffer.Length < size)
                buffer = new byte[size];

            Marshal.Copy(data, buffer, 0, size);
            onFrameEncoded?.Invoke(buffer, size);
        }

        private unsafe void AsyncPush(AsyncGPUReadbackRequest request, ulong id)
        {
            if (!request.hasError)
            {
                var data = request.GetData<byte>();
                var dataPtr = (IntPtr)NativeArrayUnsafeUtility.GetUnsafePtr(data);

                lock (pushLock)
                {
                    if (id < lastFrameId) Debug.LogError("MacH264Encoder: TryPush() - Pushing data out of order");
                    if (input && PushDataToVideoEncoder(encoderId, input.width, input.height, dataPtr) != Result.SUCCESS)
                        Debug.LogError("MacH264Encoder: AsyncPush() - Failed to push data to encoder");
                    lastFrameId = id;
                }

                data.Dispose();
            }
        }

        #endregion
    }

    #endregion
}