using System.Collections.Generic;
using UnityEngine;
using UnityEngine.Playables;

namespace MaterialTrack.Group
{
    public class MaterialGroupMixer : PlayableBehaviour, IMixer
    {
        private MaterialGroup materialGroup;
        
        // Dictionary to store original material states for restoration
        private Dictionary<int, Material> originalMaterials = new Dictionary<int, Material>();
        private bool firstFrameHappened = false;

        // IMixer implementation
        private readonly RenderTextureCache renderTextureCache = new RenderTextureCache();
        private readonly Texture2DCache texture2DCache = new Texture2DCache();

        public RenderTextureCache RenderTextureCache => renderTextureCache;
        public Texture2DCache Texture2DCache => texture2DCache;

        public IEnumerable<Material> Materials
        {
            get
            {
                if (materialGroup != null && materialGroup.materials != null)
                {
                    return materialGroup.materials;
                }
                return new Material[0];
            }
        }

        public override void OnPlayableDestroy(Playable playable)
        {
            RestoreMaterials();
            firstFrameHappened = false;
            // renderTextureCache.Dispose(); // Not available
            // texture2DCache.Dispose(); // Not available
        }

        private void RestoreMaterials()
        {
            if (materialGroup == null || materialGroup.materials == null) return;

            foreach (var mat in materialGroup.materials)
            {
                if (mat == null) continue;
                
                if (originalMaterials.TryGetValue(mat.GetInstanceID(), out Material original))
                {
                    mat.CopyPropertiesFromMaterial(original);
                }
            }
            
            // Cleanup copies
            foreach(var copy in originalMaterials.Values)
            {
                if (copy != null) Object.DestroyImmediate(copy);
            }
            originalMaterials.Clear();
        }

        public override void ProcessFrame(Playable playable, FrameData info, object playerData)
        {
            materialGroup = playerData as MaterialGroup;

            if (materialGroup == null || materialGroup.materials == null)
                return;

            int inputCount = playable.GetInputCount();
            if (inputCount == 0)
                return;

            // Initialize / Backup on first frame
            if (!firstFrameHappened)
            {
                foreach (var mat in materialGroup.materials)
                {
                    if (mat == null) continue;
                    if (!originalMaterials.ContainsKey(mat.GetInstanceID()))
                    {
                        originalMaterials[mat.GetInstanceID()] = new Material(mat);
                    }
                }
                firstFrameHappened = true;
            }

            // Calculate the mixed behaviour
            RendererBehaviour mixedBehaviour = null;
            float totalWeight = 0f;

            for (int i = 0; i < inputCount; i++)
            {
                float weight = playable.GetInputWeight(i);
                if (weight <= 0) continue;

                var scriptPlayable = (ScriptPlayable<RendererBehaviour>)playable.GetInput(i);
                var behaviour = scriptPlayable.GetBehaviour();
                
                weight *= behaviour.weightMultiplier;
                totalWeight += weight;

                if (mixedBehaviour == null)
                {
                    mixedBehaviour = new RendererBehaviour();
                    mixedBehaviour.propertyName = behaviour.propertyName;
                    mixedBehaviour.propertyType = behaviour.propertyType;
                    mixedBehaviour.textureTarget = behaviour.textureTarget;
                    mixedBehaviour.vector = behaviour.vector;
                    mixedBehaviour.texture = behaviour.texture;
                }
                else
                {
                    mixedBehaviour.Lerp(mixedBehaviour, behaviour, weight / totalWeight);
                }
            }

            if (mixedBehaviour != null)
            {
                // Apply to all materials
                foreach (var mat in materialGroup.materials)
                {
                    if (mat == null) continue;
                    ApplyBehaviourToMaterial(mixedBehaviour, mat);
                }
            }
        }

        private void ApplyBehaviourToMaterial(RendererBehaviour behaviour, Material target)
        {
            switch (behaviour.propertyType)
            {
                case UnityEngine.Rendering.ShaderPropertyType.Color:
                    target.SetColor(behaviour.propertyName, behaviour.vector);
                    break;
                case UnityEngine.Rendering.ShaderPropertyType.Vector:
                    target.SetVector(behaviour.propertyName, behaviour.vector);
                    break;
                case UnityEngine.Rendering.ShaderPropertyType.Float:
                case UnityEngine.Rendering.ShaderPropertyType.Range:
                    target.SetFloat(behaviour.propertyName, behaviour.vector.x);
                    break;
                case UnityEngine.Rendering.ShaderPropertyType.Texture:
                    if (behaviour.textureTarget == RendererBehaviour.TextureTarget.TilingOffset)
                    {
                        target.SetTextureScale(behaviour.propertyName, new Vector2(behaviour.vector.x, behaviour.vector.y));
                        target.SetTextureOffset(behaviour.propertyName, new Vector2(behaviour.vector.z, behaviour.vector.w));
                    }
                    else
                    {
                        target.SetTexture(behaviour.propertyName, behaviour.texture);
                    }
                    break;
            }
        }
    }
}