using System;
using System.Collections.Generic;
using Unity.Collections;
using UnityEngine;
using UnityEngine.Rendering;
using UnityEngine.Rendering.Universal;

namespace BadDog.Rendering.AreaLight
{
    public partial class AreaLightingPass : ScriptableRenderPass
    {
        private static readonly ProfilingSampler m_ProfilingSampler = new ProfilingSampler("AreaLighting Setup");
        private static readonly ProfilingSampler m_ShadowProfilingSampler = new ProfilingSampler("AreaLight Shadows");

        private const int k_ShadowmapBufferBits = 16;
        private const int k_MaxAreaLightShadowCount = 8;
        
        private static readonly Matrix4x4[] s_EmptyShadowMatrices = new Matrix4x4[k_MaxAreaLightShadowCount]
        {
            Matrix4x4.identity, Matrix4x4.identity, Matrix4x4.identity, Matrix4x4.identity,
            Matrix4x4.identity, Matrix4x4.identity, Matrix4x4.identity, Matrix4x4.identity
        };
        
        private static readonly Vector4[] s_EmptyShadowParams = new Vector4[k_MaxAreaLightShadowCount]
        {
            Vector4.zero, Vector4.zero, Vector4.zero, Vector4.zero,
            Vector4.zero, Vector4.zero, Vector4.zero, Vector4.zero
        };
        private AreaLightSettings m_LightSettings;
        private AreaLightShadowSettings m_ShadowSettings;
        
        private RTHandle m_AreaLightsShadowmapHandle;
        private Matrix4x4[] m_AreaLightShadowSliceIndexTo_WorldShadowMatrix;  // Indexed by areaLightIndex
        private Vector4[] m_AreaLightIndexToShadowParams;                     // Indexed by areaLightIndex
        private ShadowSliceData[] m_AreaLightsShadowSlices;                   // Indexed by shadowSliceIndex
        private int m_ValidShadowCastingLightsCount = 0;
        private short[] m_VisibleLightIndexToAreaLightIndex;
        private short[] m_AreaLightIndexToVisibleLightIndex;
        private bool[] m_VisibleLightIndexToIsCastingShadows;
        private readonly List<short> m_ShadowSliceToAreaLightIndex = new();   // shadowSliceIndex -> areaLightIndex
        private RendererList[] m_ShadowRendererLists;                         // Indexed by shadowSliceIndex
        private bool m_CreateEmptyShadowmap;
        private bool m_EmptyShadowmapNeedsClear = true;
        private int m_RenderTargetWidth;
        private int m_RenderTargetHeight;
        private ShadowData m_ShadowDataCache;
        private RenderTextureDescriptor m_AreaLightShadowDescriptor;
        private static readonly List<int> s_AreaLightRequestedResolutions = new();

        private partial class PassData
        {
            internal AreaLightingPass pass;
            internal bool emptyShadowmap;
            internal Vector2Int allocatedShadowAtlasSize;
            internal RendererList[] shadowRendererLists;
            internal NativeArray<VisibleLight> visibleLights;
            internal Matrix4x4 viewMatrix;
            internal Matrix4x4 projectionMatrix;
            internal Camera camera;
        }

        private PassData m_PassData;

        partial void ConfigureShadowSettingsRenderGraph(int maxShadowCastingLights);

        public AreaLightingPass()
        {
            renderPassEvent = RenderPassEvent.AfterRenderingShadows;

            m_AreaLightShadowSliceIndexTo_WorldShadowMatrix = new Matrix4x4[k_MaxAreaLightShadowCount];
            m_AreaLightIndexToShadowParams = new Vector4[k_MaxAreaLightShadowCount];
            m_AreaLightsShadowSlices = new ShadowSliceData[k_MaxAreaLightShadowCount];
            m_ShadowRendererLists = new RendererList[k_MaxAreaLightShadowCount];

            m_PassData = new PassData
            {
                shadowRendererLists = m_ShadowRendererLists
            };
        }
        
