using System;
using System.Runtime.InteropServices;

#if UNITY_STANDALONE_WIN
namespace UnityEngine.Rendering.HighDefinition.AMD.XeSS
{
    public class XeSSUpscalerPlugin: UpscalerPlugin
    {
        public override string name => "XeSS 2.1";
        public override bool isSupported => SystemInfo.graphicsDeviceType == GraphicsDeviceType.Direct3D12;
        public override bool isMLBased => true;

        private static bool _libraryLoaded = false;
        internal static int BaseEventId = -1;

        public override bool Load()
        {
            if (_libraryLoaded)
                return true;

            if (!XeSSLibrary.InitApi())
                return false;

            BaseEventId = XeSSLibrary.GetBaseEventId();
            _libraryLoaded = BaseEventId >= 0;
            return _libraryLoaded;
        }

        public override bool IsLoaded()
        {
            return _libraryLoaded;
        }

        ~XeSSUpscalerPlugin()
        {
            if (_libraryLoaded)
            {
                XeSSLibrary.ShutdownApi();
                _libraryLoaded = false;
            }
        }

        public override UpscalerContext CreateContext(in FSR2CommandInitializationData initSettings)
        {
            XeSSUpscalerContext context = new();
            context.Init(in initSettings);
            return context;
        }

        public override void DestroyContext(UpscalerContext context)
        {
            ((XeSSUpscalerContext)context).Destroy();
        }

        public override bool GetRenderResolutionFromQualityMode(FSR2Quality qualityMode, uint displayWidth, uint displayHeight, out uint renderWidth, out uint renderHeight)
        {
            float ratio = GetUpscaleRatioFromQualityMode(qualityMode);
            renderWidth = (uint)Mathf.RoundToInt(displayWidth / ratio);
            renderHeight = (uint)Mathf.RoundToInt(displayHeight / ratio);
            return true;
        }

        public override float GetUpscaleRatioFromQualityMode(FSR2Quality qualityMode)
        {
            return XeSSLibrary.ConvertQuality(qualityMode) switch
            {
                XeSSLibrary.QualitySetting.NativeAA => 1.0f,
                XeSSLibrary.QualitySetting.UltraQualityPlus => 1.3f,
                XeSSLibrary.QualitySetting.UltraQuality => 1.5f,
                XeSSLibrary.QualitySetting.Quality => 1.7f,
                XeSSLibrary.QualitySetting.Balanced => 2.0f,
                XeSSLibrary.QualitySetting.Performance => 2.3f,
                XeSSLibrary.QualitySetting.UltraPerformance => 3.0f,
                _ => 1.0f,
            };
        }
    }

    public class XeSSUpscalerContext : UpscalerContext
    {
        private readonly NativeData<XeSSLibrary.InitParams> _nativeInitParams = new();
        private readonly NativeData<XeSSLibrary.ExecuteParams> _nativeExecParams = new();

        private bool _initialized;

        internal void Init(in FSR2CommandInitializationData initSettings)
        {
            XeSSLibrary.InitFlags initFlags = XeSSLibrary.InitFlags.ResponsivePixelMask;
            if (initSettings.GetFlag(FfxFsr2InitializationFlags.DepthInverted)) initFlags |= XeSSLibrary.InitFlags.InvertedDepth;
            if (initSettings.GetFlag(FfxFsr2InitializationFlags.EnableDisplayResolutionMotionVectors)) initFlags |= XeSSLibrary.InitFlags.HighResMotionVectors;
            if (initSettings.GetFlag(FfxFsr2InitializationFlags.EnableMotionVectorsJitterCancellation)) initFlags |= XeSSLibrary.InitFlags.JitteredMotionVectors;
            if (!initSettings.GetFlag(FfxFsr2InitializationFlags.EnableHighDynamicRange)) initFlags |= XeSSLibrary.InitFlags.LDRInputColor;
            if (initSettings.GetFlag(FfxFsr2InitializationFlags.EnableAutoExposure)) initFlags |= XeSSLibrary.InitFlags.EnableAutoExposure;

            //Debug.Log($"Setting up XeSS with input size: {initSettings.maxRenderSizeWidth}x{initSettings.maxRenderSizeHeight}, output size: {initSettings.displaySizeWidth}x{initSettings.displaySizeHeight}, flags: {initFlags}");

            ref var initParams = ref _nativeInitParams.Value;
            initParams.maxRenderSizeWidth = initSettings.maxRenderSizeWidth;
            initParams.maxRenderSizeHeight = initSettings.maxRenderSizeHeight;
            initParams.displaySizeWidth = initSettings.displaySizeWidth;
            initParams.displaySizeHeight = initSettings.displaySizeHeight;
            initParams.flags = (int)initFlags;
            initParams.featureSlot = _nativeExecParams.Value.featureSlot = XeSSLibrary.CreateFeatureSlot();

            using CommandBuffer cmd = new();
            cmd.IssuePluginEventAndData(XeSSLibrary.GetRenderEventAndDataFunc(), XeSSUpscalerPlugin.BaseEventId + (int)XeSSLibrary.PluginEvent.Init, _nativeInitParams.Ptr);
            Graphics.ExecuteCommandBuffer(cmd);

            _initialized = true;
        }

