using System;
using System.Runtime.InteropServices;
using UnityEngine.Experimental.Rendering;

namespace UnityEngine.Rendering.PostProcessing
{
	internal class PSSRUpscaler : Upscaler
	{
#if UNITY_PS5
		public static bool IsSupported => UnityEngine.PS5.Utility.isTrinityMode && PluginInitialized;
		
		private static readonly bool PluginInitialized;

		private PSSRPlugin.DispatchParams _dispatchParams;
		private IntPtr _dispatchParamsBuffer;

		private RenderTexture _inputColor;
		private readonly RenderTexture[] _inputDepth = new RenderTexture[2];
		private readonly RenderTexture[] _inputMotionVectors = new RenderTexture[2];
		private RenderTexture _outputColor;
		private Texture2D _outputIntermediate;
		
		private bool _contextInitialized;
		private uint _frameCount;

		static PSSRUpscaler()
		{
			if (PSSRPlugin.InitPssr() < 0)
			{
				Debug.LogError("Failed to initialize PSSR plugin!");
				PluginInitialized = false;
			}
			
			PluginInitialized = true;
		}	

		public override void CreateContext(PostProcessRenderContext context, Upscaling config)
		{
			if (!PluginInitialized)
			{
				Debug.LogWarning("PSSR plugin is not initialized!");
				return;
			}

			PSSRPlugin.InitParams initParams;
			initParams.displayWidth = (uint)config.UpscaleSize.x;
			initParams.displayHeight = (uint)config.UpscaleSize.y;
			initParams.maxRenderWidth = (uint)config.MaxRenderSize.x;
			initParams.maxRenderHeight = (uint)config.MaxRenderSize.y;
			initParams.autoKeepCopies = 0u;		// We use double buffered depth and motion vector copies created by Unity

			CreateRenderTexture(ref _inputColor, "PSSR Input Color", config.MaxRenderSize, context.sourceFormat, true);
			CreateRenderTextureArray( _inputDepth, "PSSR Input Depth", config.MaxRenderSize, GraphicsFormat.R32_SFloat, true);
			CreateRenderTextureArray(_inputMotionVectors, "PSSR Input Motion Vectors", config.MaxRenderSize, GraphicsFormat.R16G16_SFloat, true);
			CreateRenderTexture(ref _outputColor, "PSSR Output Color", config.UpscaleSize, RenderTextureFormat.RGB111110Float);

			if (PSSRPlugin.CreatePssrContext(ref initParams, out IntPtr outputColorTexturePtr) >= 0 && outputColorTexturePtr != IntPtr.Zero)
			{
				// PSSR requires an output color texture in a very particular format (k11_11_10Float with kStandard256B tile mode and a specific alignment) that Unity cannot create directly.
				// So instead we let the plugin create that texture and then import it into Unity as a generic 32bpp texture from a native pointer.
				_outputIntermediate = Texture2D.CreateExternalTexture(config.UpscaleSize.x, config.UpscaleSize.y, TextureFormat.RGBA32, false, true, outputColorTexturePtr);
				_dispatchParamsBuffer = Marshal.AllocHGlobal(Marshal.SizeOf<PSSRPlugin.DispatchParams>());
				_contextInitialized = true;
			}
		}

		public override void DestroyContext()
		{
			base.DestroyContext();
			
			// TODO: this still race conditions sometimes
			if (_contextInitialized)
			{
				var cmd = new CommandBuffer();
				cmd.IssuePluginEventAndData(PSSRPlugin.GetRenderEventAndDataFunc(), 2, IntPtr.Zero);
				Graphics.ExecuteCommandBuffer(cmd);
				cmd.Release();
				
				_contextInitialized = false;
			}

			if (_outputIntermediate != null)
			{
				Object.Destroy(_outputIntermediate);
				_outputIntermediate = null;
			}

			DestroyRenderTexture(ref _outputColor);
			DestroyRenderTextureArray(_inputMotionVectors);
			DestroyRenderTextureArray(_inputDepth);
			DestroyRenderTexture(ref _inputColor);

			if (_dispatchParamsBuffer != IntPtr.Zero)
			{
				Marshal.FreeHGlobal(_dispatchParamsBuffer);
				_dispatchParamsBuffer = IntPtr.Zero;
			}
		}

