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

namespace FidelityFX
{
    internal abstract class Fsr2Pipeline: IDisposable
    {
        protected ComputeShader ComputeShader;
        protected int KernelIndex;

        private readonly Fsr2Callbacks _callbacks;

        protected Fsr2Pipeline(Fsr2Callbacks callbacks)
        {
            _callbacks = callbacks;
        }

        public virtual void Dispose()
        {
            UnloadComputeShader();
        }
        
        public abstract void ScheduleDispatch(CommandBuffer commandBuffer, Fsr2.DispatchDescription dispatchParams, int dispatchX, int dispatchY);
        
        protected void LoadComputeShader(string name)
        {
            LoadComputeShader(name, ref ComputeShader, out KernelIndex);
        }
        
        protected void LoadComputeShader(string name, ref ComputeShader shaderRef, out int kernelIndex)
        {
            if (shaderRef == null)
                shaderRef = _callbacks.LoadComputeShader(name);

            kernelIndex = shaderRef.FindKernel("CS");
        }

        protected void UnloadComputeShader()
        {
            UnloadComputeShader(ref ComputeShader);
        }
        
        protected void UnloadComputeShader(ref ComputeShader shaderRef)
        {
            if (shaderRef == null)
                return;

            _callbacks.UnloadComputeShader(shaderRef);
            shaderRef = null;
        }
    }

    internal class Fsr2RcasPipeline : Fsr2Pipeline
    {
        private readonly ComputeBuffer _fsr2Constants;
        private readonly ComputeBuffer _rcasConstants;
        
        public Fsr2RcasPipeline(Fsr2Callbacks callbacks, ComputeBuffer fsr2Constants, ComputeBuffer rcasConstants)
            : base(callbacks)
        {
            _fsr2Constants = fsr2Constants;
            _rcasConstants = rcasConstants;
            
            LoadComputeShader("FSR2/ffx_fsr2_rcas_pass");
        }

        public override void ScheduleDispatch(CommandBuffer commandBuffer, Fsr2.DispatchDescription dispatchParams, int dispatchX, int dispatchY)
        {
            // Run the RCAS sharpening filter on the upscaled image
            ComputeShader.SetTexture(KernelIndex, "r_input_exposure", dispatchParams.Exposure);
            ComputeShader.SetTexture(KernelIndex, "r_rcas_input", dispatchParams.Input);
            ComputeShader.SetTexture(KernelIndex, "rw_upscaled_output", dispatchParams.Output);
            ComputeShader.SetConstantBuffer("cbFSR2", _fsr2Constants, 0, Marshal.SizeOf<Fsr2.Fsr2Constants>());
            ComputeShader.SetConstantBuffer("cbRCAS", _rcasConstants, 0, Marshal.SizeOf<Fsr2.RcasConstants>());

            commandBuffer.DispatchCompute(ComputeShader, KernelIndex, dispatchX, dispatchY, 1);
        }
    }
}