using System.Collections.Generic;
using UnityEngine;
using GamePlay; // for PoolItem

public class NotePool : MonoBehaviour
{
    public static NotePool Instance { get; private set; }

    public GameObject[] notePrefabs;  // 短音符预制体（按颜色存储）
    public GameObject[] holdNotePrefabs; // 长音符片段预制体
    public GameObject[] holdNoteEndPrefabs; // 新增：尾部专用长音符预制体（可选）
    public GameObject[] startNotePrefabs;  // 长音符start片段预制体
    private int poolSize = 12;         // 每种类型的音符池大小
    private int maxPoolSize = 60;      // 池的最大容量

    // 颜色到索引的快速映射，避免每次 switch
    private Dictionary<string, int> colorIndexMap;

    private Dictionary<int, Stack<GameObject>> notePools; // 短音符对象池
    private Dictionary<int, Stack<GameObject>> holdNotePools; // 长音符片段对象池
    private Dictionary<int, Stack<GameObject>> holdNoteEndPools; // 新增：尾部长音符对象池
    private Dictionary<int, Stack<GameObject>> startNotePools; // 长音符start片段对象池

    // 容器物体用于在层级中组织池中的对象，便于调试和管理
    private Transform notePoolContainer;
    private Transform holdPoolContainer;
    private Transform startPoolContainer;

    // 控制调试输出，默认关闭以提升性能
    public bool verboseLogging = false;

    private void Awake()
    {
        if (Instance == null)
        {
            Instance = this;
        }
        else
        {
            Destroy(gameObject);
            return;
        }

        // 初始化颜色映射
        colorIndexMap = new Dictionary<string, int>
        {
            { "red", 0 },
            { "green", 1 },
            { "yellow", 2 },
            { "purple", 3 },
            { "blue", 4 }
        };

        // 创建容器
        notePoolContainer = new GameObject("_NotePool_Notes").transform;
        notePoolContainer.SetParent(transform, false);
        holdPoolContainer = new GameObject("_NotePool_HoldSegments").transform;
        holdPoolContainer.SetParent(transform, false);
        startPoolContainer = new GameObject("_NotePool_StartSegments").transform;
        startPoolContainer.SetParent(transform, false);

        notePools = new Dictionary<int, Stack<GameObject>>();
        holdNotePools = new Dictionary<int, Stack<GameObject>>();
        holdNoteEndPools = new Dictionary<int, Stack<GameObject>>();
        startNotePools = new Dictionary<int, Stack<GameObject>>();

        int colorCount = Mathf.Max(1, notePrefabs != null ? notePrefabs.Length : 0);

        for (int i = 0; i < colorCount; i++)
        {
            notePools[i] = new Stack<GameObject>();
            holdNotePools[i] = new Stack<GameObject>();
            holdNoteEndPools[i] = new Stack<GameObject>();
            startNotePools[i] = new Stack<GameObject>();

            for (int j = 0; j < poolSize; j++)
            {
                AddToPool(notePools[i], notePrefabs[i], notePoolContainer);
                AddToPool(startNotePools[i], startNotePrefabs[i], startPoolContainer);
            }

            for (int j = 0; j < poolSize * 5; j++)
            {
                AddToPool(holdNotePools[i], holdNotePrefabs[i], holdPoolContainer);
                // 如果专用尾部预制体存在则初始化尾部池，否则使用中段预制体
                if (holdNoteEndPrefabs != null && holdNoteEndPrefabs.Length > i && holdNoteEndPrefabs[i] != null)
                    AddToPool(holdNoteEndPools[i], holdNoteEndPrefabs[i], holdPoolContainer);
                else
                    AddToPool(holdNoteEndPools[i], holdNotePrefabs[i], holdPoolContainer);
            }
        }
    }

    private GameObject GetObjectFromPool(Stack<GameObject> pool, GameObject prefab, Transform container)
    {
        if (pool == null)
            return InstantiateAndPrepare(prefab);

        if (pool.Count > 0)
        {
            GameObject obj = pool.Pop();
            if (obj == null)
            {
                // 如果池中对象已被销毁，创建新对象
                obj = InstantiateAndPrepare(prefab);
            }
            else
            {
                // 优化：使用 PoolItem 标识来检查是否来自同一 prefab
                GamePlay.PoolItem pi = obj.GetComponent<GamePlay.PoolItem>();
                if (pi == null || prefab == null || pi.prefabName != prefab.name)
                {
                    // 销毁错误类型的对象并替换为正确的 prefab 实例
                    if (pi == null && verboseLogging) Debug.LogWarning("PoolItem missing on pooled object, replacing.");
                    Destroy(obj);
                    obj = InstantiateAndPrepare(prefab);
                }
            }
            obj.transform.SetParent(null);
            obj.SetActive(true);

            // mark as taken from pool
            GamePlay.PoolItem takenPi = obj.GetComponent<GamePlay.PoolItem>();
            if (takenPi != null) takenPi.inPool = false;

            return obj;
        }
        else
        {
            // 池空时直接实例化新对象
            GameObject obj = InstantiateAndPrepare(prefab);
            obj.SetActive(true);
            return obj;
        }
    }

    private GameObject InstantiateAndPrepare(GameObject prefab)
    {
        if (prefab == null) return null;
        GameObject obj = Instantiate(prefab);
        GamePlay.PoolItem pi = obj.GetComponent<GamePlay.PoolItem>();
        if (pi == null) pi = obj.AddComponent<GamePlay.PoolItem>();
        pi.prefabName = prefab.name;
        pi.inPool = false;
        return obj;
    }