        internal void Destroy()
        {
            if (_initialized)
            {
                using CommandBuffer cmd = new();
                cmd.IssuePluginEventAndData(XeSSLibrary.GetRenderEventAndDataFunc(), XeSSUpscalerPlugin.BaseEventId + (int)XeSSLibrary.PluginEvent.Destroy, (IntPtr)_nativeInitParams.Value.featureSlot);
                Graphics.ExecuteCommandBuffer(cmd);

                _initialized = false;
            }

            _nativeInitParams.Dispose();
            _nativeExecParams.Dispose();
        }

        public override void Execute(CommandBuffer cmd, in FSR2CommandExecutionData executeData, in FSR2TextureTable textures)
        {
            cmd.BeginSample("XeSS");

            SetTexture(cmd, XeSSLibrary.Texture.ColorInput, textures.colorInput, true);
            SetTexture(cmd, XeSSLibrary.Texture.Depth, textures.depth);
            SetTexture(cmd, XeSSLibrary.Texture.MotionVectors, textures.motionVectors);
            SetTexture(cmd, XeSSLibrary.Texture.ExposureTexture, textures.exposureTexture);
            SetTexture(cmd, XeSSLibrary.Texture.ReactiveMask, textures.reactiveMask);
            SetTexture(cmd, XeSSLibrary.Texture.BiasColorMask, textures.biasColorMask);
            SetTexture(cmd, XeSSLibrary.Texture.TransparencyMask, textures.transparencyMask);
            SetTexture(cmd, XeSSLibrary.Texture.ColorOutput, textures.colorOutput);

            ref var execParams = ref _nativeExecParams.Value;
            execParams.jitterOffsetX = executeData.jitterOffsetX;
            execParams.jitterOffsetY = executeData.jitterOffsetY;
            execParams.MVScaleX = executeData.MVScaleX;
            execParams.MVScaleY = executeData.MVScaleY;
            execParams.renderSizeWidth = executeData.renderSizeWidth;
            execParams.renderSizeHeight = executeData.renderSizeHeight;
            execParams.preExposure = executeData.preExposure;
            execParams.reset = executeData.reset;

            cmd.IssuePluginEventAndData(XeSSLibrary.GetRenderEventAndDataFunc(), XeSSUpscalerPlugin.BaseEventId + (int)XeSSLibrary.PluginEvent.Execute, _nativeExecParams.Ptr);
            cmd.EndSample("XeSS");
        }

        private void SetTexture(CommandBuffer cmd, XeSSLibrary.Texture textureSlot, Texture texture, bool clearTextureTable = false)
        {
            if (texture == null)
                return;

            uint userData = (_nativeInitParams.Value.featureSlot & ushort.MaxValue) << 16 | ((uint)textureSlot & (uint)short.MaxValue) << 1 | (clearTextureTable ? 1u : 0u);
            cmd.IssuePluginCustomTextureUpdateV2(XeSSLibrary.GetSetTextureEventFunc(), texture, userData);
        }
    }

    internal static class XeSSLibrary
    {
        private const string DllName = "XeSS2UnityPlugin";

        [DllImport(DllName, EntryPoint = "XeSSUP_InitApi")]
        public static extern bool InitApi();

