// %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 UnityEngine;
using UnityEngine.Rendering;

namespace MagicLeap.Spectator
{
    #region C++ API

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

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

#if !UNITY_EDITOR && UNITY_ANDROID

        private const string H264EncoderLib = "h264_encoder";

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

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

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

        [DllImport(H264EncoderLib, CallingConvention = CallingConvention.Cdecl)]
        protected static extern Result H264Encoder_Start(IntPtr encoder, Callback callback);

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

        [DllImport(H264EncoderLib, CallingConvention = CallingConvention.Cdecl)]
        protected static extern IntPtr H264Encoder_Update();

#else

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

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

        protected static Result H264Encoder_AssignInput(IntPtr encoder, IntPtr vkImgPtr) => Result.NOT_INITIALIZED;

        protected static Result H264Encoder_Start(IntPtr encoder, Callback callback) => Result.NOT_INITIALIZED;

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

        protected static IntPtr H264Encoder_Update() => IntPtr.Zero;

#endif
    }

    #endregion

    #region C# API

    public class MLH264Encoder : MLH264EncoderAPI, IVideoEncoder
    {
        #region Static

        private static ConcurrentDictionary<IntPtr, MLH264Encoder> encoderMap = new();

        [MonoPInvokeCallback(typeof(Callback))]
        private static void GlobalCallback(IntPtr ptr, IntPtr data, int size)
        {
            if (encoderMap.TryGetValue(ptr, out var encoder))
                encoder.OnFrameEncoded(data, size);
        }

        #endregion

        #region Private variables

        private IntPtr ptr = IntPtr.Zero;
        private byte[] buffer = null;
        private RenderTexture input = null;
        private bool hasResources = false;

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

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

            ptr = H264Encoder_Create(width, height, width, bitrate, 30, -1);

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

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

        #endregion

        #region Public methods

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

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

            H264Encoder_Start(encoder.ptr, GlobalCallback);

            return encoder;
        }

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

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

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

            ReleaseResources();

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

        public void AttachToCamera(Camera camera, Material material)
        {
            if (camera.targetTexture == null)
            {
                Debug.LogError("MLH264Encoder: 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.IssuePluginEventAndData(H264Encoder_Update(), 0, ptr);

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

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

            var cmd = CommandBufferPool.Get();

            if (!material) cmd.Blit(frame, input);
            else cmd.Blit(frame, input, material);

            cmd.IssuePluginEventAndData(H264Encoder_Update(), 0, ptr);
            Graphics.ExecuteCommandBuffer(cmd);

            cmd.Clear();
            CommandBufferPool.Release(cmd);
        }

        public bool TryRequestSyncFrame()
        {
            return H264Encoder_RequestSyncFrame(ptr) == Result.SUCCESS;
        }

        #endregion

        #region Private methods

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

            int start = Environment.TickCount;

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

            // Assign our input texture to our encoder
            IntPtr inputPtr = input.GetNativeTexturePtr();
            if (H264Encoder_AssignInput(ptr, inputPtr) != Result.SUCCESS)
                Debug.LogError("MLH264Encoder: CreateResources() - Error encountered setting input source");
            else Debug.Log("MLH264Encoder: CreateResources() - Set input source successfully");

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

        private void ReleaseResources()
        {
            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 void OnFrameEncoded(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);
        }

        #endregion
    }

    #endregion
}