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)qualityMode);
            renderWidth = (uint)rw;
            renderHeight = (uint)rh;
            return true;
        }

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

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

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

            Fsr3Upscaler.InitializationFlags flags = 0;
            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;

            _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)
        {
            var dispatchDescription = new Fsr3Upscaler.DispatchDescription
            {
                Color = new ResourceView(textures.colorInput),
                Depth = new ResourceView(textures.depth),
                MotionVectors = new ResourceView(textures.motionVectors),
                Exposure = new ResourceView(textures.exposureTexture),
                Reactive = new ResourceView(textures.biasColorMask),
                TransparencyAndComposition = new ResourceView(textures.transparencyMask),
                Output = new ResourceView(textures.colorOutput),
                
                JitterOffset = new Vector2(_executeData.jitterOffsetX, _executeData.jitterOffsetY),
                MotionVectorScale = new Vector2(_executeData.MVScaleX, _executeData.MVScaleY),
                RenderSize = new Vector2Int((int)_executeData.renderSizeWidth, (int)_executeData.renderSizeHeight),
                InputResourceSize = new Vector2Int((int)_executeData.renderSizeWidth, (int)_executeData.renderSizeHeight),
                EnableSharpening = _executeData.enableSharpening != 0,
                Sharpness = _executeData.sharpness,
                FrameTimeDelta = _executeData.frameTimeDelta / 1000f,
                PreExposure = _executeData.preExposure,
                Reset = _executeData.reset != 0,
                CameraNear = _executeData.cameraNear,
                CameraFar = _executeData.cameraFar,
                CameraFovAngleVertical = _executeData.cameraFovAngleVertical,
                ViewSpaceToMetersFactor = 1.0f, // 1 unit is 1 meter in Unity
            };
            
            _context.Dispatch(dispatchDescription, cmd);
        }
    }
}