/// <summary>
/// Shiny SSRR - Screen Space Raytraced Reflections - (c) 2021 Kronnect
/// </summary>

using UnityEngine;
using UnityEngine.Rendering;

namespace ShinySSRR {

    public enum OutputMode {
        Final,
        OnlyReflections,
        SideBySideComparison
    }

    public enum RaytracingPreset {
        Fast = 10,
        Medium = 20,
        High = 30,
        Superb = 35,
        Ultra = 40
    }

    [ExecuteInEditMode]
    [RequireComponent(typeof(Camera))]
    [ImageEffectAllowedInSceneView]
    public class ShinySSRR : MonoBehaviour {

        class SSRPass {

            enum Pass {
                CopyExact = 0,
                SSRSurf = 1,
                Resolve = 2,
                BlurHoriz = 3,
                BlurVert = 4,
                Debug = 5,
                Combine = 6,
                CombineWithCompare = 7,
                GBuffPass = 8,
                Copy = 9
            }

            const string SHINY_CBUFNAME = "Shiny_SSRR";
            const float GOLDEN_RATIO = 0.618033989f;
            const int MIP_COUNT = 5;
            Material sMat;
            Texture noiseTex;
            ShinySSRR settings;
            readonly Plane[] frustumPlanes = new Plane[6];
            CommandBuffer cmd;
            int[] rtPyramid;

            public void Setup(ShinySSRR settings, bool deferred) {
                this.settings = settings;
                if (cmd == null) {
                    cmd = new CommandBuffer {
                        name = SHINY_CBUFNAME
                    };
                }
                if (sMat == null) {
                    Shader shader = Shader.Find("Hidden/Kronnect/SSR");
                    sMat = new Material(shader);
                }
                if (noiseTex == null) {
                    noiseTex = Resources.Load<Texture>("SSR/blueNoiseSSR64");
                }
                sMat.SetTexture(ShaderParams.NoiseTex, noiseTex);

                // set global settings
                sMat.SetVector(ShaderParams.SSRSettings2, new Vector4(settings.jitter, settings.contactHardening, settings.reflectionsMultiplier, settings.vignetteSize));
                sMat.SetVector(ShaderParams.SSRSettings4, new Vector4(settings.separationPos, settings.reflectionsMinIntensity, settings.reflectionsMaxIntensity, settings.specularSoftenPower));
                sMat.SetVector(ShaderParams.SSRBlurStrength, new Vector4(settings.blurStrength.x, settings.blurStrength.y, 0, 0));
                sMat.SetVector(ShaderParams.SSRSettings5, new Vector4(settings.thicknessFine * settings.thickness, settings.smoothnessThreshold, 0, 0));
                if (settings.specularControl) {
                    sMat.EnableKeyword(ShaderParams.SKW_DENOISE);
                } else {
                    sMat.DisableKeyword(ShaderParams.SKW_DENOISE);
                }
                sMat.SetFloat(ShaderParams.MinimumBlur, settings.minimumBlur);

                if (deferred) {
                    if (settings.jitter > 0) {
                        sMat.EnableKeyword(ShaderParams.SKW_JITTER);
                    } else {
                        sMat.DisableKeyword(ShaderParams.SKW_JITTER);
                    }
                    if (settings.refineThickness) {
                        sMat.EnableKeyword(ShaderParams.SKW_REFINE_THICKNESS);
                    } else {
                        sMat.DisableKeyword(ShaderParams.SKW_REFINE_THICKNESS);
                    }
                    sMat.SetVector(ShaderParams.SSRSettings, new Vector4(settings.thickness, settings.sampleCount, settings.binarySearchIterations, settings.maxRayLength));
                    sMat.SetVector(ShaderParams.MaterialData, new Vector4(0, settings.fresnel, settings.fuzzyness, settings.decay));
                }

                if (rtPyramid == null || rtPyramid.Length != MIP_COUNT) {
                    rtPyramid = new int[MIP_COUNT];
                }
                for (int k = 0; k < rtPyramid.Length; k++) {
                    rtPyramid[k] = Shader.PropertyToID("_BlurRTMip" + k);
                }
            }

            public void Execute(RenderTexture source, RenderTexture destination, Camera cam) {
                ExecuteInternal(source, destination, cam);
                Graphics.Blit(source, destination, sMat, (int)Pass.CopyExact);
            }

