using System;
using System.Collections.Generic;

// This is placed in an AMD sub-namespace, to trick HDRP into using classes from here instead of from the UnityEngine.AMD namespace.
// This allows us to intercept HDRP's upscaler calls and forward them to our own implementations, without having to modify the HDRP code itself. 
namespace UnityEngine.Rendering.HighDefinition.AMD
{
    public static class AMDUnityPlugin
    {
        private static readonly List<UpscalerPlugin> AvailablePlugins = new()
        {
            new FSR2.FSR2UpscalerPlugin(),
            new FSR3.FSR3UpscalerPlugin(),
        };

        private static UpscalerPlugin _activePlugin = AvailablePlugins[0];
        public static UpscalerPlugin ActivePlugin => _activePlugin;

        public static bool Load() => ActivePlugin.Load();

        public static bool IsLoaded() => ActivePlugin.IsLoaded() || ActivePlugin.isSupported;

        /// <summary>
        /// Get the list of available upscaler plugins, as a read-only list.
        /// Use this to display the available options and retrieve the index of the plugin you want to activate.
        /// </summary>
        public static IReadOnlyList<UpscalerPlugin> GetAvailablePlugins()
        {
            return AvailablePlugins.AsReadOnly();
        }
        
        /// <summary>
        /// Adds an upscaler plugin to the list of available plugins.
        /// If a plugin of the same type already exists in the list, then the new plugin won't be added and the index of the existing plugin will be returned instead.
        /// </summary>
        /// <param name="plugin">The plugin to add</param>
        /// <returns>The index in the plugins list where the plugin can be found. Returns a negative number if the plugin could not be added.</returns>
        public static int AddPlugin(UpscalerPlugin plugin)
        {
            if (plugin == null)
                return -1;

            // Check if we already have a plugin of the same type
            var pluginType = plugin.GetType();
            for (int i = 0; i < AvailablePlugins.Count; ++i)
            {
                if (AvailablePlugins[i].GetType() == pluginType)
                    return i;
            }
            
            AvailablePlugins.Add(plugin);
            return AvailablePlugins.Count - 1;
        }
        
        /// <summary>
        /// Sets the currently active upscaler plugin.
        /// This deactivates the previously active plugin and creates a new device for the newly activated plugin.
        /// If the chosen plugin is already active, then this won't do anything.
        /// </summary>
        /// <param name="pluginIndex">The index in the available plugin list of the upscaler plugin to activate.</param>
        /// <returns>Whether the plugin was successfully changed and activated.</returns>
        public static bool SetActivePlugin(int pluginIndex)
        {
            if (pluginIndex < 0 || pluginIndex >= AvailablePlugins.Count)
                return false;

            var newPlugin = AvailablePlugins[pluginIndex];
            if (newPlugin == null || !newPlugin.isSupported)
                return false;

            if (newPlugin == _activePlugin)
                return true;

            if (!newPlugin.IsLoaded() && !newPlugin.Load())
                return false;
            
            // Hot-swap the upscaler contexts, so HDRP can keep calling the same FSR2Contexts without knowing the underlying implementation has changed.
            GraphicsDevice.device.RecreateFeatures(_activePlugin, newPlugin);
            
            _activePlugin = newPlugin;
            return true;
        }
    }

    public class GraphicsDevice
    {
        private static GraphicsDevice sGraphicsDevice = null;

        public static GraphicsDevice device => sGraphicsDevice;

        public static uint version => 0x00;
        
        private readonly HashSet<FSR2Context> _contexts = new();

        public static GraphicsDevice CreateGraphicsDevice()
        {
            // Lazy initialization of the plugin, because HDRP fails to call AMDUnityPlugin.Load() itself
            UpscalerPlugin plugin = AMDUnityPlugin.ActivePlugin;
            if (plugin != null && !plugin.IsLoaded())
            {
                if (!plugin.Load())
                {
                    Debug.LogWarning("Failed to initialize upscaler plugin");
                    return null;
                }
            }
            
            if (sGraphicsDevice != null)
            {
                sGraphicsDevice.Shutdown();
                sGraphicsDevice.Initialize();
                return sGraphicsDevice;
            }
            
            GraphicsDevice graphicsDevice = new GraphicsDevice();
            if (graphicsDevice.Initialize())
            {
                sGraphicsDevice = graphicsDevice;
                return graphicsDevice;
            }
            
            Debug.LogWarning("Failed to initialize upscaler GraphicsDevice");
            return null;
        }

        private bool Initialize()
        {
#if UNITY_EDITOR
            UnityEditor.AssemblyReloadEvents.beforeAssemblyReload += Shutdown;
#endif
            
            return true;
        }

        private void Shutdown()
        {
#if UNITY_EDITOR
            UnityEditor.AssemblyReloadEvents.beforeAssemblyReload -= Shutdown;
#endif

            foreach (FSR2Context context in _contexts)
            {
                context.Destroy(AMDUnityPlugin.ActivePlugin);
            }
            
            _contexts.Clear();
        }

        internal void RecreateFeatures(UpscalerPlugin oldPlugin, UpscalerPlugin newPlugin)
        {
            foreach (FSR2Context context in _contexts)
            {
                context.Destroy(oldPlugin);
                context.Create(newPlugin);
            }
        }

