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 1.3";
        public override bool isSupported => SystemInfo.graphicsDeviceType == GraphicsDeviceType.Direct3D12;

        private XeSSLibrary.Version _version;
        
        public override bool Load()
        {
            return XeSSLibrary.GetUpscalerVersion(out _version) == 0;
        }

        public override bool IsLoaded()
        {
            return _version.major != 0 && _version.minor != 0;
        }

        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 IntPtr _contextHandle = IntPtr.Zero;

        private XeSSLibrary.ExecuteParams _executeParams;
        private IntPtr _paramsBuffer;

        internal void Init(in FSR2CommandInitializationData initSettings)
        {
            Vector2Int outputResolution = new Vector2Int((int)initSettings.displaySizeWidth, (int)initSettings.displaySizeHeight);

            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;
            // TODO: use NDC velocity? I think motion vectors might be in normalized device coordinates...

            Debug.Log($"Setting up XeSS with input size: {initSettings.maxRenderSizeWidth}x{initSettings.maxRenderSizeHeight}, output size: {initSettings.displaySizeWidth}x{initSettings.displaySizeHeight}, flags: {initFlags}");
            
            _contextHandle = XeSSLibrary.CreateContext(outputResolution, XeSSLibrary.QualitySetting.Quality, initFlags);
            _paramsBuffer = Marshal.AllocHGlobal(Marshal.SizeOf<XeSSLibrary.ExecuteParams>());
        }

        internal void Destroy()
        {
            if (_paramsBuffer != IntPtr.Zero)
            {
                Marshal.FreeHGlobal(_paramsBuffer);
                _paramsBuffer = IntPtr.Zero;
            }
            
            if (_contextHandle != IntPtr.Zero)
            {
                //XeSSLibrary.DestroyContext(_contextHandle);
                var cmd = new CommandBuffer();
                cmd.IssuePluginEventAndData(XeSSLibrary.GetRenderEventAndDataFunc(), 2, _contextHandle);
                Graphics.ExecuteCommandBuffer(cmd);
                cmd.Release();
                
                _contextHandle = IntPtr.Zero;
            }
        }
        
        public override void Execute(CommandBuffer cmd, in FSR2CommandExecutionData executeData, in FSR2TextureTable textures)
        {
            _executeParams.contextHandle = _contextHandle;

            _executeParams.pColorTexture = textures.colorInput.ToNativePtr();
            _executeParams.pDepthTexture = textures.depth.ToNativePtr();
            _executeParams.pVelocityTexture = textures.motionVectors.ToNativePtr();
            _executeParams.pExposureScaleTexture = textures.exposureTexture.ToNativePtr();
            _executeParams.pResponsivePixelMaskTexture = textures.biasColorMask.ToNativePtr();
            _executeParams.pOutputTexture = textures.colorOutput.ToNativePtr();

            _executeParams.velocityScaleX = executeData.MVScaleX;
            _executeParams.velocityScaleY = executeData.MVScaleY;
            _executeParams.jitterScaleX = 1f;
            _executeParams.jitterScaleY = -1f;
            _executeParams.jitterOffsetX = executeData.jitterOffsetX;
            _executeParams.jitterOffsetY = executeData.jitterOffsetY;
            _executeParams.exposureScale = executeData.preExposure;
            _executeParams.resetHistory = (uint)executeData.reset;
            _executeParams.inputWidth = executeData.renderSizeWidth;
            _executeParams.inputHeight = executeData.renderSizeHeight;

            Marshal.StructureToPtr(_executeParams, _paramsBuffer, false);

            cmd.IssuePluginEventAndData(XeSSLibrary.GetRenderEventAndDataFunc(), 1, _paramsBuffer);
        }
    }

    internal static class XeSSLibrary
    {
        private const string DllName = "XeSSUnityPlugin";
        
        [DllImport(DllName, CallingConvention = CallingConvention.StdCall)]
        public static extern IntPtr GetRenderEventAndDataFunc();
        
        [DllImport(DllName, CallingConvention = CallingConvention.StdCall)]
        public static extern int GetUpscalerVersion(out Version version);
        
        [DllImport(DllName, CallingConvention = CallingConvention.StdCall)]
        public static extern IntPtr CreateContext(Vector2Int outputResolution, QualitySetting qualitySetting, InitFlags initFlags);
        
        [DllImport(DllName, CallingConvention = CallingConvention.StdCall)]
        public static extern void DestroyContext(IntPtr contextHandle);

        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
        }

        [StructLayout(LayoutKind.Sequential, Pack = 8)]
        public struct Version
        {
            /** A major version increment indicates a new API and potentially a
             * break in functionality. */
            public ushort major;
            /** A minor version increment indicates incremental changes such as
             * optional inputs or flags. This does not break existing functionality. */
            public ushort minor;
            /** A patch version increment may include performance or quality tweaks or fixes for known issues.
             * There's no change in the interfaces.
             * Versions beyond 90 used for development builds to change the interface for the next release.
             */
            public ushort patch;
            /** Reserved for future use. */
            public ushort reserved;
        }
        
        [StructLayout(LayoutKind.Sequential, Pack = 8)]
        public struct ExecuteParams
        {
            public IntPtr contextHandle;

            /** Input color texture. Must be in NON_PIXEL_SHADER_RESOURCE state.*/
            public IntPtr pColorTexture;

            /** Input motion vector texture. Must be in NON_PIXEL_SHADER_RESOURCE state.*/
            public IntPtr pVelocityTexture;

            /** Optional depth texture. Required if XESS_INIT_FLAG_HIGH_RES_MV has not been specified.
             * Must be in NON_PIXEL_SHADER_RESOURCE state.*/
            public IntPtr pDepthTexture;

            /** Optional 1x1 exposure scale texture. Required if XESS_INIT_FLAG_EXPOSURE_TEXTURE has been
             * specified. Must be in NON_PIXEL_SHADER_RESOURCE state */
            public IntPtr pExposureScaleTexture;

            /** Optional responsive pixel mask texture. Required if XESS_INIT_FLAG_RESPONSIVE_PIXEL_MASK
             * has been specified. Must be in NON_PIXEL_SHADER_RESOURCE state */
            public IntPtr pResponsivePixelMaskTexture;

            /** Output texture in target resolution. Must be in UNORDERED_ACCESS state.*/
            public IntPtr pOutputTexture;

            public float jitterScaleX;
            public float jitterScaleY;
            public float velocityScaleX;
            public float velocityScaleY;

            /** Jitter X coordinate in the range [-0.5, 0.5]. */
            public float jitterOffsetX;

            /** Jitter Y coordinate in the range [-0.5, 0.5]. */
            public float jitterOffsetY;

            /** Optional input color scaling. Default is 1. */
            public float exposureScale;

            public float dummy0;

            /** Resets the history accumulation in this frame. */
            public uint resetHistory;

            /** Input color width. */
            public uint inputWidth;

            /** Input color height. */
            public uint inputHeight;

            public uint dummy1;
        }
        
        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;
            }
        }

        public static IntPtr ToNativePtr(this Texture texture)
        {
            return texture != null ? texture.GetNativeTexturePtr() : IntPtr.Zero;
        }
    }
}
#endif  // UNITY_STANDALONE_WIN