            void ExecuteInternal(RenderTexture source, RenderTexture renderTexture, Camera cam) {
                // ignore SceneView depending on setting
                if (cam.cameraType == CameraType.SceneView) {
                    if (!settings.showInSceneView) return;
                } else {
                    // ignore any camera other than GameView
                    if (cam.cameraType != CameraType.Game) return;
                }

                RenderTextureDescriptor sourceDesc = source.descriptor;
                sourceDesc.colorFormat = settings.lowPrecision ? RenderTextureFormat.ARGB32 : RenderTextureFormat.ARGBHalf;
                sourceDesc.width /= settings.downsampling;
                sourceDesc.height /= settings.downsampling;
                sourceDesc.msaaSamples = 1;

                float goldenFactor = GOLDEN_RATIO;
                if (settings.animatedJitter) {
                    goldenFactor *= (Time.frameCount % 480);
                }
                Shader.SetGlobalVector(ShaderParams.SSRSettings3, new Vector4(sourceDesc.width, sourceDesc.height, goldenFactor, settings.depthBias));

                if (settings.isDeferredActive) {
                    // init command buffer
                    cmd.Clear();

                    // pass UNITY_MATRIX_V
                    sMat.SetMatrix(ShaderParams.WorldToViewDir, cam.worldToCameraMatrix);

                    // prepare ssr target
                    cmd.GetTemporaryRT(ShaderParams.RayCast, sourceDesc, FilterMode.Point);

                    // raytrace using gbuffers
                    cmd.Blit(source, ShaderParams.RayCast, sMat, (int)Pass.GBuffPass);

                } else {

                    // early exit if no reflection objects
                    int count = Reflections.instances.Count;
                    if (count == 0) return;

                    bool firstSSR = true;

                    GeometryUtility.CalculateFrustumPlanes(cam, frustumPlanes);

                    for (int k = 0; k < count; k++) {
                        Reflections go = Reflections.instances[k];
                        if (go == null) continue;
                        int rendererCount = go.ssrRenderers.Count;
                        for (int j = 0; j < rendererCount; j++) {
                            Reflections.SSR_Renderer ssrRenderer = go.ssrRenderers[j];
                            Renderer goRenderer = ssrRenderer.renderer;

                            if (goRenderer == null || !goRenderer.isVisible) continue;

                            // if object is part of static batch, check collider bounds (if existing)
                            if (goRenderer.isPartOfStaticBatch) {
                                if (ssrRenderer.hasStaticBounds) {
                                    // check artifically computed bounds
                                    if (!GeometryUtility.TestPlanesAABB(frustumPlanes, ssrRenderer.staticBounds)) continue;
                                } else if (ssrRenderer.collider != null) {
                                    // check if object is visible by current camera using collider bounds
                                    if (!GeometryUtility.TestPlanesAABB(frustumPlanes, ssrRenderer.collider.bounds)) continue;
                                }
                            } else {
                                // check if object is visible by current camera using renderer bounds
                                if (!GeometryUtility.TestPlanesAABB(frustumPlanes, goRenderer.bounds)) continue;
                            }

                            if (!ssrRenderer.isInitialized) {
                                ssrRenderer.Init(sMat);
                                ssrRenderer.UpdateMaterialProperties(go, settings);
                            }
#if UNITY_EDITOR
                        else if (!Application.isPlaying) {
                                ssrRenderer.UpdateMaterialProperties(go, settings);
                            }
#endif
                            if (ssrRenderer.exclude) continue;

                            if (firstSSR) {
                                firstSSR = false;

                                // init command buffer
                                cmd.Clear();

                                // prepare ssr target
                                cmd.GetTemporaryRT(ShaderParams.RayCast, sourceDesc, FilterMode.Point);
                                cmd.SetRenderTarget(ShaderParams.RayCast);
                                cmd.ClearRenderTarget(true, true, new Color(0, 0, 0, 0));
                            }
                            for (int s = 0; s < ssrRenderer.ssrMaterials.Length; s++) {
                                if (go.subMeshMask <= 0 || ((1 << s) & go.subMeshMask) != 0) {
                                    Material ssrMat = ssrRenderer.ssrMaterials[s];
                                    cmd.DrawRenderer(goRenderer, ssrMat, s, (int)Pass.SSRSurf);
                                }
                            }
                        }
                    }

                    if (firstSSR) return;
                }

                // Resolve reflections
                RenderTextureDescriptor copyDesc = sourceDesc;
                copyDesc.depthBufferBits = 0;

                cmd.GetTemporaryRT(ShaderParams.ReflectionsTex, copyDesc);
                cmd.Blit(source, ShaderParams.ReflectionsTex, sMat, (int)Pass.Resolve);
                RenderTargetIdentifier input = ShaderParams.ReflectionsTex;

                // Pyramid blur
                copyDesc.width /= settings.blurDownsampling;
                copyDesc.height /= settings.blurDownsampling;
                for (int k = 0; k < MIP_COUNT; k++) {
                    copyDesc.width = Mathf.Max(2, copyDesc.width / 2);
                    copyDesc.height = Mathf.Max(2, copyDesc.height / 2);
                    cmd.GetTemporaryRT(rtPyramid[k], copyDesc, FilterMode.Bilinear);
                    cmd.GetTemporaryRT(ShaderParams.BlurRT, copyDesc, FilterMode.Bilinear);
                    cmd.Blit(input, ShaderParams.BlurRT, sMat, (int)Pass.BlurHoriz);
                    cmd.Blit(ShaderParams.BlurRT, rtPyramid[k], sMat, (int)Pass.BlurVert);
                    cmd.ReleaseTemporaryRT(ShaderParams.BlurRT);
                    input = rtPyramid[k];
                }

                // Output
                int finalPass;
                if (settings.outputMode == OutputMode.Final) {
                    finalPass = (int)Pass.Combine;
                } else if (settings.outputMode == OutputMode.SideBySideComparison) {
                    finalPass = (int)Pass.CombineWithCompare;
                } else {
                    finalPass = (int)Pass.Debug;
                }
                cmd.Blit(ShaderParams.ReflectionsTex, source, sMat, finalPass);

                if (settings.stopNaN) {
                    RenderTextureDescriptor nanDesc = source.descriptor;
                    nanDesc.depthBufferBits = 0;
                    nanDesc.msaaSamples = 1;
                    cmd.GetTemporaryRT(ShaderParams.NaNBuffer, nanDesc);
                    cmd.Blit(source, ShaderParams.NaNBuffer, sMat, (int)Pass.CopyExact);
                    cmd.Blit(ShaderParams.NaNBuffer, source, sMat, (int)Pass.CopyExact);
                }

                // Clean up
                for (int k = 0; k < rtPyramid.Length; k++) {
                    cmd.ReleaseTemporaryRT(rtPyramid[k]);
                }
                cmd.ReleaseTemporaryRT(ShaderParams.RayCast);

                Graphics.ExecuteCommandBuffer(cmd);

            }

