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

public class RemoveTreesUnderTerrainLayer : EditorWindow
{
    public Terrain terrain;
    public string targetLayerName = "Road";
    public float threshold = 0.5f; // Road图层强度大于该值则认为是Road

    [MenuItem("Tools/Terrain/Remove Trees Under Terrain Layer")]
    public static void ShowWindow()
    {
        GetWindow<RemoveTreesUnderTerrainLayer>("Remove Trees by Terrain Layer");
    }

    private void OnGUI()
    {
        terrain = (Terrain)EditorGUILayout.ObjectField("目标 Terrain", terrain, typeof(Terrain), true);
        targetLayerName = EditorGUILayout.TextField("目标 Terrain Layer 名", targetLayerName);
        threshold = EditorGUILayout.Slider("图层权重阈值", threshold, 0f, 1f);

        if (GUILayout.Button("移除指定图层下的树"))
        {
            if (terrain != null)
            {
                RemoveTreesByLayer();
            }
            else
            {
                Debug.LogError("未指定 Terrain！");
            }
        }
    }

    private void RemoveTreesByLayer()
    {
        TerrainData terrainData = terrain.terrainData;
        TreeInstance[] allTrees = terrainData.treeInstances;
        List<TreeInstance> remainingTrees = new List<TreeInstance>();

        int mapWidth = terrainData.alphamapWidth;
        int mapHeight = terrainData.alphamapHeight;
        float[,,] alphaMaps = terrainData.GetAlphamaps(0, 0, mapWidth, mapHeight);

        int targetLayerIndex = -1;
        var terrainLayers = terrainData.terrainLayers;

        for (int i = 0; i < terrainLayers.Length; i++)
        {
            if (terrainLayers[i].name == targetLayerName)
            {
                targetLayerIndex = i;
                break;
            }
        }

        if (targetLayerIndex == -1)
        {
            Debug.LogError($"找不到 Terrain Layer：{targetLayerName}");
            return;
        }

        Vector3 size = terrainData.size;
        Vector3 terrainPos = terrain.transform.position;
        int removed = 0;

        foreach (var tree in allTrees)
        {
            Vector3 worldPos = Vector3.Scale(tree.position, size) + terrainPos;

            int mapX = Mathf.Clamp(Mathf.RoundToInt((tree.position.x) * mapWidth), 0, mapWidth - 1);
            int mapZ = Mathf.Clamp(Mathf.RoundToInt((tree.position.z) * mapHeight), 0, mapHeight - 1);

            float layerWeight = alphaMaps[mapZ, mapX, targetLayerIndex];

            if (layerWeight < threshold)
            {
                remainingTrees.Add(tree);
            }
            else
            {
                removed++;
            }
        }

        terrainData.treeInstances = remainingTrees.ToArray();
        Debug.Log($"✅ 清除完成：移除 {removed} 棵树，保留 {remainingTrees.Count} 棵。");
    }
}