using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.InteropServices;
using UnityEditor;
using UnityEngine;
using UnityEngine.Rendering;

public class FSR2Thing : MonoBehaviour
{
    [SerializeField]
    private bool performSharpenPass = true;

    [SerializeField, Range(0, 1)]
    private float sharpness = 0.8f;
    
    [HideInInspector]
    public Camera gameCamera;

    [HideInInspector]
    public Camera outputCamera;

    [HideInInspector]
    public float renderScale;

    private RenderTexture _upscaleRT;
    private RenderTexture _rcasOutput;
    private Texture2D _exposure;
    
    private Material _testMaterial;
    private Material TestMaterial
    {
        get
        {
            if (_testMaterial == null)
            {
                var testShader = Shader.Find("FSR2/FSRTest");
                _testMaterial = new Material(testShader);
            }

            return _testMaterial;
        }
    }

    private ComputeShader _rcasComputeShader;
    private ComputeShader RCASComputeShader
    {
        get
        {
            if (_rcasComputeShader == null)
            {
                // TODO: this is nasty, I don't like this. How do we manage/bind compute shaders best? => make a Callbacks class/interface with overridable implementation
                // GPUInstancer used Resources, we modified that to load stuff from asset bundles instead. Maybe provide a custom loader callback interface?
                _rcasComputeShader = Resources.Load<ComputeShader>("FSR2/ffx_fsr2_rcas_pass");
            }

            return _rcasComputeShader;
        }
    }

    private ComputeBuffer _fsr2ConstantsBuffer;
    private readonly Fsr2Constants[] _fsr2ConstantsArray = { new Fsr2Constants() };
    
    private ComputeBuffer _rcasConstantsBuffer;
    private readonly RcasConstants[] _rcasConstantsArray = new RcasConstants[1];

    private void OnEnable()
    {
        RenderPipelineManager.endContextRendering += OnEndContextRendering;

        _upscaleRT = new RenderTexture(Screen.width, Screen.height, 24, RenderTextureFormat.ARGBHalf);
        _upscaleRT.Create();

        _rcasOutput = new RenderTexture(Screen.width, Screen.height, 24, RenderTextureFormat.ARGBHalf);
        _rcasOutput.enableRandomWrite = true;
        _rcasOutput.Create();

        _exposure = new Texture2D(1, 1);
        _exposure.name = "FSR2 Exposure";
        _exposure.SetPixel(0, 0, Color.white);
        _exposure.Apply();

        _fsr2ConstantsBuffer = new ComputeBuffer(1, Marshal.SizeOf<Fsr2Constants>(), ComputeBufferType.Constant);
        _rcasConstantsBuffer = new ComputeBuffer(1, Marshal.SizeOf<RcasConstants>(), ComputeBufferType.Constant);
    }

    private void OnDisable()
    {
        if (_rcasConstantsBuffer != null)
        {
            _rcasConstantsBuffer.Release();
            _rcasConstantsBuffer = null;
        }

        if (_fsr2ConstantsBuffer != null)
        {
            _fsr2ConstantsBuffer.Release();
            _fsr2ConstantsBuffer = null;
        }

        if (_exposure != null)
        {
            Destroy(_exposure);
            _exposure = null;
        }

        if (_rcasOutput != null)
        {
            _rcasOutput.Release();
            _rcasOutput = null;
        }

        if (_upscaleRT != null)
        {
            _upscaleRT.Release();
            _upscaleRT = null;
        }

        RenderPipelineManager.endContextRendering -= OnEndContextRendering;
    }

    // For scriptable rendering pipeline
    private void OnEndContextRendering(ScriptableRenderContext context, List<Camera> cameras)
    {
        Debug.Log($"OnEndContentRendering, cameras = {string.Join(", ", cameras.Select(c => c.name))}");
    }

