using UnityEngine;

/// <summary>
/// Manages grayscale texture assignment for LUT-based gradient coloring.
/// Supports both MaterialPropertyBlock mode (for shared materials) and Material Instance mode (for Timeline animation).
/// </summary>
[ExecuteInEditMode]
[RequireComponent(typeof(Renderer))]
public class PotionTextureSetup : MonoBehaviour
{
    [Header("Grayscale Texture")]
    [Tooltip("Grayscale texture used for LUT (Look-Up Table) gradient coloring. Should be a linear gradient from black to white.")]
    public Texture2D grayscaleTexture;

    private static readonly int GrayscaleTexID = Shader.PropertyToID("_GrayscaleTex");

    private Renderer rend;
    private Liquid liquidComponent;
    private Texture2D lastTexture;
    private MaterialPropertyBlock mpb;
    private bool usesMaterialInstances = false;

    void Awake()
    {
        InitializeComponents();
    }

    void OnEnable()
    {
        InitializeComponents();
        ApplyTexture();
    }

    void Start()
    {
        // Validate components
        if (!ValidateComponents())
        {
            enabled = false;
            return;
        }

        CheckIfUsingInstances();
        ApplyTexture();
    }

    void OnValidate()
    {
        InitializeComponents();
        ApplyTexture();
    }

    void Update()
    {
        // Only update if texture reference changed
        if (grayscaleTexture != lastTexture)
        {
            ApplyTexture();
        }
    }

    /// <summary>
    /// Validates that all required components are present and properly configured.
    /// </summary>
    /// <returns>True if valid, false otherwise.</returns>
    bool ValidateComponents()
    {
        if (rend == null)
        {
            Debug.LogError($"[PotionTextureSetup] Renderer component missing on '{gameObject.name}'. " +
                          "Cannot apply grayscale texture without a Renderer.", this);
            return false;
        }

        if (rend.sharedMaterials == null || rend.sharedMaterials.Length == 0)
        {
            Debug.LogWarning($"[PotionTextureSetup] Renderer on '{gameObject.name}' has no materials assigned. " +
                            "Assign materials to use grayscale texture feature.", this);
            return false;
        }

        // Check if any material has the _GrayscaleTex property
        bool hasGrayscaleProperty = false;
        foreach (var mat in rend.sharedMaterials)
        {
            if (mat != null && mat.HasProperty("_GrayscaleTex"))
            {
                hasGrayscaleProperty = true;
                break;
            }
        }

        if (!hasGrayscaleProperty)
        {
            Debug.LogWarning($"[PotionTextureSetup] None of the materials on '{gameObject.name}' have '_GrayscaleTex' property. " +
                            "This component will have no effect. Use a compatible shader (e.g., Liquid_Effect, Stylized_Tint).", this);
        }

        return true;
    }

    /// <summary>
    /// Initializes component references and MaterialPropertyBlock with GPU instancing enabled.
    /// </summary>
    void InitializeComponents()
    {
        if (rend == null) rend = GetComponent<Renderer>();
        if (liquidComponent == null) liquidComponent = GetComponent<Liquid>();

        if (mpb == null)
        {
            mpb = new MaterialPropertyBlock();
        }

        CheckIfUsingInstances();
    }

    void CheckIfUsingInstances()
    {
        if (rend == null) return;

        Material[] materials = rend.sharedMaterials;
        usesMaterialInstances = false;

        foreach (var mat in materials)
        {
            if (mat != null && mat.name.Contains("(Instance)"))
            {
                usesMaterialInstances = true;
                break;
            }
        }
    }

    private void ApplyTexture()
    {
        if (grayscaleTexture == null) return;

        // If Liquid component exists, delegate to it (handles instance detection internally)
        if (liquidComponent != null)
        {
            liquidComponent.SetGrayscale(grayscaleTexture);
            lastTexture = grayscaleTexture;
            return;
        }

        if (rend == null) return;

        Material[] materials = rend.sharedMaterials;

        if (usesMaterialInstances)
        {
            // Write directly to material instances (for Timeline mode)
            for (int i = 0; i < materials.Length; i++)
            {
                if (materials[i] != null && materials[i].HasProperty("_GrayscaleTex"))
                {
                    materials[i].SetTexture(GrayscaleTexID, grayscaleTexture);
                }
            }
        }
        else
        {
            // ⚡ GPU INSTANCING: Use MPB with instancing enabled
            if (mpb == null) return;

            for (int i = 0; i < materials.Length; i++)
            {
                if (materials[i] != null && materials[i].HasProperty("_GrayscaleTex"))
                {
                    rend.GetPropertyBlock(mpb, i);
                    mpb.SetTexture(GrayscaleTexID, grayscaleTexture);
                    rend.SetPropertyBlock(mpb, i);
                }
            }
        }

        lastTexture = grayscaleTexture;
    }

    /// <summary>
    /// Sets a new grayscale texture and applies it immediately.
    /// </summary>
    /// <param name="texture">The grayscale texture to apply.</param>
    /// <remarks>
    /// Use this method to dynamically change the grayscale texture at runtime.
    /// The texture will be applied according to the current mode (MPB or Material Instance).
    /// </remarks>
    public void SetGrayscaleTexture(Texture2D texture)
    {
        grayscaleTexture = texture;
        ApplyTexture();
    }