            public void Cleanup() {
                if (sMat != null) {
                    DestroyImmediate(sMat);
                }
            }
        }

        [Header("General Settings")]

        [Tooltip("Show reflections in SceneView window")]
        public bool showInSceneView = true;

        [Tooltip("Downsampling multiplier applied to the final blurred reflections")]
        [Range(1, 8)] public int downsampling = 1;

        [Tooltip("Bias applied to depth checking. Increase if reflections desappear at the distance when downsampling is used")]
        [Min(0)] public float depthBias = 0.01f;

        [Tooltip("Show final result / debug view or compare view")]
        public OutputMode outputMode = OutputMode.Final;

        [Tooltip("Position of the dividing line")]
        [Range(-0.01f, 1.01f)] public float separationPos = 0.5f;

        [Tooltip("HDR reflections")]
        public bool lowPrecision;

        [Tooltip("Prevents out of range colors when composing reflections in the destination buffer. This operation performs a ping-pong copy of the frame buffer which can be expensive. Use only if required.")]
        public bool stopNaN;

        [Tooltip("Max number of samples used during the raymarch loop")]
        [Range(4, 128)] public int sampleCount = 16;

        [HideInInspector]
        public float stepSize; // no longer used; kept for backward compatibility during upgrade

        [Tooltip("Maximum reflection distance")]
        public float maxRayLength;

        [Tooltip("Assumed thickness of geometry in the depth buffer before binary search")]
        public float thickness = 0.2f;

        [Tooltip("Number of refinements steps when a reflection hit is found")]
        [Range(0, 16)] public int binarySearchIterations = 6;

        [Tooltip("Increase accuracy of reflection hit after binary search by discarding points further than a reduced thickness.")]
        public bool refineThickness;

        [Tooltip("Assumed thickness of geometry in the depth buffer after binary search")]
        [Range(0.005f, 1f)]
        public float thicknessFine = 0.05f;

        [Tooltip("Jitter helps smoothing edges")]
        [Range(0, 1f)] public float jitter = 0.3f;

        [Tooltip("Animates jitter every frame")]
        public bool animatedJitter = true;

        [Header("Reflection Intensity")]

        [Tooltip("Minimum smoothness to receive reflections")]
        [Range(0,1)]
        public float smoothnessThreshold;