    // For legacy built-in render pipeline
    private void OnRenderImage(RenderTexture src, RenderTexture dest)
    {
        // Do a dumb upscale first
        Graphics.Blit(gameCamera.targetTexture, _upscaleRT, TestMaterial);

        _fsr2ConstantsArray[0].preExposure = 1.0f;
        _fsr2ConstantsBuffer.SetData(_fsr2ConstantsArray);
        
        if (performSharpenPass)
        {
            int sharpnessIndex = Mathf.RoundToInt(Mathf.Clamp01(sharpness) * (RcasConfigs.Count - 1));
            _rcasConstantsArray[0] = RcasConfigs[sharpnessIndex];
            _rcasConstantsBuffer.SetData(_rcasConstantsArray);
            
            // Run the RCAS sharpening filter on the upscaled image
            int rcasKernel = RCASComputeShader.FindKernel("CS");
            RCASComputeShader.SetTexture(rcasKernel, "r_exposure", _exposure);
            RCASComputeShader.SetTexture(rcasKernel, "r_rcas_input", _upscaleRT);
            RCASComputeShader.SetTexture(rcasKernel, "rw_upscaled_output", _rcasOutput);
            RCASComputeShader.SetConstantBuffer("cbFSR2", _fsr2ConstantsBuffer, 0, Marshal.SizeOf<Fsr2Constants>());
            RCASComputeShader.SetConstantBuffer("cbRCAS", _rcasConstantsBuffer, 0, Marshal.SizeOf<RcasConstants>());

            const int threadGroupWorkRegionDimRcas = 16;
            int threadGroupsX = (Screen.width + threadGroupWorkRegionDimRcas - 1) / threadGroupWorkRegionDimRcas;
            int threadGroupsY = (Screen.height + threadGroupWorkRegionDimRcas - 1) / threadGroupWorkRegionDimRcas;
            
            RCASComputeShader.Dispatch(rcasKernel, threadGroupsX, threadGroupsY, 1);

            // Output sharpened image to screen
            Graphics.Blit(_rcasOutput, dest);
        }
        else
        {
            Graphics.Blit(_upscaleRT, dest);
        }
    }

    [Serializable, StructLayout(LayoutKind.Sequential)]
    private struct RcasConstants
    {
        public RcasConstants(uint sharpness, uint halfSharp)
        {
            this.sharpness = sharpness;
            this.halfSharp = halfSharp;
            dummy0 = dummy1 = 0;
        }
        
        public readonly uint sharpness;
        public readonly uint halfSharp;
        public readonly uint dummy0;
        public readonly uint dummy1;
    }

    [Serializable, StructLayout(LayoutKind.Sequential)]
    private struct Fsr2Constants
    {
        public Vector2Int renderSize;
        public Vector2Int displaySize;
        public uint lumaMipDimensionsX, lumaMipDimensionsY;
        public uint lumaMipLevelToUse;
        public uint frameIndex;
        public Vector2 displaySizeRcp;
        public Vector2 jitterOffset;
        public Vector4 deviceToViewDepth;
        public Vector2 depthClipUVScale;
        public Vector2 postLockStatusUVScale;
        public Vector2 reactiveMaskDimRcp;
        public Vector2 motionVectorScale;
        public Vector2 downscaleFactor;
        public float preExposure;
        public float tanHalfFOV;
        public Vector2 motionVectorJitterCancellation;
        public float jitterPhaseCount;
        public float lockInitialLifetime;
        public float lockTickDelta;
        public float deltaTime;
        public float dynamicResChangeFactor;
        public float lumaMipRcp;
    }

    private static readonly List<RcasConstants> RcasConfigs = new()
    {
        new(1048576000u, 872428544u),
        new(1049178080u, 877212745u),
        new(1049823372u, 882390168u),
        new(1050514979u, 887895276u),
        new(1051256227u, 893859143u),
        new(1052050675u, 900216232u),
        new(1052902144u, 907032080u),
        new(1053814727u, 914306687u),
        new(1054792807u, 922105590u),
        new(1055841087u, 930494326u),
        new(1056964608u, 939538432u),
        new(1057566688u, 944322633u),
        new(1058211980u, 949500056u),
        new(1058903587u, 955005164u),
        new(1059644835u, 960969031u),
        new(1060439283u, 967326120u),
        new(1061290752u, 974141968u),
        new(1062203335u, 981416575u),
        new(1063181415u, 989215478u),
        new(1064229695u, 997604214u),
        new(1065353216u, 1006648320),
    };
}