        [DllImport(DllName, EntryPoint = "XeSSUP_ShutdownApi")]
        public static extern void ShutdownApi();

        [DllImport(DllName, EntryPoint = "XeSSUP_GetRenderEventCallback")]
        public static extern IntPtr GetRenderEventAndDataFunc();

        [DllImport(DllName, EntryPoint = "XeSSUP_GetSetTextureEventCallback")]
        public static extern IntPtr GetSetTextureEventFunc();

        [DllImport(DllName, EntryPoint = "XeSSUP_CreateFeatureSlot")]
        public static extern uint CreateFeatureSlot();

        [DllImport(DllName, EntryPoint = "XeSSUP_GetBaseEventId")]
        public static extern int GetBaseEventId();

        public enum PluginEvent
        {
            Destroy,
            Execute,
            PostExecute,
            Init,
        };

        public enum QualitySetting
        {
            UltraPerformance = 100,
            Performance = 101,
            Balanced = 102,
            Quality = 103,
            UltraQuality = 104,
            UltraQualityPlus = 105,
            NativeAA = 106,
        }

        [Flags]
        public enum InitFlags
        {
            None = 0,
            /** Use motion vectors at target resolution. */
            HighResMotionVectors = 1 << 0,
            /** Use inverted (increased precision) depth encoding */
            InvertedDepth = 1 << 1,
            /** Use exposure texture to scale input color. */
            ExposureScaleTexture = 1 << 2,
            /** Use responsive pixel mask texture. */
            ResponsivePixelMask = 1 << 3,
            /** Use velocity in NDC */
            UseNDCVelocity = 1 << 4,
            /** Use external descriptor heap */
            ExternalDescriptorHeap = 1 << 5,
            /** Disable tonemapping for input and output */
            LDRInputColor = 1 << 6,
            /** Remove jitter from input velocity*/
            JitteredMotionVectors = 1 << 7,
            /** Enable automatic exposure calculation. */
            EnableAutoExposure = 1 << 8
        }

        public enum Texture
        {
            ColorInput,
            ColorOutput,
            Depth,
            MotionVectors,
            TransparencyMask,
            ExposureTexture,
            ReactiveMask,
            BiasColorMask,
        };

        [Serializable, StructLayout(LayoutKind.Sequential)]
        public struct InitParams
        {
            public uint maxRenderSizeWidth;
            public uint maxRenderSizeHeight;
            public uint displaySizeWidth;
            public uint displaySizeHeight;
            public int flags;

            public uint featureSlot;
        }

        [Serializable, StructLayout(LayoutKind.Sequential)]
        public struct ExecuteParams
        {
            public float jitterOffsetX;
            public float jitterOffsetY;
            public float MVScaleX;
            public float MVScaleY;
            public uint renderSizeWidth;
            public uint renderSizeHeight;
            public float preExposure;
            public int reset;

            public uint featureSlot;
        }

        public static QualitySetting ConvertQuality(FSR2Quality quality)
        {
            switch (quality)
            {
                case FSR2Quality.Quality:
                    return QualitySetting.Quality;
                case FSR2Quality.Balanced:
                    return QualitySetting.Balanced;
                case FSR2Quality.Performance:
                    return QualitySetting.Performance;
                case FSR2Quality.UltraPerformance:
                    return QualitySetting.UltraPerformance;
                default:
                    return QualitySetting.NativeAA;
            }
        }
    }

    internal class NativeData<T> : IDisposable where T : struct
    {
        private IntPtr _marshalledValue = Marshal.AllocHGlobal(Marshal.SizeOf(typeof(T)));
        public T Value = new T();

        public IntPtr Ptr
        {
            get
            {
                Marshal.StructureToPtr(Value, _marshalledValue, true);
                return _marshalledValue;
            }
        }

        public void Dispose()
        {
            Dispose(true);
            GC.SuppressFinalize(this);
        }

        protected virtual void Dispose(bool disposing)
        {
            if (_marshalledValue == IntPtr.Zero)
                return;

            Marshal.FreeHGlobal(_marshalledValue);
            _marshalledValue = IntPtr.Zero;
        }

        ~NativeData() => Dispose(false);
    }
}
#endif  // UNITY_STANDALONE_WIN