        [Tooltip("Reflection multiplier")]
        [Range(0, 2)]
        public float reflectionsMultiplier = 1f;

        [Tooltip("Reflection min intensity")]
        [Range(0, 1)]
        public float reflectionsMinIntensity;

        [Tooltip("Reflection max intensity")]
        [Range(0, 1)]
        public float reflectionsMaxIntensity = 1f;

        [Range(0, 1)]
        [Tooltip("Reduces reflection based on view angle")]
        public float fresnel = 0.75f;

        [Tooltip("Reflection decay with distance to reflective point")]
        public float decay = 2f;

        [Tooltip("Reduces intensity of specular reflections")]
        public bool specularControl;

        [Min(0), Tooltip("Power of the specular filter")]
        public float specularSoftenPower = 15f;

        [Tooltip("Controls the attenuation range of effect on screen borders")]
        [Range(0.5f, 2f)]
        public float vignetteSize = 1.1f;

        [Header("Reflection Sharpness")]

        [Min(0)]
        [Tooltip("Ray dispersion with distance")]
        public float fuzzyness;

        [Tooltip("Makes sharpen reflections near objects")]
        public float contactHardening;

        [Range(0, 4f)]
        [Tooltip("Produces sharper reflections based on distance")]
        public float minimumBlur = 0.25f;

        [Tooltip("Downsampling multiplier applied to the blur")]
        [Range(1, 8)] public int blurDownsampling = 1;

        [Tooltip("Custom directional blur strength")]
        public Vector2 blurStrength = Vector2.one;

        SSRPass renderPass;

        Camera cam;

        public bool isDeferredActive {
            get {
                if (cam != null) {
                    return cam.renderingPath != RenderingPath.Forward;
                }
                return false;
            }
        }

        void OnDisable() {
            if (renderPass != null) {
                renderPass.Cleanup();
            }
        }

        void OnEnable() {
            if (maxRayLength == 0) {
                maxRayLength = Mathf.Max(0.1f, stepSize * sampleCount);
            }
            if (renderPass == null) {
                renderPass = new SSRPass();
            }
            cam = GetComponent<Camera>();
            cam.depthTextureMode |= DepthTextureMode.Depth;
        }

        void OnValidate() {
            decay = Mathf.Max(1f, decay);
            if (maxRayLength == 0) {
                maxRayLength = stepSize * sampleCount;
            }
            maxRayLength = Mathf.Max(0.1f, maxRayLength);
            fuzzyness = Mathf.Max(0, fuzzyness);
            thickness = Mathf.Max(0.01f, thickness);
            thicknessFine = Mathf.Max(0.01f, thicknessFine);
            contactHardening = Mathf.Max(0, contactHardening);
            reflectionsMaxIntensity = Mathf.Max(reflectionsMinIntensity, reflectionsMaxIntensity);
            blurStrength.x = Mathf.Max(blurStrength.x, 0f);
            blurStrength.y = Mathf.Max(blurStrength.y, 0f);
        }

        void OnRenderImage(RenderTexture source, RenderTexture destination) {
            renderPass.Setup(this, isDeferredActive);
            renderPass.Execute(source, destination, cam);
        }

        public void ApplyRaytracingPreset(RaytracingPreset preset) {
            switch (preset) {
                case RaytracingPreset.Fast:
                    sampleCount = 16;
                    maxRayLength = 6;
                    binarySearchIterations = 4;
                    downsampling = 3;
                    thickness = 0.5f;
                    refineThickness = false;
                    jitter = 0.3f;
                    break;
                case RaytracingPreset.Medium:
                    sampleCount = 24;
                    maxRayLength = 12;
                    binarySearchIterations = 5;
                    downsampling = 2;
                    refineThickness = false;
                    break;
                case RaytracingPreset.High:
                    sampleCount = 48;
                    maxRayLength = 24;
                    binarySearchIterations = 6;
                    downsampling = 1;
                    refineThickness = false;
                    thicknessFine = 0.05f;
                    break;
                case RaytracingPreset.Superb:
                    sampleCount = 88;
                    maxRayLength = 48;
                    binarySearchIterations = 7;
                    downsampling = 1;
                    refineThickness = true;
                    thicknessFine = 0.02f;
                    break;
                case RaytracingPreset.Ultra:
                    sampleCount = 128;
                    maxRayLength = 64;
                    binarySearchIterations = 8;
                    downsampling = 1;
                    refineThickness = true;
                    thicknessFine = 0.02f;
                    break;
            }
        }

    }

}