using UnityEngine;
using UnityEditor;
using System.Collections.Generic;
using System.Linq;

public class MaterialReplacerWindow : EditorWindow
{
    private List<string> targetNames = new List<string>();
    private List<GameObject> selectedParents = new List<GameObject>();
    private List<Material> randomMaterials = new List<Material>();
    private Material singleMaterial;
    private float minDistance = 2f;
    private bool useRandomMaterials = false;

    private Vector2 scrollPos;
    private Dictionary<GameObject, Material> lastAssignedMaterials = new Dictionary<GameObject, Material>();

    [MenuItem("Tools/Material Replacer")]
    public static void ShowWindow()
    {
        GetWindow<MaterialReplacerWindow>("Material Replacer");
    }

    private void OnGUI()
    {
        scrollPos = EditorGUILayout.BeginScrollView(scrollPos);

        GUILayout.Label("Material Replace Tool", EditorStyles.boldLabel);

        // Target names
        EditorGUILayout.LabelField("Target Names:");
        for (int i = 0; i < targetNames.Count; i++)
        {
            EditorGUILayout.BeginHorizontal();
            targetNames[i] = EditorGUILayout.TextField(targetNames[i]);
            if (GUILayout.Button("X", GUILayout.Width(20)))
            {
                targetNames.RemoveAt(i);
                i--;
            }
            EditorGUILayout.EndHorizontal();
        }

        if (GUILayout.Button("Add Target Name"))
        {
            targetNames.Add("");
        }

        EditorGUILayout.Space();

        if (GUILayout.Button("Load Selected Objects (Including Children for Apply)"))
        {
            selectedParents = Selection.gameObjects.Distinct().ToList();
        }

        EditorGUILayout.LabelField("Loaded Objects: " + selectedParents.Count);

        if (selectedParents.Count > 0)
        {
            if (GUILayout.Button("Clear All Loaded Objects"))
            {
                selectedParents.Clear();
            }

            int removeIndex = -1;
            for (int i = 0; i < selectedParents.Count; i++)
            {
                EditorGUILayout.BeginHorizontal();
                selectedParents[i] = (GameObject)EditorGUILayout.ObjectField(selectedParents[i], typeof(GameObject), true);
                if (GUILayout.Button("X", GUILayout.Width(20)))
                {
                    removeIndex = i;
                }
                EditorGUILayout.EndHorizontal();
            }

            if (removeIndex >= 0)
            {
                selectedParents.RemoveAt(removeIndex);
            }
        }

        EditorGUILayout.Space();
        useRandomMaterials = EditorGUILayout.Toggle("Use Random Materials", useRandomMaterials);

        if (useRandomMaterials)
        {
            EditorGUILayout.LabelField("Random Materials:", EditorStyles.boldLabel);

            int removeMat = -1;
            for (int i = 0; i < randomMaterials.Count; i++)
            {
                EditorGUILayout.BeginHorizontal();
                randomMaterials[i] = (Material)EditorGUILayout.ObjectField(randomMaterials[i], typeof(Material), false);
                if (GUILayout.Button("X", GUILayout.Width(20)))
                {
                    removeMat = i;
                }
                EditorGUILayout.EndHorizontal();
            }

            if (removeMat >= 0)
            {
                randomMaterials.RemoveAt(removeMat);
            }

            EditorGUILayout.BeginHorizontal();
            if (GUILayout.Button("Add Material"))
            {
                randomMaterials.Add(null);
            }
            if (GUILayout.Button("Add Selected Materials from Assets"))
            {
                var selected = Selection.objects
                    .Where(obj => obj is Material)
                    .Cast<Material>()
                    .Where(mat => !randomMaterials.Contains(mat))
                    .ToList();

                if (selected.Count > 0)
                {
                    randomMaterials.AddRange(selected);
                    Debug.Log($"Added {selected.Count} materials from selection.");
                }
                else
                {
                    Debug.Log("No new materials found in selection.");
                }
            }
            if (GUILayout.Button("Clear All"))
            {
                randomMaterials.Clear();
            }
            EditorGUILayout.EndHorizontal();

            minDistance = EditorGUILayout.FloatField("Min Distance Between Same Material", minDistance);
        }
        else
        {
            singleMaterial = (Material)EditorGUILayout.ObjectField("Material", singleMaterial, typeof(Material), false);
        }

        EditorGUILayout.Space();

        if (GUILayout.Button("Apply Materials"))
        {
            ApplyMaterials();
        }

        // Preview
        if (lastAssignedMaterials.Count > 0)
        {
            EditorGUILayout.Space();
            EditorGUILayout.LabelField("Assigned Materials Preview", EditorStyles.boldLabel);

            foreach (var pair in lastAssignedMaterials)
            {
                EditorGUILayout.BeginHorizontal();
                EditorGUILayout.ObjectField(pair.Key, typeof(GameObject), true);
                EditorGUILayout.ObjectField(pair.Value, typeof(Material), false);
                EditorGUILayout.EndHorizontal();
            }
        }

        EditorGUILayout.EndScrollView();
    }

    private void ApplyMaterials()
    {
        lastAssignedMaterials.Clear();

        if (targetNames.Count == 0 || targetNames.All(string.IsNullOrEmpty))
        {
            Debug.LogWarning("No target names specified.");
            return;
        }

        if (selectedParents.Count == 0)
        {
            Debug.LogWarning("No objects loaded.");
            return;
        }

        var allObjects = selectedParents
            .SelectMany(go => go.GetComponentsInChildren<Transform>(true))
            .Select(t => t.gameObject)
            .Distinct()
            .ToList();

        var filtered = allObjects
            .Where(obj => obj != null && targetNames.Contains(obj.name) && obj.TryGetComponent<MeshRenderer>(out _))
            .ToList();

        if (filtered.Count == 0)
        {
            Debug.LogWarning("No matching objects with name and MeshRenderer found.");
            return;
        }

        if (!useRandomMaterials)
        {
            if (singleMaterial == null)
            {
                Debug.LogWarning("No material assigned.");
                return;
            }

            foreach (var obj in filtered)
            {
                var renderer = obj.GetComponent<MeshRenderer>();
                Undo.RecordObject(renderer, "Apply Single Material");
                renderer.sharedMaterial = singleMaterial;
                lastAssignedMaterials[obj] = singleMaterial;
            }
        }
        else
        {
            if (randomMaterials.Count < 2)
            {
                Debug.LogWarning("At least two materials required for randomization.");
                return;
            }

            Dictionary<GameObject, Material> assignment = new Dictionary<GameObject, Material>();
            System.Random rng = new System.Random();

            foreach (var obj in filtered)
            {
                Vector3 pos = obj.transform.position;

                var nearby = assignment
                    .Where(pair => Vector3.Distance(pair.Key.transform.position, pos) < minDistance)
                    .Select(pair => pair.Value)
                    .ToHashSet();

                var available = randomMaterials.Where(m => !nearby.Contains(m)).ToList();

                if (available.Count == 0)
                {
                    available = randomMaterials;
                }

                var chosen = available[rng.Next(available.Count)];
                assignment[obj] = chosen;
            }

            foreach (var pair in assignment)
            {
                var renderer = pair.Key.GetComponent<MeshRenderer>();
                Undo.RecordObject(renderer, "Apply Random Material");
                renderer.sharedMaterial = pair.Value;
                lastAssignedMaterials[pair.Key] = pair.Value;
            }
        }

        Debug.Log($"Materials applied to {filtered.Count} object(s).");
    }
}