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;

		private PSSRPlugin.DispatchParams _dispatchParams;
		private IntPtr _dispatchParamsBuffer;

		private RenderTexture _inputColor;
		private RenderTexture _inputDepth;
		private RenderTexture _outputColor;
		private Texture2D _outputIntermediate;
		
		private bool _pluginInitialized;
		private bool _contextInitialized;

		public override void CreateContext(PostProcessRenderContext context, Upscaling config)
		{
			if (!PSSRPlugin.InitPssr())
			{
				Debug.LogError("Failed to initialize PSSR plugin!");
				_pluginInitialized = false;
				return;
			}

			_pluginInitialized = true;

			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;

			CreateRenderTexture(ref _inputColor, "PSSR Input Color", config.MaxRenderSize, context.sourceFormat, true);
			CreateRenderTexture(ref _inputDepth, "PSSR Input Depth", config.MaxRenderSize, GraphicsFormat.R32_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;
			}

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

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

			if (_pluginInitialized)
			{
				PSSRPlugin.ReleasePssr();
				_pluginInitialized = false;
			}
		}

		private static readonly int CameraMotionVectorsTexture = Shader.PropertyToID("_CameraMotionVectorsTexture");

		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");
			
			// TODO: if PSSR needs a copy of the previous depth and motion vectors anyway, why not just double-buffer those resources ourselves and make good use of these otherwise dumb copies?
			PrepareInputs(cmd, context, config, _inputColor, _inputDepth);
			
			_dispatchParams.color = ToNativePtr(_inputColor);
			_dispatchParams.depth = ToNativePtr(_inputDepth);
			_dispatchParams.motionVectors = ToNativePtr(Shader.GetGlobalTexture(CameraMotionVectorsTexture));
			_dispatchParams.exposure = ToNativePtr(config.exposureSource switch
			{
				Upscaling.ExposureSource.Manual when config.exposure != null => config.exposure,
				Upscaling.ExposureSource.Unity => context.autoExposureTexture,
				_ => null,
			});
			_dispatchParams.reactiveMask = ToNativePtr(config.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;		// TODO: may need to scale this (should be jitter in render pixels apparently)
			_dispatchParams.motionVectorScale = new Vector2(-scaledRenderSize.x, -scaledRenderSize.y);
			_dispatchParams.FromCamera(context.camera);
			_dispatchParams.preExposure = config.preExposure;
			_dispatchParams.resetHistory = config.Reset ? 1u : 0u;
			
			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 bool InitPssr();

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

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

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

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

				public uint renderWidth;
				public uint renderHeight;

				public Vector2 jitter;
				public Vector2 motionVectorScale;

				public Matrix4x4 camProjection;
				public Vector3 camForward;
				public Vector3 camUp;
				public Vector3 camRight;
				public double camPositionX;
				public double camPositionY;
				public double camPositionZ;
				public float pad;

				public float camNear;
				public float camFar;
				public float preExposure;
				public uint resetHistory;

				public void FromCamera(Camera cam)
				{
					camProjection = cam.projectionMatrix;
					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;
					pad = 0f;

					camNear = cam.nearClipPlane;
					camFar = cam.farClipPlane;
				}
			}
		}
#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
	}
}