        private void ConfigureAreaLightingSettings(AreaLightSettings lightSettings,
            AreaLightShadowSettings shadowSettings)
        {
            m_LightSettings = lightSettings;
            m_ShadowSettings = shadowSettings;
            
            PreIntegratedFGD.instance.SetShader(PreIntegratedFGD.FGDIndex.FGD_GGXAndDisneyDiffuse, lightSettings.preIntegratedFGD_GGXDisneyDiffuse);
            PreIntegratedFGD.instance.SetShader(PreIntegratedFGD.FGDIndex.FGD_CharlieAndFabricLambert, lightSettings.preIntegratedFGD_CharlieFabricLambert);
            PreIntegratedFGD.instance.SetShader(PreIntegratedFGD.FGDIndex.FGD_Marschner, lightSettings.preIntegratedFGD_Marschner);

            ConfigureShadowSettingsRenderGraph(shadowSettings.maxShadowCastingLights);
        }

        public void Setup(AreaLightSettings lightSettings, AreaLightShadowSettings shadowSettings, RenderingData renderingData)
        {
            ConfigureAreaLightingSettings(lightSettings, shadowSettings);
            PrepareAreaLightShadows(renderingData);
        }

        public override void OnCameraCleanup(CommandBuffer cmd)
        {
            CoreUtils.SetKeyword(cmd, AreaLightingShaderIDs.ENABLE_BG_AREA_LIGHTING_KEYWORD, false);
        }

        public void Dispose()
        {
            m_AreaLightsShadowmapHandle?.Release();
        }
        

