using System.Collections.Generic;
using System.Runtime.InteropServices;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Mathematics;
using UnityEngine;
using UnityEngine.Rendering;

namespace BadDog.Rendering.AreaLight
{
    /// <summary>
    /// Per-light area light data sent to shader.
    /// </summary>
    [StructLayout(LayoutKind.Sequential)]
    public struct BGAreaLightRenderingData
    {
        public Vector4 positionWS;              // float4 (xyz: position, w: isAreaLight flag, 1.0=area light, 0.0=not area light)
        public Vector4 right;                   // float4 (matches GPU float3 + padding)
        public Vector4 up;                      // float4 (matches GPU float3 + padding)
        public Vector4 forward;                 // float4 (matches GPU float3 + padding)
        public Vector4 size;                    // float4 (matches GPU float2 + padding)
        public Vector4 colorIntensity;          // float4 (matches GPU float3 color + float intensity)
        public Vector4 rangeParams;             // float4 (range, rangeAttenuationScale, rangeAttenuationBias, isRectLight)
        public Vector4 renderingLayerMask;      // float4 (x stores uint bitmask via bit reinterpretation)
    }

    /// <summary>
    /// Static manager for all BGAreaLight components in the scene.
    /// </summary>
    public static class BGAreaLightManager
    {
        static Dictionary<Light, BGAreaLight> s_AreaLightDictionary = new Dictionary<Light, BGAreaLight>();
        static NativeArray<BGAreaLightRenderingData> s_AreaLightRenderingDataArray;
        static GraphicsBuffer s_AreaLightDataBuffer = null;
        static GraphicsBuffer s_EmptyBuffer = null;
        static readonly Dictionary<Light, int> s_LightToAreaIndex = new Dictionary<Light, int>();

#if UNITY_EDITOR
        [UnityEditor.InitializeOnLoadMethod]
#else
        [RuntimeInitializeOnLoadMethod(RuntimeInitializeLoadType.SubsystemRegistration)]
#endif
        static void Initialize()
        {
            Cleanup();
        }

        static BGAreaLightManager()
        {
            Application.quitting -= Cleanup;
            Application.quitting += Cleanup;
        }

        public static void RegisterAreaLight(BGAreaLight areaLight)
        {
            if (areaLight == null)
            {
                return;
            }

            Light light = areaLight.GetLight();
            if (light == null)
            {
                Debug.LogWarning("BGAreaLight component does not have a Light component attached.", areaLight);
                return;
            }

            if (s_AreaLightDictionary.ContainsKey(light))
            {
                return;
            }

            s_AreaLightDictionary.Add(light, areaLight);
        }

        public static void UnRegisterAreaLight(BGAreaLight areaLight)
        {
            if (areaLight == null)
            {
                return;
            }

            Light light = areaLight.GetLight();
            if (light != null && s_AreaLightDictionary.ContainsKey(light))
            {
                s_AreaLightDictionary.Remove(light);
            }
        }

        private static bool TryGetAreaLight(Light light, out BGAreaLight areaLight)
        {
            return TryGetRegisteredAreaLight(light, out areaLight);
        }

        internal static bool TryGetRegisteredAreaLight(Light light, out BGAreaLight areaLight)
        {
            if (light == null)
            {
                areaLight = null;
                return false;
            }

            return s_AreaLightDictionary.TryGetValue(light, out areaLight);
        }