        public FSR2Context CreateFeature(CommandBuffer cmd, in FSR2CommandInitializationData initSettings)
        {
            FSR2Context context = new FSR2Context(in initSettings);
            if (!context.Create(AMDUnityPlugin.ActivePlugin))
                return null;
            
            _contexts.Add(context);
            return context;
        }

        public void DestroyFeature(CommandBuffer cmd, FSR2Context fsrContext)
        {
            fsrContext.Destroy(AMDUnityPlugin.ActivePlugin);
            _contexts.Remove(fsrContext);
        }

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

        public bool GetRenderResolutionFromQualityMode(FSR2Quality qualityMode, uint displayWidth, uint displayHeight, out uint renderWidth, out uint renderHeight)
        {
            UpscalerPlugin plugin = AMDUnityPlugin.ActivePlugin;
            if (plugin == null)
            {
                renderWidth = displayWidth;
                renderHeight = displayHeight;
                return false;
            }

            return plugin.GetRenderResolutionFromQualityMode(qualityMode, displayWidth, displayHeight, out renderWidth, out renderHeight);
        }

        public float GetUpscaleRatioFromQualityMode(FSR2Quality qualityMode)
        {
            UpscalerPlugin plugin = AMDUnityPlugin.ActivePlugin;
            if (plugin == null)
            {
                return 1.0f;
            }

            return plugin.GetUpscaleRatioFromQualityMode(qualityMode);
        }
    }

    public class FSR2Context
    {
        private UpscalerContext _wrappedContext;

        private FSR2CommandInitializationData _initData;
        public ref FSR2CommandInitializationData initData => ref _initData;

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

        public FSR2Context(in FSR2CommandInitializationData initSettings)
        {
            _initData = initSettings;
        }
        
        public bool Create(UpscalerPlugin plugin)
        {
            _wrappedContext = plugin?.CreateContext(in _initData);
            return _wrappedContext != null;
        }

        public void Destroy(UpscalerPlugin plugin)
        {
            if (_wrappedContext == null)
                return;
            
            plugin?.DestroyContext(_wrappedContext);
            _wrappedContext = null;
        }

        public void Execute(CommandBuffer cmd, in FSR2TextureTable textures)
        {
            _wrappedContext?.Execute(cmd, in executeData, in textures);
        }
    }
    
    public abstract class UpscalerPlugin
    {
        public abstract string name { get; }
        
        public abstract bool isSupported { get; }
        
        public abstract bool Load();

        public abstract bool IsLoaded();

        public abstract UpscalerContext CreateContext(in FSR2CommandInitializationData initSettings);

        public abstract void DestroyContext(UpscalerContext context);
        
        public abstract bool GetRenderResolutionFromQualityMode(FSR2Quality qualityMode, uint displayWidth, uint displayHeight, out uint renderWidth, out uint renderHeight);

        public abstract float GetUpscaleRatioFromQualityMode(FSR2Quality qualityMode);
    }

    public abstract class UpscalerContext
    {
        public abstract void Execute(CommandBuffer cmd, in FSR2CommandExecutionData executeData, in FSR2TextureTable textures);
    }

    public struct FSR2CommandInitializationData
    {
        public uint displaySizeHeight;

        public uint displaySizeWidth;

        public FfxFsr2InitializationFlags ffxFsrFlags;

        public uint maxRenderSizeHeight;

        public uint maxRenderSizeWidth;

        public readonly bool GetFlag(FfxFsr2InitializationFlags flag)
        {
            return (ffxFsrFlags & flag) == flag;
        }

        public void SetFlag(FfxFsr2InitializationFlags flag, bool value)
        {
            if (value)
                ffxFsrFlags |= flag;
            else
                ffxFsrFlags &= ~flag;
        }
    }

    [Flags]
    public enum FfxFsr2InitializationFlags
    {
        EnableHighDynamicRange = 1,
        EnableDisplayResolutionMotionVectors = 2,
        EnableMotionVectorsJitterCancellation = 4,
        DepthInverted = 8,
        EnableDepthInfinite = 16, // 0x00000010
        EnableAutoExposure = 32, // 0x00000020
        EnableDynamicResolution = 64, // 0x00000040
        EnableTexture1DUsage = 128, // 0x00000080
    }

    public struct FSR2TextureTable
    {
        public Texture biasColorMask;
        public Texture colorInput;
        public Texture colorOutput;
        public Texture depth;
        public Texture exposureTexture;
        public Texture motionVectors;
        public Texture reactiveMask;
        public Texture transparencyMask;
    }

    public enum FSR2Quality
    {
        Quality,
        Balanced,
        Performance,
        UltraPerformance,
    }

    public struct FSR2CommandExecutionData
    {
        public float jitterOffsetX;
        public float jitterOffsetY;
        public float MVScaleX;
        public float MVScaleY;
        public uint renderSizeWidth;
        public uint renderSizeHeight;
        public int enableSharpening;
        public float sharpness;
        public float frameTimeDelta;
        public float preExposure;
        public int reset;
        public float cameraNear;
        public float cameraFar;
        public float cameraFovAngleVertical;
    }
}