		public override void Render(PostProcessRenderContext context, Upscaling config)
		{
			var cmd = context.command;
			if (!PluginInitialized || !_contextInitialized)
			{
				cmd.BlitFullscreenTriangle(context.source, context.destination);
				return;
			}
			
			cmd.BeginSample("PSSR");

			if (config.Reset || _frameCount == 0)
			{
				cmd.SetRenderTarget(_inputDepth[0]);
				cmd.ClearRenderTarget(false, true, Color.clear);
				cmd.SetRenderTarget(_inputDepth[1]);
				cmd.ClearRenderTarget(false, true, Color.clear);
				cmd.SetRenderTarget(_inputMotionVectors[0]);
				cmd.ClearRenderTarget(false, true, Color.clear);
				cmd.SetRenderTarget(_inputMotionVectors[1]);
				cmd.ClearRenderTarget(false, true, Color.clear);
			}

			int frameIndex = (int)(_frameCount++ % 2);
			
			// We need to provide copies of the previous depth and motion vector buffers anyway, so we can turn this otherwise wasteful copying into a benefit
			PrepareInputs(cmd, context, config, _inputColor, _inputDepth[frameIndex], _inputMotionVectors[frameIndex]);

			Texture reactiveMask = config.reactiveMask;
			if (config.autoGenerateReactiveMask || config.autoGenerateTransparencyAndComposition)
			{
				reactiveMask = GenerateReactiveMask(cmd, context, config);
			}

			var flags = PSSRPlugin.OptionFlags.None;
			if (SystemInfo.usesReversedZBuffer) flags |= PSSRPlugin.OptionFlags.ReverseDepth;
			if (config.exposureSource == Upscaling.ExposureSource.Auto) flags |= PSSRPlugin.OptionFlags.AutoExposure;
			
			_dispatchParams.color = ToNativePtr(_inputColor);
			_dispatchParams.depth = ToNativePtr(_inputDepth[frameIndex]);
			_dispatchParams.prevDepth = ToNativePtr(_inputDepth[frameIndex ^ 1]);
			_dispatchParams.motionVectors = ToNativePtr(_inputMotionVectors[frameIndex]);
			_dispatchParams.prevMotionVectors = ToNativePtr(_inputMotionVectors[frameIndex ^ 1]);
			_dispatchParams.exposure = ToNativePtr(config.exposureSource switch
			{
				Upscaling.ExposureSource.Manual when config.exposure != null => config.exposure,
				Upscaling.ExposureSource.Unity => context.autoExposureTexture,
				_ => null,
			});
			_dispatchParams.reactiveMask = ToNativePtr(reactiveMask);
			_dispatchParams.outputColor = ToNativePtr(_outputIntermediate);

			var scaledRenderSize = config.GetScaledRenderSize(context.camera);
			_dispatchParams.renderWidth = (uint)scaledRenderSize.x;
			_dispatchParams.renderHeight = (uint)scaledRenderSize.y;
			_dispatchParams.jitter = config.JitterOffset;
			_dispatchParams.motionVectorScale = new Vector2(-scaledRenderSize.x, -scaledRenderSize.y);
			_dispatchParams.FromCamera(context.camera);
			_dispatchParams.preExposure = config.preExposure;
			_dispatchParams.resetHistory = config.Reset ? 1u : 0u;
			_dispatchParams.flags = flags;
			
			Marshal.StructureToPtr(_dispatchParams, _dispatchParamsBuffer, false);
			cmd.IssuePluginEventAndData(PSSRPlugin.GetRenderEventAndDataFunc(), 1, _dispatchParamsBuffer);
			
			// Convert the native texture into Unity-land using a blind texture data copy
			cmd.CopyTexture(_outputIntermediate, _outputColor);
			
			// TODO: optional sharpening
			cmd.BlitFullscreenTriangle(_outputColor, context.destination);
			
			cmd.EndSample("PSSR");
		}
		
		private static IntPtr ToNativePtr(Texture texture)
		{
			return texture != null ? texture.GetNativeTexturePtr() : IntPtr.Zero;
		}

		private static class PSSRPlugin
		{
			private const string LibraryName = "PSSRPlugin.prx";

			[DllImport(LibraryName)]
			public static extern IntPtr GetRenderEventAndDataFunc();

			[DllImport(LibraryName)]
			public static extern int InitPssr();

			[DllImport(LibraryName)]
			public static extern void ReleasePssr();

			[DllImport(LibraryName)]
			public static extern int CreatePssrContext(ref InitParams initParams, out IntPtr outputColorTexturePtr);

			[DllImport(LibraryName)]
			public static extern void DestroyPssrContext();

			[Serializable, StructLayout(LayoutKind.Sequential)]
			public struct InitParams
			{
				public uint displayWidth;
				public uint displayHeight;
				public uint maxRenderWidth;
				public uint maxRenderHeight;

				public uint autoKeepCopies;
			}

			[Serializable, StructLayout(LayoutKind.Sequential)]
			public struct DispatchParams
			{
				public IntPtr color;
				public IntPtr depth;
				public IntPtr prevDepth;
				public IntPtr motionVectors;
				public IntPtr prevMotionVectors;
				public IntPtr exposure;
				public IntPtr reactiveMask;
				public IntPtr outputColor;

				public uint renderWidth;
				public uint renderHeight;

				public Vector2 jitter;
				public Vector2 motionVectorScale;

				public Matrix4x4 camProjectionNoJitter;
				public Vector3 camForward;
				public Vector3 camUp;
				public Vector3 camRight;
				public double camPositionX;
				public double camPositionY;
				public double camPositionZ;
				public float camNear;
				public float camFar;
				public float preExposure;
				public uint resetHistory;
				public OptionFlags flags;

				public void FromCamera(Camera cam)
				{
					camProjectionNoJitter = cam.nonJitteredProjectionMatrix;
					camForward = cam.transform.forward;
					camUp = cam.transform.up;
					camRight = cam.transform.right;
					camPositionX = cam.transform.position.x;
					camPositionY = cam.transform.position.y;
					camPositionZ = cam.transform.position.z;
					camNear = cam.nearClipPlane;
					camFar = cam.farClipPlane;
				}
			}
			
			[Flags]
			public enum OptionFlags: uint
			{
				None = 0u,
				ViewSpaceDepth = 1u << 14,
				ReverseDepth = 1u << 15,
				AutoExposure = 1u << 16,
				PassThrough = 1u << 31,
			}
		}
#else
		public static bool IsSupported => false;
		
		public override void CreateContext(PostProcessRenderContext context, Upscaling config)
		{
			throw new NotImplementedException();
		}
		
		public override void DestroyContext()
		{
			base.DestroyContext();
		}
		
		public override void Render(PostProcessRenderContext context, Upscaling config)
		{
			throw new NotImplementedException();
		}
#endif
	}
}