        private void PrepareAreaLightShadows(RenderingData renderingData)
        {
            m_CreateEmptyShadowmap = false;
            m_ValidShadowCastingLightsCount = 0;
            m_RenderTargetWidth = m_ShadowSettings.shadowAtlasResolution;
            m_RenderTargetHeight = m_ShadowSettings.shadowAtlasResolution;

            var cameraData = renderingData.cameraData;
            var lightData = renderingData.lightData;
            var shadowData = renderingData.shadowData;
            m_ShadowDataCache = shadowData;
            NativeArray<VisibleLight> visibleLights = lightData.visibleLights;
            if (!visibleLights.IsCreated || visibleLights.Length == 0)
            {
                m_CreateEmptyShadowmap = true;
                m_EmptyShadowmapNeedsClear = true;
                AllocateShadowmapIfNeeded();
                return;
            }

            EnsureMappingArrays(visibleLights.Length);
            ResetMappingArrays(visibleLights.Length);

            PrepareAreaLightRequestedResolutions(visibleLights, lightData.mainLightIndex);

            var atlasLayout = new AreaLightsShadowAtlasLayout(lightData, shadowData, cameraData, s_AreaLightRequestedResolutions, m_ShadowSettings.maxShadowCastingLights, m_ShadowSettings.shadowAtlasResolution);

            int targetShadowCount = Math.Min(atlasLayout.GetTotalShadowSlicesCount(), m_ShadowSettings.maxShadowCastingLights);
            if (targetShadowCount <= 0)
            {
                m_CreateEmptyShadowmap = true;
                m_EmptyShadowmapNeedsClear = true;
                AllocateShadowmapIfNeeded();
                return;
            }

            // Trim shadow atlas dimensions if possible (to avoid allocating texture space that will not be used)
            int atlasMaxX = 0;
            int atlasMaxY = 0;
            for (int i = 0; i < targetShadowCount; ++i)
            {
                var request = atlasLayout.GetSortedShadowResolutionRequest(i);
                int maxX = request.offsetX + request.allocatedResolution;
                int maxY = request.offsetY + request.allocatedResolution;
                if (maxX > atlasMaxX) atlasMaxX = maxX;
                if (maxY > atlasMaxY) atlasMaxY = maxY;
            }
            
            // Use power-of-two dimensions (might perform better on some hardware)
            m_RenderTargetWidth = Mathf.NextPowerOfTwo(atlasMaxX);
            m_RenderTargetHeight = Mathf.NextPowerOfTwo(atlasMaxY);

            int areaLightIndex = 0;
            int shadowSliceIndex = 0;

            for (int visibleLightIndex = 0; visibleLightIndex < visibleLights.Length && areaLightIndex < m_LightSettings.maxAreaLights; ++visibleLightIndex)
            {
                if (visibleLightIndex == lightData.mainLightIndex)
                {
                    continue;
                }

                ref readonly VisibleLight visibleLight = ref visibleLights.UnsafeElementAt(visibleLightIndex);
                Light light = visibleLight.light;

                // Only registered BGAreaLight rectangle lights participate in the area light system.
                if (visibleLight.lightType != LightType.Rectangle || light == null
                    || !BGAreaLightManager.TryGetRegisteredAreaLight(light, out _))
                {
                    continue;
                }

                // Default shadow params: no shadow (w = -1 means no shadow)
                // z = 0 means Spot-like light type identifier (for shader compatibility)
                m_AreaLightIndexToShadowParams[areaLightIndex] = new Vector4(0f, 0f, 0f, -1f);
                m_AreaLightShadowSliceIndexTo_WorldShadowMatrix[areaLightIndex] = Matrix4x4.identity;
                m_AreaLightIndexToVisibleLightIndex[areaLightIndex] = (short)visibleLightIndex;
                m_VisibleLightIndexToAreaLightIndex[visibleLightIndex] = (short)areaLightIndex;

                // Check if this light casts shadows (light type and null check already done above)
                bool canCastShadows = AreaShadowUtils.CanLightCastShadows(light);
                // Ensure we don't exceed array bounds: shadowSliceIndex < maxShadowCastingLights
                if (m_ShadowSettings.shadowFilter != AreaLightShadowSettings.BGAreaLightShadowMode.Off
                    && canCastShadows
                    && atlasLayout.HasSpaceForLight(visibleLightIndex)
                    && shadowSliceIndex < targetShadowCount
                    && shadowSliceIndex < m_ShadowSettings.maxShadowCastingLights)
                {
                    var resolutionRequest = atlasLayout.GetSliceShadowResolutionRequest(visibleLightIndex, 0);
                    int resolution = Mathf.Clamp(resolutionRequest.allocatedResolution, AreaShadowUtils.MinimalAreaLightShadowResolution(), m_ShadowSettings.shadowAtlasResolution);

                    if (TryComputeAreaLightShadowSlice(visibleLight, visibleLightIndex, resolution, out var viewMatrix, out var projMatrix, out var splitData, out var shadowTransform))
                    {
                        // Store shadow slice data using shadowSliceIndex for rendering
                        ref ShadowSliceData slice = ref m_AreaLightsShadowSlices[shadowSliceIndex];
                        slice.viewMatrix = viewMatrix;
                        slice.projectionMatrix = projMatrix;
                        slice.splitData = splitData;
                        slice.offsetX = resolutionRequest.offsetX;
                        slice.offsetY = resolutionRequest.offsetY;
                        slice.resolution = resolution;
                        slice.shadowTransform = shadowTransform;

                        float shadowStrength = light.shadowStrength;
                        
                        // Store shadow parameters using areaLightIndex (matches BGAreaLightManager index)
                        // x: shadowStrength, y: softShadows (1.0), z: lightType (0.0 = Spot-like), w: areaLightIndex
                        m_AreaLightIndexToShadowParams[areaLightIndex] = new Vector4(shadowStrength, 1f, 0f, (float)areaLightIndex);
                        m_AreaLightShadowSliceIndexTo_WorldShadowMatrix[areaLightIndex] = shadowTransform;
                        
                        // Map shadowSliceIndex to areaLightIndex for rendering loop
                        m_ShadowSliceToAreaLightIndex.Add((short)areaLightIndex);
                        m_VisibleLightIndexToIsCastingShadows[visibleLightIndex] = true;
                        
                        shadowSliceIndex++;
                    }
                }

                areaLightIndex++;
            }

            m_ValidShadowCastingLightsCount = shadowSliceIndex;

            // Clear remaining entries
            for (int i = areaLightIndex; i < m_AreaLightShadowSliceIndexTo_WorldShadowMatrix.Length; ++i)
            {
                m_AreaLightShadowSliceIndexTo_WorldShadowMatrix[i] = Matrix4x4.identity;
                m_AreaLightIndexToShadowParams[i] = Vector4.zero;
            }

            if (m_ValidShadowCastingLightsCount <= 0)
            {
                m_CreateEmptyShadowmap = true;
                m_ShadowSliceToAreaLightIndex.Clear();
                m_EmptyShadowmapNeedsClear = true;
                AllocateShadowmapIfNeeded();
                return;
            }

            m_CreateEmptyShadowmap = false;
            m_EmptyShadowmapNeedsClear = true;

            float oneOverAtlasWidth = 1f / m_RenderTargetWidth;
            float oneOverAtlasHeight = 1f / m_RenderTargetHeight;

            // Apply atlas transform to shadow matrices
            for (int i = 0; i < shadowSliceIndex; ++i)
            {
                int areaLightIdx = m_ShadowSliceToAreaLightIndex[i];
                ref ShadowSliceData slice = ref m_AreaLightsShadowSlices[i];

                Matrix4x4 sliceTransform = Matrix4x4.identity;
                sliceTransform.m00 = slice.resolution * oneOverAtlasWidth;
                sliceTransform.m11 = slice.resolution * oneOverAtlasHeight;
                sliceTransform.m03 = slice.offsetX * oneOverAtlasWidth;
                sliceTransform.m13 = slice.offsetY * oneOverAtlasHeight;

                Matrix4x4 worldToShadow = sliceTransform * slice.shadowTransform;
                m_AreaLightShadowSliceIndexTo_WorldShadowMatrix[areaLightIdx] = worldToShadow;
                slice.shadowTransform = worldToShadow;
            }

            UpdateTextureDescriptorIfNeeded();
            AllocateShadowmapIfNeeded();
        }

