using UnityEngine;
using UnityEngine.Experimental.Rendering;

namespace FidelityFX.OpticalFlow
{
    internal class OpticalFlowResources
    {
        public readonly RenderTexture[][] OpticalFlowInputLevels = 
        {
            new RenderTexture[2],
            new RenderTexture[2],
            new RenderTexture[2],
            new RenderTexture[2],
            new RenderTexture[2],
            new RenderTexture[2],
            new RenderTexture[2],
        };
        
        public readonly RenderTexture[][] OpticalFlowLevels = 
        {
            new RenderTexture[2],
            new RenderTexture[2],
            new RenderTexture[2],
            new RenderTexture[2],
            new RenderTexture[2],
            new RenderTexture[2],
            new RenderTexture[2],
        };
        
        public RenderTexture OpticalFlowSCDHistogram;
        public RenderTexture OpticalFlowSCDPreviousHistogram;
        public RenderTexture OpticalFlowSCDTemp;
        
        public void Create(OpticalFlow.ContextDescription contextDescription)
        {
            Vector2Int opticalFlowInputTextureSize = contextDescription.resolution;
            
            const int minBlockSize = 8;
            Vector2Int opticalFlowTextureSize = OpticalFlow.GetOpticalFlowTextureSize(contextDescription.resolution, minBlockSize);
            
            // TODO: this seems useless, delete unless something comes up
            Vector2Int opticalFlowHistogramTextureSize = OpticalFlow.GetOpticalFlowHistogramSize(0);
            Vector2Int globalMotionSearchMaxDispatchSize = OpticalFlow.GetGlobalMotionSearchDispatchSize(0);
            int globalMotionSearchTextureWidth = 4 + (globalMotionSearchMaxDispatchSize.x * globalMotionSearchMaxDispatchSize.y);

            for (int level = 0; level < OpticalFlow.OpticalFlowMaxPyramidLevels; ++level)
            {
                CreateDoubleBufferedResource(OpticalFlowInputLevels[level], $"OPTICALFLOW_OpticalFlowInputLevel{level}-", opticalFlowInputTextureSize, GraphicsFormat.R8_UInt);
                CreateDoubleBufferedResource(OpticalFlowLevels[level], $"OPTICALFLOW_OpticalFlowLevel{level}-", opticalFlowTextureSize, GraphicsFormat.R16G16_SInt);

                opticalFlowTextureSize = NextLevelSize(opticalFlowTextureSize);
                opticalFlowInputTextureSize = new Vector2Int(opticalFlowInputTextureSize.x >> 1, opticalFlowInputTextureSize.y >> 1);
            }

            OpticalFlowSCDHistogram = CreateResource("OPTICALFLOW_OpticalFlowSCDHistogram", new Vector2Int(OpticalFlow.GetSCDHistogramTextureWidth(), 1), GraphicsFormat.R32_UInt);
            OpticalFlowSCDPreviousHistogram = CreateResource("OPTICALFLOW_OpticalFlowSCDPreviousHistogram", new Vector2Int(OpticalFlow.GetSCDHistogramTextureWidth(), 1), GraphicsFormat.R32_SFloat);
            OpticalFlowSCDTemp = CreateResource("OPTICALFLOW_OpticalFlowSCDTemp", new Vector2Int(3, 1), GraphicsFormat.R32_UInt);
        }

        public void Destroy()
        {
            DestroyResource(ref OpticalFlowSCDTemp);
            DestroyResource(ref OpticalFlowSCDPreviousHistogram);
            DestroyResource(ref OpticalFlowSCDHistogram);

            for (int level = 0; level < OpticalFlow.OpticalFlowMaxPyramidLevels; ++level)
            {
                DestroyResource(OpticalFlowLevels[level]);
                DestroyResource(OpticalFlowInputLevels[level]);
            }
        }

        private static Vector2Int NextLevelSize(Vector2Int size)
        {
            return new Vector2Int(AlignUp(size.x, 2) / 2, AlignUp(size.y, 2) / 2);
        }

        private static int AlignUp(int x, int y)
        {
            return (x + (y - 1)) & ~(y - 1);
        }

        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 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(RenderTexture[] resource)
        {
            for (int i = 0; i < resource.Length; ++i)
                DestroyResource(ref resource[i]);
        }
    }
}