using System.Collections.Generic;
using FidelityFX;

namespace UnityEngine.Rendering.HighDefinition.AMD.FSR3
{
    /// <summary>
    /// Custom upscaler plugin that uses the open source port of FSR3 to Unity.
    /// This eschews using native plugins, instead using portable code that will work on any platform that supports compute shaders.
    /// </summary>
    public class FSR3UpscalerPlugin: UpscalerPlugin
    {
        private static FSR3GraphicsDevice sGraphicsDeviceInstance;

        public override string name => "FSR 3.0";

        public override bool isSupported => SystemInfo.supportsComputeShaders;

        public override bool Load() => true;

        public override bool IsLoaded() => true;

        public override GraphicsDevice CreateGraphicsDevice()
        {
            if (sGraphicsDeviceInstance != null)
            {
                sGraphicsDeviceInstance.Shutdown();
                sGraphicsDeviceInstance.Initialize();
                return sGraphicsDeviceInstance;
            }

            var graphicsDevice = new FSR3GraphicsDevice();
            if (graphicsDevice.Initialize())
            {
                sGraphicsDeviceInstance = graphicsDevice;
                return graphicsDevice;
            }

            Debug.LogWarning("Failed to initialize FSR3 Graphics Device");
            return null;
        }

        public override void DestroyGraphicsDevice()
        {
            if (sGraphicsDeviceInstance != null)
            {
                sGraphicsDeviceInstance.Shutdown();
                sGraphicsDeviceInstance = null;
            }
        }

        public override GraphicsDevice device => sGraphicsDeviceInstance;
    }

    public class FSR3GraphicsDevice : GraphicsDevice
    {
        private readonly Stack<FSR3Context> _contextPool = new();

        private Fsr3UpscalerAssets _assets;
        
        internal bool Initialize()
        {
            if (_assets != null)
                return true;
            
            _assets = Resources.Load<Fsr3UpscalerAssets>("Fsr3UpscalerAssets");
            return _assets != null;
        }

        internal void Shutdown()
        {
            foreach (var context in _contextPool)
            {
                context.Reset();
            }
            
            if (_assets != null)
            {
                Resources.UnloadAsset(_assets);
                _assets = null;
            }
        }
        
        public override FSR2Context CreateFeature(CommandBuffer cmd, in FSR2CommandInitializationData initSettings)
        {
            var context = _contextPool.Count != 0 ? _contextPool.Pop() : new FSR3Context();
            context.Init(initSettings, _assets);
            return context;
        }

        public override void DestroyFeature(CommandBuffer cmd, FSR2Context fsrContext)
        {
            var context = (FSR3Context)fsrContext;
            context.Reset();
            _contextPool.Push(context);
        }

        public override void ExecuteFSR2(CommandBuffer cmd, FSR2Context fsrContext, in FSR2TextureTable textures)
        {
            ((FSR3Context)fsrContext).Draw(cmd, in textures);
        }

        public override bool GetRenderResolutionFromQualityMode(FSR2Quality qualityMode, uint displayWidth, uint displayHeight, out uint renderWidth, out uint renderHeight)
        {
            Fsr3Upscaler.GetRenderResolutionFromQualityMode(out int rw, out int rh, (int)displayWidth, (int)displayHeight, (Fsr3Upscaler.QualityMode)((int)qualityMode + 2));
            renderWidth = (uint)rw;
            renderHeight = (uint)rh;
            return true;
        }

        public override float GetUpscaleRatioFromQualityMode(FSR2Quality qualityMode)
        {
            return Fsr3Upscaler.GetUpscaleRatioFromQualityMode((Fsr3Upscaler.QualityMode)((int)qualityMode + 2));
        }
    }

    public class FSR3Context : FSR2Context
    {
        private FSR2CommandInitializationData _initData;
        public override ref FSR2CommandInitializationData initData => ref _initData;

        private FSR2CommandExecutionData _executeData;
        public override ref FSR2CommandExecutionData executeData => ref _executeData;

        private readonly Fsr3UpscalerContext _context = new();
        private readonly Fsr3Upscaler.DispatchDescription _dispatchDescription = new();