        public static NativeArray<BGAreaLightRenderingData> BuildAreaLightDataArray(NativeArray<VisibleLight> visibleLights, int mainLightIndex, int maxAreaLights = 8)
        {
            s_LightToAreaIndex.Clear();

            int visibleCount = visibleLights.Length;
            int areaLightCount = 0;

            for (int i = 0; i < visibleCount && areaLightCount < maxAreaLights; ++i)
            {
                if (i == mainLightIndex)
                {
                    continue;
                }

                Light light = visibleLights[i].light;
                if (light != null && light.type == LightType.Rectangle && TryGetAreaLight(light, out _))
                {
                    areaLightCount++;
                }
            }

            if (areaLightCount <= 0)
            {
                if (s_AreaLightRenderingDataArray.IsCreated)
                {
                    s_AreaLightRenderingDataArray.Dispose();
                }
                return default;
            }

            if (!s_AreaLightRenderingDataArray.IsCreated || s_AreaLightRenderingDataArray.Length != areaLightCount)
            {
                if (s_AreaLightRenderingDataArray.IsCreated)
                {
                    s_AreaLightRenderingDataArray.Dispose();
                }

                s_AreaLightRenderingDataArray = new NativeArray<BGAreaLightRenderingData>(
                    areaLightCount,
                    Allocator.Persistent,
                    NativeArrayOptions.UninitializedMemory);
            }

            BGAreaLightRenderingData defaultData = new BGAreaLightRenderingData
            {
                positionWS = Vector4.zero,
                right = new Vector4(1, 0, 0, 0),
                up = new Vector4(0, 1, 0, 0),
                forward = new Vector4(0, 0, 1, 0),
                size = Vector4.zero,
                colorIntensity = Vector4.zero,
                rangeParams = new Vector4(0, 1, 0, 0),
                renderingLayerMask = Vector4.zero
            };

            int areaIndex = 0;

            for (int i = 0; i < visibleCount && areaIndex < areaLightCount; ++i)
            {
                if (i == mainLightIndex)
                {
                    continue;
                }

                Light light = visibleLights[i].light;
                if (light == null || light.type != LightType.Rectangle || !TryGetAreaLight(light, out BGAreaLight areaLight))
                {
                    continue;
                }

                BGAreaLightRenderingData renderData = defaultData;

                // Mark as area light by default
                Transform lightTransform = light.transform;
                Vector3 pos = lightTransform.position;
                renderData.positionWS = new Vector4(pos.x, pos.y, pos.z, 1.0f);

                Vector3 rightVec = lightTransform.right.normalized;
                Vector3 upVec = lightTransform.up.normalized;
                Vector3 forwardVec = lightTransform.forward.normalized;

                renderData.right = new Vector4(rightVec.x, rightVec.y, rightVec.z, 0);
                renderData.up = new Vector4(upVec.x, upVec.y, upVec.z, 0);
                renderData.forward = new Vector4(forwardVec.x, forwardVec.y, forwardVec.z, 0);

                Vector2 size = areaLight.GetSize();
                float attenuationScale = areaLight.GetRangeAttenuationScale();
                float attenuationBias = areaLight.GetRangeAttenuationBias();
                float isRect = areaLight.GetIsRectLight() ? 1f : 0f;

                renderData.size = new Vector4(size.x, size.y, 0, 0);

                Color linearColor = light.color.linear;
                renderData.colorIntensity = new Vector4(linearColor.r, linearColor.g, linearColor.b, light.intensity);

                renderData.rangeParams = new Vector4(
                    light.range,
                    attenuationScale,
                    attenuationBias,
                    isRect);
                renderData.renderingLayerMask = new Vector4(math.asfloat((uint)light.renderingLayerMask), 0f, 0f, 0f);

                s_AreaLightRenderingDataArray[areaIndex] = renderData;
                s_LightToAreaIndex[light] = areaIndex;
                areaIndex++;
            }

            return s_AreaLightRenderingDataArray;
        }

        public static unsafe GraphicsBuffer UpdateGraphicsBuffer(NativeArray<BGAreaLightRenderingData> areaLightDataArray)
        {
            if (!areaLightDataArray.IsCreated || areaLightDataArray.Length == 0)
            {
                CleanupGraphicsBuffer();
                return null;
            }
            
            int count = areaLightDataArray.Length;
            int structSize = UnsafeUtility.SizeOf<BGAreaLightRenderingData>();
            
            if (s_AreaLightDataBuffer == null || !s_AreaLightDataBuffer.IsValid() || s_AreaLightDataBuffer.count != count)
            {
                CleanupGraphicsBuffer();
                s_AreaLightDataBuffer = new GraphicsBuffer(
                    GraphicsBuffer.Target.Structured,
                    count,
                    structSize);
            }
            
            s_AreaLightDataBuffer.SetData(areaLightDataArray);
            
            return s_AreaLightDataBuffer;
        }

        public static GraphicsBuffer GetEmptyBuffer()
        {
            if (s_EmptyBuffer == null || !s_EmptyBuffer.IsValid())
            {
                s_EmptyBuffer = new GraphicsBuffer(
                    GraphicsBuffer.Target.Structured,
                    1,
                    UnsafeUtility.SizeOf<BGAreaLightRenderingData>());
                
                var defaultData = new BGAreaLightRenderingData
                {
                    positionWS = Vector4.zero,  // w=0 means not an area light
                    right = new Vector4(1, 0, 0, 0),
                    up = new Vector4(0, 1, 0, 0),
                    forward = new Vector4(0, 0, 1, 0),
                    size = Vector4.zero,
                    colorIntensity = Vector4.zero, // intensity = 0 (w component)
                    rangeParams = new Vector4(0, 1, 0, 0), // range=0, scale=1, bias=0, isRect=0
                    renderingLayerMask = Vector4.zero
                };
                s_EmptyBuffer.SetData(new[] { defaultData });
            }
            
            return s_EmptyBuffer;
        }

        static void CleanupGraphicsBuffer()
        {
            if (s_AreaLightDataBuffer != null && s_AreaLightDataBuffer.IsValid())
            {
                s_AreaLightDataBuffer.Release();
                s_AreaLightDataBuffer = null;
            }
            
            if (s_EmptyBuffer != null && s_EmptyBuffer.IsValid())
            {
                s_EmptyBuffer.Release();
                s_EmptyBuffer = null;
            }
        }

        static void Cleanup()
        {
            s_AreaLightDictionary.Clear();
            
            if (s_AreaLightRenderingDataArray.IsCreated)
            {
                s_AreaLightRenderingDataArray.Dispose();
            }
            
            CleanupGraphicsBuffer();
        }

        public static void Clear()
        {
            Cleanup();
        }
    }
}