using ArmASR;
using UnityEngine.Experimental.Rendering;

namespace UnityEngine.Rendering.PostProcessing
{
	internal abstract class ASRUpscaler: Upscaler
	{
		public static bool IsSupported => SystemInfo.supportsComputeShaders;
		
		protected abstract Asr.Variant Variant { get; }
		
		private AsrContext _asrContext;
		
		private Asr.DispatchDescription _dispatchDescription = new();
		private Asr.GenerateReactiveDescription _genReactiveDescription = new();
		
		public override void CreateContext(PostProcessRenderContext context, Upscaling config)
		{
			// Initialize ASR context
			Asr.InitializationFlags flags = 0;
			if (context.camera.allowHDR) flags |= Asr.InitializationFlags.EnableHighDynamicRange;
			if (config.exposureSource == Upscaling.ExposureSource.Auto) flags |= Asr.InitializationFlags.EnableAutoExposure;
			if (RuntimeUtilities.IsDynamicResolutionEnabled(context.camera)) flags |= Asr.InitializationFlags.EnableDynamicResolution;

			_asrContext = Asr.CreateContext(Variant, config.UpscaleSize, config.MaxRenderSize, context.resources.asrUpscalerShaders, flags);
		}

		public override void DestroyContext()
		{
			base.DestroyContext();
			
			if (_asrContext != null)
			{
				_asrContext.Destroy();
				_asrContext = null;
			}
		}

		public override void Render(PostProcessRenderContext context, Upscaling config)
		{
			var cmd = context.command;
			cmd.BeginSample("ASR");
			
			SetupDispatchDescription(context, config);
			
			if (config.autoGenerateReactiveMask)
			{
				SetupAutoReactiveDescription(context, config);
				
				var scaledRenderSize = _genReactiveDescription.RenderSize;
				cmd.GetTemporaryRT(AsrShaderIDs.UavAutoReactive, scaledRenderSize.x, scaledRenderSize.y, 0, default, GraphicsFormat.R8_UNorm, 1, true);
				_asrContext.GenerateReactiveMask(_genReactiveDescription, cmd);
				_dispatchDescription.Reactive = new ResourceView(AsrShaderIDs.UavAutoReactive);
			}
			
			_asrContext.Dispatch(_dispatchDescription, cmd);
			
			cmd.EndSample("ASR");
		}
		
        private void SetupDispatchDescription(PostProcessRenderContext context, Upscaling config)
        {
            var camera = context.camera;
            
            // Set up the main ASR dispatch parameters
            _dispatchDescription.Color = new ResourceView(context.source);
            _dispatchDescription.Depth = new ResourceView(Upscaling.GetDepthTexture(context.camera), RenderTextureSubElement.Depth);
            _dispatchDescription.MotionVectors = new ResourceView(BuiltinRenderTextureType.MotionVectors);
            _dispatchDescription.Exposure = ResourceView.Unassigned;
            _dispatchDescription.Reactive = ResourceView.Unassigned;
            _dispatchDescription.TransparencyAndComposition = ResourceView.Unassigned;

            if (config.exposureSource == Upscaling.ExposureSource.Manual && config.exposure != null) _dispatchDescription.Exposure = new ResourceView(config.exposure);
            if (config.exposureSource == Upscaling.ExposureSource.Unity) _dispatchDescription.Exposure = new ResourceView(context.autoExposureTexture);
            if (config.reactiveMask != null) _dispatchDescription.Reactive = new ResourceView(config.reactiveMask);
            if (config.transparencyAndCompositionMask != null) _dispatchDescription.TransparencyAndComposition = new ResourceView(config.transparencyAndCompositionMask);

            var scaledRenderSize = config.GetScaledRenderSize(context.camera);
            
            _dispatchDescription.Output = new ResourceView(context.destination);
            _dispatchDescription.PreExposure = config.preExposure;
            _dispatchDescription.EnableSharpening = config.performSharpenPass;
            _dispatchDescription.Sharpness = config.sharpness;
            _dispatchDescription.JitterOffset = config.JitterOffset;
            _dispatchDescription.MotionVectorScale.x = -scaledRenderSize.x;
            _dispatchDescription.MotionVectorScale.y = -scaledRenderSize.y;
            _dispatchDescription.RenderSize = scaledRenderSize;
            _dispatchDescription.InputResourceSize = scaledRenderSize;
            _dispatchDescription.FrameTimeDelta = Time.unscaledDeltaTime;
            _dispatchDescription.CameraNear = camera.nearClipPlane;
            _dispatchDescription.CameraFar = camera.farClipPlane;
            _dispatchDescription.CameraFovAngleVertical = camera.fieldOfView * Mathf.Deg2Rad;
            _dispatchDescription.ViewSpaceToMetersFactor = 1.0f; // 1 unit is 1 meter in Unity
            _dispatchDescription.Reset = config.Reset;

            if (SystemInfo.usesReversedZBuffer)
            {
                // Swap the near and far clip plane distances as FSR2 expects this when using inverted depth
                (_dispatchDescription.CameraNear, _dispatchDescription.CameraFar) = (_dispatchDescription.CameraFar, _dispatchDescription.CameraNear);
            }
        }

        private void SetupAutoReactiveDescription(PostProcessRenderContext context, Upscaling config)
        {
            // Set up the parameters to auto-generate a reactive mask
            _genReactiveDescription.ColorOpaqueOnly = new ResourceView(config.ColorOpaqueOnly);
            _genReactiveDescription.ColorPreUpscale = new ResourceView(context.source);
            _genReactiveDescription.OutReactive = new ResourceView(AsrShaderIDs.UavAutoReactive);
            _genReactiveDescription.RenderSize = config.GetScaledRenderSize(context.camera);
            _genReactiveDescription.Scale = config.generateReactiveParameters.scale;
            _genReactiveDescription.CutoffThreshold = config.generateReactiveParameters.cutoffThreshold;
            _genReactiveDescription.BinaryValue = config.generateReactiveParameters.binaryValue;
            _genReactiveDescription.Flags = (Asr.GenerateReactiveFlags)config.generateReactiveParameters.flags;
        }
	}

	internal class ASRUpscaler_Quality : ASRUpscaler
	{
		protected override Asr.Variant Variant => Asr.Variant.Quality;
	}

	internal class ASRUpscaler_Balanced : ASRUpscaler
	{
		protected override Asr.Variant Variant => Asr.Variant.Balanced;
	}
	
	internal class ASRUpscaler_Performance : ASRUpscaler
	{
		protected override Asr.Variant Variant => Asr.Variant.Performance;
	}
}