using System;
using System.Runtime.InteropServices;
using UnityEngine;
using UnityEngine.Experimental.Rendering;

namespace FidelityFX.FrameGen
{
    internal class FrameInterpolationResources
    {
        public RenderTexture ReconstructedDepthInterpolatedFrame;
        public RenderTexture GameMotionVectorFieldX;
        public RenderTexture GameMotionVectorFieldY;
        public RenderTexture InpaintingPyramid;
        public ComputeBuffer Counters;
        public RenderTexture OpticalFlowMotionVectorFieldX;
        public RenderTexture OpticalFlowMotionVectorFieldY;
        public RenderTexture PreviousInterpolationSource;
        public RenderTexture DisocclusionMask;
        
        public readonly RenderTexture[] DilatedDepth = new RenderTexture[2];
        public readonly RenderTexture[] DilatedMotionVectors = new RenderTexture[2];
        public readonly RenderTexture[] ReconstructedDepth = new RenderTexture[2];
        
        public void Create(in FrameInterpolation.ContextDescription contextDescription)
        {
            Vector2Int displaySize = contextDescription.displaySize;
            Vector2Int displaySizeDiv2 = new Vector2Int(displaySize.x / 2, displaySize.y / 2);
            Vector2Int maxRenderSize = contextDescription.maxRenderSize;

            ReconstructedDepthInterpolatedFrame = CreateResource("FI_ReconstructedDepthInterpolatedFrame", maxRenderSize, GraphicsFormat.R32_UInt);
            GameMotionVectorFieldX = CreateResource("FI_GameMotionVectorFieldX", maxRenderSize, GraphicsFormat.R32_UInt);
            GameMotionVectorFieldY = CreateResource("FI_GameMotionVectorFieldY", maxRenderSize, GraphicsFormat.R32_UInt);
            InpaintingPyramid = CreateResourceMips("FI_InpaintingPyramid", displaySizeDiv2, GraphicsFormat.R16G16B16A16_SFloat);
            Counters = CreateBuffer<uint>("FI_Counters", 2);    // structured buffer containing 2 UINT values
            OpticalFlowMotionVectorFieldX = CreateResource("FI_OpticalFlowMotionVectorFieldX", maxRenderSize, GraphicsFormat.R32_UInt);
            OpticalFlowMotionVectorFieldY = CreateResource("FI_OpticalFlowMotionVectorFieldY", maxRenderSize, GraphicsFormat.R32_UInt);
            PreviousInterpolationSource = CreateResource("FI_PreviousInterpolationSource", displaySize, contextDescription.backBufferFormat);
            DisocclusionMask = CreateResource("FI_DisocclusionMask", maxRenderSize, GraphicsFormat.R8G8_UNorm);
            
            CreateDoubleBufferedResource(DilatedDepth, "FI_DilatedDepth_", maxRenderSize, GraphicsFormat.R32_SFloat);
            CreateDoubleBufferedResource(DilatedMotionVectors, "FI_DilatedMVs_", maxRenderSize, GraphicsFormat.R16G16_SFloat);
            CreateDoubleBufferedResource(ReconstructedDepth, "FI_ReconstructedDepth_", maxRenderSize, GraphicsFormat.R32_UInt);
        }
        
        public void Destroy()
        {
            DestroyResource(ReconstructedDepth);
            DestroyResource(DilatedMotionVectors);
            DestroyResource(DilatedDepth);
            
            DestroyResource(ref DisocclusionMask);
            DestroyResource(ref PreviousInterpolationSource);
            DestroyResource(ref OpticalFlowMotionVectorFieldY);
            DestroyResource(ref OpticalFlowMotionVectorFieldX);
            DestroyResource(ref Counters);
            DestroyResource(ref InpaintingPyramid);
            DestroyResource(ref GameMotionVectorFieldY);
            DestroyResource(ref GameMotionVectorFieldX);
            DestroyResource(ref ReconstructedDepthInterpolatedFrame);
        }
        
        private static RenderTexture CreateResource(string name, Vector2Int size, GraphicsFormat format)
        {
            var rt = new RenderTexture(size.x, size.y, 0, format) { name = name, enableRandomWrite = true };
            rt.Create();
            return rt;
        }

        private static RenderTexture CreateResourceMips(string name, Vector2Int size, GraphicsFormat format)
        {
            int mipCount = 1 + Mathf.FloorToInt(Mathf.Log(Math.Max(size.x, size.y), 2.0f));
            var rt = new RenderTexture(size.x, size.y, 0, format, mipCount) { name = name, enableRandomWrite = true, useMipMap = true, autoGenerateMips = false };
            rt.Create();
            return rt;
        }
        
        private static ComputeBuffer CreateBuffer<TElem>(string name, int count)
        {
            return new ComputeBuffer(count, Marshal.SizeOf<TElem>());
        }
        
        private static void CreateDoubleBufferedResource(RenderTexture[] resource, string name, Vector2Int size, GraphicsFormat format)
        {
            for (int i = 0; i < 2; ++i)
            {
                resource[i] = new RenderTexture(size.x, size.y, 0, format) { name = name + (i + 1), enableRandomWrite = true };
                resource[i].Create();
            }
        }

        private static void DestroyResource(ref RenderTexture resource)
        {
            if (resource == null)
                return;
            
            resource.Release();
            resource = null;
        }

        private static void DestroyResource(ref ComputeBuffer resource)
        {
            if (resource == null)
                return;
            
            resource.Release();
            resource = null;
        }

        private static void DestroyResource(RenderTexture[] resource)
        {
            for (int i = 0; i < resource.Length; ++i)
                DestroyResource(ref resource[i]);
        }
    }
}