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

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

    public GameObject[] notePrefabs;  // Documentation text normalized.
    public GameObject[] holdNotePrefabs; // Documentation text normalized.
    public GameObject[] holdNoteEndPrefabs; // Documentation text normalized.
    public GameObject[] startNotePrefabs;  // Documentation text normalized.

    [Header("Pool Size")]
    public int poolSize = 24;
    public int maxPoolSize = 300;
    [Tooltip("Multiplier for hold segment pool size (per color).")]
    public int holdSegmentPoolMultiplier = 8;
    [Tooltip("Multiplier for hold end pool size (per color).")]
    public int holdEndPoolMultiplier = 8;

    // Documentation text normalized.
    private Dictionary<string, int> colorIndexMap;

    private Dictionary<int, Stack<GameObject>> notePools; // Documentation text normalized.
    private Dictionary<int, Stack<GameObject>> holdNotePools;
    private Dictionary<int, Stack<GameObject>> holdNoteEndPools;
    private Dictionary<int, Stack<GameObject>> startNotePools;
    // Documentation text normalized.
    private Transform notePoolContainer;
    private Transform holdPoolContainer;
    private Transform startPoolContainer;

    [Header("Debug")]
    public bool verboseLogging = false;

    [Header("Prewarm Settings")]
    [Tooltip("If true, the pool will be prewarmed across multiple frames to avoid a large GC/instantiation spike at scene load.")]
    public bool prewarmOnStart = true;
    [Tooltip("Number of instantiated pool items to create per frame during prewarming.")]
    // Raised default so heavy pools finish prewarming faster and reduce hitch at first spawn.
    public int prewarmPerFrame = 200;

    // internal prewarm coroutine handle
    private Coroutine prewarmCoroutine = null;

    // Expose whether prewarm has completed
    public bool IsPrewarmed => prewarmCoroutine == null;

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

        // Documentation text normalized.
        colorIndexMap = new Dictionary<string, int>
        {
            { "red", 0 },
            { "green", 1 },
            { "yellow", 2 },
            { "purple", 3 },
            { "blue", 4 }
        };

        // Documentation text normalized.
        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>();
        }

        // minimal immediate entries to ensure safety (one each)
        for (int i = 0; i < colorCount; i++)
        {
            if (notePrefabs != null && i < notePrefabs.Length && notePrefabs[i] != null)
                AddToPool(notePools[i], notePrefabs[i], notePoolContainer);
            if (startNotePrefabs != null && i < startNotePrefabs.Length && startNotePrefabs[i] != null)
                AddToPool(startNotePools[i], startNotePrefabs[i], startPoolContainer);
            if (holdNotePrefabs != null && i < holdNotePrefabs.Length && holdNotePrefabs[i] != null)
                AddToPool(holdNotePools[i], holdNotePrefabs[i], holdPoolContainer);
            if (holdNoteEndPrefabs != null && i < holdNoteEndPrefabs.Length && holdNoteEndPrefabs[i] != null)
                AddToPool(holdNoteEndPools[i], holdNoteEndPrefabs[i], holdPoolContainer);
        }

        if (prewarmOnStart)
        {
            prewarmCoroutine = StartCoroutine(PrewarmCoroutine(colorCount));
        }

        // Also ensure particle and judge prefab warm-up if AnimationController exists in scene.
        var anim = AnimationController.Global ?? SceneObjectLookupCache.FindAny<AnimationController>();
        if (anim != null)
        {
            anim.PrewarmParticles(6);
            if (verboseLogging) Debug.Log("NotePool: requested AnimationController prewarm");
        }
    }

    private IEnumerator PrewarmCoroutine(int colorCount)
    {
        // Calculate targets
        int noteTarget = Mathf.Min(poolSize, maxPoolSize); // per color
        int startTarget = Mathf.Min(poolSize, maxPoolSize);
        int holdTarget = Mathf.Min(poolSize * Mathf.Max(1, holdSegmentPoolMultiplier), maxPoolSize);
        int endTarget = Mathf.Min(poolSize * Mathf.Max(1, holdEndPoolMultiplier), maxPoolSize);

        while (true)
        {
            int createdThisFrame = 0;
            bool allReached = true;

            for (int i = 0; i < colorCount; i++)
            {
                // note pool
                while (notePools[i].Count < noteTarget && createdThisFrame < prewarmPerFrame)
                {
                    if (notePrefabs != null && i < notePrefabs.Length && notePrefabs[i] != null)
                        AddToPool(notePools[i], notePrefabs[i], notePoolContainer);
                    createdThisFrame++;
                }
                if (notePools[i].Count < noteTarget) allReached = false;

                // start pool
                while (startNotePools[i].Count < startTarget && createdThisFrame < prewarmPerFrame)
                {
                    if (startNotePrefabs != null && i < startNotePrefabs.Length && startNotePrefabs[i] != null)
                        AddToPool(startNotePools[i], startNotePrefabs[i], startPoolContainer);
                    createdThisFrame++;
                }
                if (startNotePools[i].Count < startTarget) allReached = false;

                // hold middle pool
                while (holdNotePools[i].Count < holdTarget && createdThisFrame < prewarmPerFrame)
                {
                    if (holdNotePrefabs != null && i < holdNotePrefabs.Length && holdNotePrefabs[i] != null)
                        AddToPool(holdNotePools[i], holdNotePrefabs[i], holdPoolContainer);
                    createdThisFrame++;
                }
                if (holdNotePools[i].Count < holdTarget) allReached = false;

                // hold end pool
                while (holdNoteEndPools[i].Count < endTarget && createdThisFrame < prewarmPerFrame)
                {
                    if (holdNoteEndPrefabs != null && i < holdNoteEndPrefabs.Length && holdNoteEndPrefabs[i] != null)
                        AddToPool(holdNoteEndPools[i], holdNoteEndPrefabs[i], holdPoolContainer);
                    else if (holdNotePrefabs != null && i < holdNotePrefabs.Length && holdNotePrefabs[i] != null)
                        AddToPool(holdNoteEndPools[i], holdNotePrefabs[i], holdPoolContainer);
                    createdThisFrame++;
                }
                if (holdNoteEndPools[i].Count < endTarget) allReached = false;

                if (createdThisFrame >= prewarmPerFrame)
                    break;
            }

            if (allReached) break;
            yield return null;
        }

        prewarmCoroutine = null;
    }

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

        if (pool.Count > 0)
        {
            GameObject obj = pool.Pop();
            GamePlay.PoolItem pi = null;
            if (obj == null)
            {
                // Documentation text normalized.
            }
            else
            {
                // Documentation text normalized.
                pi = obj.GetComponent<GamePlay.PoolItem>();
                if (pi == null || prefab == null || pi.prefabName != prefab.name)
                {
                    // Documentation text normalized.
                    if (pi == null && verboseLogging) Debug.LogWarning("PoolItem missing on pooled object, replacing.");
                    Destroy(obj);
                    obj = InstantiateAndPrepare(prefab);
                    pi = obj != null ? obj.GetComponent<GamePlay.PoolItem>() : null;
                }
            }
            obj.transform.SetParent(null);
            obj.SetActive(true);

            // mark as taken from pool (reuse the PoolItem already resolved above)
            if (pi != null)
            {
                pi.inPool = false;
                if (pi.initialLocalScaleCaptured)
                {
                    obj.transform.localScale = pi.initialLocalScale;
                }
            }

            return obj;
        }
        else
        {
            // Documentation text normalized.
            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;
        if (!pi.initialLocalScaleCaptured)
        {
            pi.initialLocalScale = obj.transform.localScale;
            pi.initialLocalScaleCaptured = true;
        }
        return obj;
    }

    private void ReturnObjectToPool(Stack<GameObject> pool, GameObject obj, Transform container)
    {
        if (obj == null || pool == null) return; // Documentation text normalized.

        // Defensive: unregister/unlock any TrackKeyManager state referencing this object to avoid stale queues/locks
        try
        {
            // Try to determine track index from Note or HoldNote components
            int trackIdx = -1;
            var noteComp = obj.GetComponent<Note>();
            if (noteComp != null)
            {
                trackIdx = noteComp.GetTrackIndex();
            }
            else
            {
                var holdComp = obj.GetComponent<HoldNote>();
                if (holdComp != null)
                {
                    trackIdx = holdComp.trackIndex;
                }
            }

            string id = obj.GetInstanceID().ToString();
            if (TrackKeyManager.Instance != null && trackIdx >= 0)
            {
                try { TrackKeyManager.Instance.UnregisterKey(trackIdx, id); } catch { }
                try { TrackKeyManager.Instance.UnlockTrackForJudge(trackIdx, id); } catch { }
            }
        }
        catch { }

        // Documentation text normalized.
        GamePlay.PoolItem pi = obj.GetComponent<GamePlay.PoolItem>();
        if (pi == null)
        {
            pi = obj.AddComponent<GamePlay.PoolItem>();
        }
        if (!pi.initialLocalScaleCaptured)
        {
            pi.initialLocalScale = obj.transform.localScale;
            pi.initialLocalScaleCaptured = true;
        }
        if (pi != null && pi.inPool) return;

        // Documentation text normalized.
        obj.transform.SetParent(container, false);
        obj.transform.localPosition = Vector3.zero;
        obj.transform.localRotation = Quaternion.identity;
        obj.transform.localScale = pi.initialLocalScaleCaptured ? pi.initialLocalScale : obj.transform.localScale;

        // Documentation text normalized.
        obj.SetActive(false);

        // Documentation text normalized.

        if (pool.Count < maxPoolSize)
        {
            pool.Push(obj);
        }
        else
        {
            // Documentation text normalized.
            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;
            if (!pi.initialLocalScaleCaptured)
            {
                pi.initialLocalScale = obj.transform.localScale;
                pi.initialLocalScaleCaptured = 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(); // Documentation text normalized.
        }

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

    // Documentation text normalized.
    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(); // Documentation text normalized.
        }

        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(); // Documentation text normalized.
        }

        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;
    }
}