        partial void AllocateShadowmapIfNeeded();

        private void UpdateTextureDescriptorIfNeeded()
        {
            if (   m_AreaLightShadowDescriptor.width != m_RenderTargetWidth
                || m_AreaLightShadowDescriptor.height != m_RenderTargetHeight
                || m_AreaLightShadowDescriptor.depthBufferBits != k_ShadowmapBufferBits
                || m_AreaLightShadowDescriptor.colorFormat != RenderTextureFormat.Shadowmap)
            {
                m_AreaLightShadowDescriptor = new RenderTextureDescriptor(m_RenderTargetWidth, m_RenderTargetHeight, RenderTextureFormat.Shadowmap, k_ShadowmapBufferBits);
            }
        }

        void PrepareAreaLightRequestedResolutions(NativeArray<VisibleLight> visibleLights, int mainLightIndex)
        {
            int requiredLength = visibleLights.Length;
            if (s_AreaLightRequestedResolutions.Capacity < requiredLength)
            {
                s_AreaLightRequestedResolutions.Capacity = requiredLength;
            }

            s_AreaLightRequestedResolutions.Clear();

            int defaultResolution = Mathf.Max(m_ShadowSettings.shadowAtlasResolution, AreaShadowUtils.MinimalAreaLightShadowResolution());

            for (int i = 0; i < visibleLights.Length; ++i)
            {
                ref readonly VisibleLight visibleLight = ref visibleLights.UnsafeElementAt(i);
                bool isValidAreaLight = AreaShadowUtils.IsValidShadowCastingLight(visibleLight, mainLightIndex, i);

                int requestedResolution = 0;

                if (isValidAreaLight)
                {
                    requestedResolution = defaultResolution;

                    Light unityLight = visibleLight.light;
                    if (unityLight != null && BGAreaLightManager.TryGetRegisteredAreaLight(unityLight, out var areaLight))
                    {
                        int custom = areaLight.GetCustomShadowResolution();
                        if (custom > 0)
                        {
                            int minRes = AreaShadowUtils.MinimalAreaLightShadowResolution();
                            requestedResolution = Mathf.Max(minRes, Mathf.NextPowerOfTwo(custom));
                        }
                    }
                }

                s_AreaLightRequestedResolutions.Add(requestedResolution);
            }
        }