        internal void Init(in FSR2CommandInitializationData initSettings, Fsr3UpscalerAssets assets)
        {
            _initData = initSettings;

            Fsr3Upscaler.InitializationFlags flags = Fsr3Upscaler.InitializationFlags.EnableFP16Usage;
            if (initSettings.GetFlag(FfxFsr2InitializationFlags.EnableHighDynamicRange)) flags |= Fsr3Upscaler.InitializationFlags.EnableHighDynamicRange;
            if (initSettings.GetFlag(FfxFsr2InitializationFlags.EnableDisplayResolutionMotionVectors)) flags |= Fsr3Upscaler.InitializationFlags.EnableDisplayResolutionMotionVectors;
            if (initSettings.GetFlag(FfxFsr2InitializationFlags.EnableMotionVectorsJitterCancellation)) flags |= Fsr3Upscaler.InitializationFlags.EnableMotionVectorsJitterCancellation;
            if (initSettings.GetFlag(FfxFsr2InitializationFlags.DepthInverted)) flags |= Fsr3Upscaler.InitializationFlags.EnableDepthInverted;
            if (initSettings.GetFlag(FfxFsr2InitializationFlags.EnableDepthInfinite)) flags |= Fsr3Upscaler.InitializationFlags.EnableDepthInfinite;
            if (initSettings.GetFlag(FfxFsr2InitializationFlags.EnableAutoExposure)) flags |= Fsr3Upscaler.InitializationFlags.EnableAutoExposure;
            if (initSettings.GetFlag(FfxFsr2InitializationFlags.EnableDynamicResolution)) flags |= Fsr3Upscaler.InitializationFlags.EnableDynamicResolution;

            //Debug.Log($"Initializing FSR3 with max render size: {initSettings.maxRenderSizeWidth}x{initSettings.maxRenderSizeHeight}, display size: {initSettings.displaySizeWidth}x{initSettings.displaySizeHeight}, flags: {flags}");
            
            _context.Create(new Fsr3Upscaler.ContextDescription
            {
                DisplaySize = new Vector2Int((int)initSettings.displaySizeWidth, (int)initSettings.displaySizeHeight),
                MaxRenderSize = new Vector2Int((int)initSettings.maxRenderSizeWidth, (int)initSettings.maxRenderSizeHeight),
                Flags = flags,
                Shaders = assets.shaders,
            });
        }

        internal void Reset()
        {
            _context.Destroy();
            
            _initData = new FSR2CommandInitializationData();
            _executeData = new FSR2CommandExecutionData();
        }

        internal void Draw(CommandBuffer cmd, in FSR2TextureTable textures)
        {
            _dispatchDescription.Color = new ResourceView(textures.colorInput);
            _dispatchDescription.Depth = new ResourceView(textures.depth);
            _dispatchDescription.MotionVectors = new ResourceView(textures.motionVectors);
            _dispatchDescription.Exposure = new ResourceView(textures.exposureTexture);
            _dispatchDescription.Reactive = new ResourceView(textures.biasColorMask);
            _dispatchDescription.TransparencyAndComposition = new ResourceView(textures.transparencyMask);
            _dispatchDescription.Output = new ResourceView(textures.colorOutput);

            _dispatchDescription.JitterOffset = new Vector2(_executeData.jitterOffsetX, _executeData.jitterOffsetY);
            _dispatchDescription.MotionVectorScale = new Vector2(_executeData.MVScaleX, _executeData.MVScaleY);
            _dispatchDescription.RenderSize = new Vector2Int((int)_executeData.renderSizeWidth, (int)_executeData.renderSizeHeight);
            _dispatchDescription.InputResourceSize = new Vector2Int((int)_executeData.renderSizeWidth, (int)_executeData.renderSizeHeight);
            _dispatchDescription.EnableSharpening = _executeData.enableSharpening != 0;
            _dispatchDescription.Sharpness = _executeData.sharpness;
            _dispatchDescription.FrameTimeDelta = _executeData.frameTimeDelta / 1000f;
            _dispatchDescription.PreExposure = _executeData.preExposure;
            _dispatchDescription.Reset = _executeData.reset != 0;
            _dispatchDescription.CameraNear = _executeData.cameraNear;
            _dispatchDescription.CameraFar = _executeData.cameraFar;
            _dispatchDescription.CameraFovAngleVertical = _executeData.cameraFovAngleVertical;
            _dispatchDescription.ViewSpaceToMetersFactor = 1.0f; // 1 unit is 1 meter in Unity
            
            _context.Dispatch(_dispatchDescription, cmd);
        }
    }
}