    /// <summary>
    /// Creates material instances for Timeline animation.
    /// This allows Timeline to animate all material properties independently per object.
    /// </summary>
    /// <remarks>
    /// After calling this method:
    /// - Each object will have its own material instances (not shared)
    /// - Timeline can animate material properties per object
    /// - Use MaterialIndexController to animate _GradientIndex independently per material slot
    /// 
    /// NOTE: Material instances increase memory usage. Only use when Timeline animation is needed.
    /// To revert, use RemoveInstances() method.
    /// </remarks>
    [ContextMenu("Prepare for Timeline")]
    public void PrepareForTimeline()
    {
        if (rend == null) rend = GetComponent<Renderer>();

        Material[] sharedMats = rend.sharedMaterials;
        bool alreadyPrepared = true;

        // Check if already using instances
        for (int i = 0; i < sharedMats.Length; i++)
        {
            if (sharedMats[i] != null && !sharedMats[i].name.Contains("(Instance)"))
            {
                alreadyPrepared = false;
                break;
            }
        }

        if (alreadyPrepared && sharedMats.Length > 0)
        {
            Debug.Log($"[PotionTextureSetup] '{gameObject.name}' is already using material instances.", this);
            return;
        }

        // Create material instances
        Material[] newInstances = new Material[sharedMats.Length];
        for (int i = 0; i < sharedMats.Length; i++)
        {
            if (sharedMats[i] != null)
            {
                newInstances[i] = new Material(sharedMats[i]);
                newInstances[i].name = sharedMats[i].name + " (Instance)";
            }
        }

        rend.sharedMaterials = newInstances;
        ApplyTexture();

        // Notify other components to switch from MPB to direct material writes
        if (liquidComponent != null)
        {
            liquidComponent.RefreshInstanceDetection();
        }

        MaterialIndexController indexController = GetComponent<MaterialIndexController>();
        if (indexController != null)
        {
            indexController.RefreshInstanceDetection();
        }

        CheckIfUsingInstances();
        ApplyTexture();

#if UNITY_EDITOR
        // Force editor refresh
        UnityEditor.EditorUtility.SetDirty(gameObject);
        UnityEditor.SceneManagement.EditorSceneManager.MarkSceneDirty(gameObject.scene);
#endif

        Debug.Log($"[PotionTextureSetup] Created {newInstances.Length} material instance(s) for '{gameObject.name}'. Timeline-ready!", this);
        Debug.Log("Note: Timeline preview in Edit Mode can be unreliable. Always test in Play Mode for accurate results.", this);

        if (newInstances.Length >= 2)
        {
            Debug.Log("IMPORTANT: For multiple materials, add MaterialIndexController component to animate _GradientIndex independently per material.", this);
        }
        else
        {
            Debug.Log("Tip: You can optionally add MaterialIndexController for a cleaner _GradientIndex slider (not required for single material).", this);
        }
    }

    /// <summary>
    /// Removes material instances and restores the original shared materials.
    /// Use this to revert changes made by PrepareForTimeline().
    /// </summary>
    /// <remarks>
    /// After calling this method:
    /// - Objects will use shared materials again (better for batching/performance)
    /// - Timeline animations will no longer work (need to call PrepareForTimeline again)
    /// - Memory usage will be reduced
    /// </remarks>
    [ContextMenu("Remove Instances (Restore Originals)")]
    public void RemoveInstances()
    {
        if (rend == null) rend = GetComponent<Renderer>();

        Material[] currentMats = rend.sharedMaterials;
        Material[] originalMats = new Material[currentMats.Length];
        bool foundAnyInstances = false;

        for (int i = 0; i < currentMats.Length; i++)
        {
            if (currentMats[i] != null && currentMats[i].name.Contains("(Instance)"))
            {
                foundAnyInstances = true;
                string originalName = currentMats[i].name.Replace(" (Instance)", "").Trim();
                Material originalMat = FindOriginalMaterial(originalName);

                if (originalMat != null)
                {
                    originalMats[i] = originalMat;
                }
                else
                {
                    originalMats[i] = currentMats[i];
                    Debug.LogWarning($"[PotionTextureSetup] Could not find original material '{originalName}' on '{gameObject.name}'. Keeping instance.", this);
                }
            }
            else
            {
                originalMats[i] = currentMats[i];
            }
        }

        if (foundAnyInstances)
        {
            rend.sharedMaterials = originalMats;
            ApplyTexture();

            if (liquidComponent != null)
            {
                liquidComponent.RefreshInstanceDetection();
            }

            MaterialIndexController indexController = GetComponent<MaterialIndexController>();
            if (indexController != null)
            {
                indexController.RefreshInstanceDetection();
            }

            CheckIfUsingInstances();
            ApplyTexture();

#if UNITY_EDITOR
            UnityEditor.EditorUtility.SetDirty(gameObject);
            UnityEditor.SceneManagement.EditorSceneManager.MarkSceneDirty(gameObject.scene);
#endif

            Debug.Log($"[PotionTextureSetup] Material instances removed from '{gameObject.name}'. Restored original materials.", this);
        }
        else
        {
            Debug.Log($"[PotionTextureSetup] No material instances found on '{gameObject.name}'. Nothing to remove.", this);
        }
    }

    /// <summary>
    /// Attempts to find the original shared material by name.
    /// </summary>
    /// <param name="materialName">Name of the material to find.</param>
    /// <returns>The original material if found, null otherwise.</returns>
    Material FindOriginalMaterial(string materialName)
    {
#if UNITY_EDITOR
        // Search in project for material with matching name
        string[] guids = UnityEditor.AssetDatabase.FindAssets($"t:Material {materialName}");
        foreach (string guid in guids)
        {
            string path = UnityEditor.AssetDatabase.GUIDToAssetPath(guid);
            Material mat = UnityEditor.AssetDatabase.LoadAssetAtPath<Material>(path);
            if (mat != null && mat.name == materialName)
            {
                return mat;
            }
        }
#endif
        return null;
    }
}