        void EnsureMappingArrays(int visibleLightCount)
        {
            if (m_AreaLightIndexToVisibleLightIndex == null || m_AreaLightIndexToVisibleLightIndex.Length != m_LightSettings.maxAreaLights)
            {
                m_AreaLightIndexToVisibleLightIndex = new short[m_LightSettings.maxAreaLights];
            }

            if (m_VisibleLightIndexToAreaLightIndex == null || m_VisibleLightIndexToAreaLightIndex.Length != visibleLightCount)
            {
                m_VisibleLightIndexToAreaLightIndex = new short[visibleLightCount];
            }

            if (m_VisibleLightIndexToIsCastingShadows == null || m_VisibleLightIndexToIsCastingShadows.Length != visibleLightCount)
            {
                m_VisibleLightIndexToIsCastingShadows = new bool[visibleLightCount];
            }

            if (m_ShadowRendererLists == null || m_ShadowRendererLists.Length != m_ShadowSettings.maxShadowCastingLights)
            {
                m_ShadowRendererLists = new RendererList[m_ShadowSettings.maxShadowCastingLights];
            }
            
            EnsureMappingArraysRenderGraph();
        }
        
        partial void EnsureMappingArraysRenderGraph();

        void ResetMappingArrays(int visibleLightCount)
        {
            if (m_VisibleLightIndexToAreaLightIndex != null)
            {
                for (int i = 0; i < m_VisibleLightIndexToAreaLightIndex.Length; ++i)
                {
                    m_VisibleLightIndexToAreaLightIndex[i] = -1;
                }
            }

            if (m_VisibleLightIndexToIsCastingShadows != null)
            {
                Array.Clear(m_VisibleLightIndexToIsCastingShadows, 0, m_VisibleLightIndexToIsCastingShadows.Length);
            }

            if (m_AreaLightIndexToVisibleLightIndex != null)
            {
                for (int i = 0; i < m_AreaLightIndexToVisibleLightIndex.Length; ++i)
                {
                    m_AreaLightIndexToVisibleLightIndex[i] = -1;
                }
            }

            if (m_ShadowRendererLists != null)
            {
                Array.Clear(m_ShadowRendererLists, 0, m_ShadowRendererLists.Length);
            }

            ResetMappingArraysRenderGraph();

            m_ShadowSliceToAreaLightIndex.Clear();
        }
        
        partial void ResetMappingArraysRenderGraph();

        bool TryComputeAreaLightShadowSlice(in VisibleLight visibleLight, int visibleLightIndex, int resolution,
            out Matrix4x4 viewMatrix, out Matrix4x4 projMatrix, out ShadowSplitData splitData, out Matrix4x4 shadowTransform)
        {
            viewMatrix = Matrix4x4.identity;
            projMatrix = Matrix4x4.identity;
            splitData = new ShadowSplitData();
            shadowTransform = Matrix4x4.identity;

            Light unityLight = visibleLight.light;
            
            if (unityLight == null)
            {
                return false;
            }

            if (!AreaShadowUtils.TryGetAreaLightParameters(unityLight, out float areaLightShadowCone,
                    out Vector2 areaSize))
            {
                return false;
            }

            float forwardOffset = AreaShadowUtils.GetAreaLightOffsetForShadows(areaSize, areaLightShadowCone);

            float normalBias;
            float nearPlane;
            if (BGAreaLightManager.TryGetRegisteredAreaLight(unityLight, out var areaLightComp))
            {
                normalBias = areaLightComp.GetShadowNormalBias();
                nearPlane = areaLightComp.GetShadowNearPlane();
            }
            else
            {
                normalBias = unityLight.shadowNormalBias;
                nearPlane = unityLight.shadowNearPlane;
            }

            AreaShadowUtils.ExtractRectangleAreaLightMatrix(
                visibleLight,
                forwardOffset,
                areaLightShadowCone,
                nearPlane,
                areaSize,
                resolution,
                normalBias,
                SystemInfo.usesReversedZBuffer,
                out shadowTransform,
                out viewMatrix,
                out projMatrix,
                out splitData);

            return true;
        }
    }
}