    private void ReturnObjectToPool(Stack<GameObject> pool, GameObject obj, Transform container)
    {
        if (obj == null || pool == null) return; // 防止空引用

        // 避免重复归还：如果对象已在池内，直接忽略
        GamePlay.PoolItem pi = obj.GetComponent<GamePlay.PoolItem>();
        if (pi != null && pi.inPool) return;

        // 将对象移动到池容器并重置变换
        obj.transform.SetParent(container, false);
        obj.transform.localPosition = Vector3.zero;
        obj.transform.localRotation = Quaternion.identity;
        obj.transform.localScale = Vector3.one;

        // 设为非激活状态以便复用
        obj.SetActive(false);

        // 标记为已在池内
        if (pi != null) pi.inPool = true;

        if (pool.Count < maxPoolSize)
        {
            pool.Push(obj);
        }
        else
        {
            // 如果池已满，销毁多余对象以节省内存
            Destroy(obj);
        }
    }

    private void AddToPool(Stack<GameObject> pool, GameObject prefab, Transform container)
    {
        if (prefab == null) return;
        if (pool.Count < maxPoolSize)
        {
            GameObject obj = Instantiate(prefab, container);
            // ensure PoolItem exists and is marked in pool
            GamePlay.PoolItem pi = obj.GetComponent<GamePlay.PoolItem>() ?? obj.AddComponent<GamePlay.PoolItem>();
            pi.prefabName = prefab.name;
            pi.inPool = true;

            obj.SetActive(false);
            pool.Push(obj);
        }
    }

    public GameObject GetNote(string color)
    {
        int colorIndex = GetColorIndexFromName(color);
        return GetObjectFromPool(notePools[colorIndex], notePrefabs[colorIndex], notePoolContainer);
    }

    public GameObject GetStartNote(string color)
    {
        int colorIndex = GetColorIndexFromName(color);
        return GetObjectFromPool(startNotePools[colorIndex], startNotePrefabs[colorIndex], startPoolContainer);
    }

    public GameObject GetHoldNoteSegment(string color)
    {
        int colorIndex = GetColorIndexFromName(color);
        return GetObjectFromPool(holdNotePools[colorIndex], holdNotePrefabs[colorIndex], holdPoolContainer);
    }

    // 新增：从尾部专用池获取尾部片段（如果未配置尾部预制体则降级为中段）
    public GameObject GetHoldNoteEndSegment(string color)
    {
        int colorIndex = GetColorIndexFromName(color);
        GameObject prefab = (holdNoteEndPrefabs != null && holdNoteEndPrefabs.Length > colorIndex && holdNoteEndPrefabs[colorIndex] != null)
            ? holdNoteEndPrefabs[colorIndex]
            : holdNotePrefabs[colorIndex];
        return GetObjectFromPool(holdNoteEndPools[colorIndex], prefab, holdPoolContainer);
    }

    public void ReturnNote(GameObject note, string color)
    {
        if (note == null) return;
        NoteController noteScript = note.GetComponent<NoteController>();
        if (noteScript != null)
        {
            noteScript.ResetState();
        }
        int idx = GetColorIndexFromName(color);
        ReturnObjectToPool(notePools[idx], note, notePoolContainer);
    }

    public void ReturnStartNote(GameObject startNote, string color)
    {
        if (startNote == null) return;
        if (string.IsNullOrEmpty(color))
        {
            if (verboseLogging) Debug.LogWarning($" 归还 StartNote 时 color 为空！音符名: {startNote.name}");
            Destroy(startNote);
            return;
        }

        HoldNote holdNote = startNote.GetComponent<HoldNote>();
        if (holdNote != null)
        {
            holdNote.ResetState(); //  确保重置音符状态
        }

        int idx = GetColorIndexFromName(color);
        ReturnObjectToPool(startNotePools[idx], startNote, startPoolContainer);
    }

    // 新增：归还尾部片段
    public void ReturnHoldNoteEndSegment(GameObject holdNoteEnd, string color)
    {
        if (holdNoteEnd == null) return;
        if (string.IsNullOrEmpty(color))
        {
            if (verboseLogging) Debug.LogWarning($" 归还 HoldNoteEndSegment 时 color 为空！音符名: {holdNoteEnd.name}");
            Destroy(holdNoteEnd);
            return;
        }

        HoldNote holdNoteScript = holdNoteEnd.GetComponent<HoldNote>();
        if (holdNoteScript != null)
        {
            holdNoteScript.ResetState(); // 确保重置音符状态
        }

        int idx = GetColorIndexFromName(color);
        ReturnObjectToPool(holdNoteEndPools[idx], holdNoteEnd, holdPoolContainer);
    }

    public void ReturnHoldNoteSegment(GameObject holdNote, string color)
    {
        if (holdNote == null) return;
        if (string.IsNullOrEmpty(color))
        {
            if (verboseLogging) Debug.LogWarning($" 归还 HoldNoteSegment 时 color 为空！音符名: {holdNote.name}");
            Destroy(holdNote);
            return;
        }

        HoldNote holdNoteScript = holdNote.GetComponent<HoldNote>();
        if (holdNoteScript != null)
        {
            holdNoteScript.ResetState(); //  确保重置音符状态
        }

        int idx = GetColorIndexFromName(color);
        ReturnObjectToPool(holdNotePools[idx], holdNote, holdPoolContainer);
    }

    private int GetColorIndexFromName(string colorName)
    {
        if (string.IsNullOrEmpty(colorName)) return 0;
        if (colorIndexMap != null && colorIndexMap.TryGetValue(colorName, out int idx)) return idx;
        if (verboseLogging) Debug.LogError($"未识别的颜色: {colorName}，默认使用红色");
        return 0;
    }
}

public class PoolItem : MonoBehaviour
{
    public string prefabName;
    public bool inPool;
}