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

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

    // store note instance IDs (as strings) in queue order per track
    // Format: trackIndex -> Queue of (noteId, noteType) tuples
    private Dictionary<int, Queue<(string noteId, string noteType)>> trackKeyMappings =
        new Dictionary<int, Queue<(string, string)>>();

    // Active notes currently inside judge zone per track
    private Dictionary<int, List<string>> trackActiveNotes = new Dictionary<int, List<string>>();
    // Hit times for notes (used to pick best candidate when multiple notes overlap)
    private readonly Dictionary<string, float> noteHitTimes = new Dictionary<string, float>();

    // Track notes currently being judged to prevent double-judging the same note
    // Format: trackIndex -> HashSet of note IDs being judged
    private Dictionary<int, HashSet<string>> trackNotesBeingJudged = new Dictionary<int, HashSet<string>>();

    // transient per-frame consumption: ensures one physical press only affects one note per track per frame
    private Dictionary<int, int> trackConsumedFrame = new Dictionary<int, int>();

    // track last frame we checked for focus loss
    private int lastCheckedFrame = -1;
    private bool lastFocusedState = true;

    // --- Anti-deadlock housekeeping ---
    // Some notes can be pooled/disabled without firing OnTriggerExit2D, leaving stale ids in the queue.
    // To guarantee the track never blocks forever, each queued id gets a TTL.
    // When TTL expires, we prune it as stale.
    private readonly Dictionary<string, float> noteIdExpiry = new Dictionary<string, float>();

    [Header("TrackKeyManager housekeeping")]
    [Tooltip("Seconds a queued note id is allowed to stay without being removed. Prevents permanent deadlocks when OnTriggerExit2D is missed.")]
    [Range(0.5f, 10f)]
    public float queuedIdTtlSeconds = 3f;

    [Tooltip("Maximum number of stale head entries to prune per track per call.")]
    [Range(1, 20)]
    public int pruneBatchSize = 5;

    private Coroutine cleanupCoroutine;
    private readonly WaitForSecondsRealtime cleanupInterval = new WaitForSecondsRealtime(1f);
    private readonly List<int> trackKeysBuffer = new List<int>(16);
    private readonly List<int> trackLockKeysBuffer = new List<int>(16);
    private readonly List<string> noteIdsToRemoveBuffer = new List<string>(32);
    private readonly HashSet<int> allTracksBuffer = new HashSet<int>();
    private readonly HashSet<string> allQueuedNoteIdsBuffer = new HashSet<string>();

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

        // start periodic cleanup to defend against stale locks that can survive missed Unregister/Unlock calls
        StartCleanupRoutine();
    }

    private void OnDestroy()
    {
        StopCleanupRoutine();
    }

    private void StartCleanupRoutine()
    {
        if (cleanupCoroutine == null)
            cleanupCoroutine = StartCoroutine(CleanupStaleLocksRoutine());
    }

    private void StopCleanupRoutine()
    {
        if (cleanupCoroutine != null)
        {
            StopCoroutine(cleanupCoroutine);
            cleanupCoroutine = null;
        }
    }

    // Periodically scan locks and queued ids to remove stale entries that can cause deadlocks.
    private IEnumerator CleanupStaleLocksRoutine()
    {
        while (true)
        {
            // run once per second (unscaled so it runs during pause)
            yield return cleanupInterval;

            float now = Time.unscaledTime;

            // 1) Prune head entries for all tracks (defensive)
            trackKeysBuffer.Clear();
            foreach (var ti in trackKeyMappings.Keys) trackKeysBuffer.Add(ti);
            for (int i = 0; i < trackKeysBuffer.Count; i++)
            {
                PruneStaleHeads(trackKeysBuffer[i]);
            }

            allQueuedNoteIdsBuffer.Clear();
            foreach (var qpair in trackKeyMappings)
            {
                var q = qpair.Value;
                if (q == null || q.Count == 0) continue;
                foreach (var item in q)
                {
                    if (!string.IsNullOrEmpty(item.noteId))
                        allQueuedNoteIdsBuffer.Add(item.noteId);
                }
            }

            // 2) Inspect locks and release any that reference ids that are expired or not present in any queue
            trackLockKeysBuffer.Clear();
            foreach (var t in trackNotesBeingJudged.Keys) trackLockKeysBuffer.Add(t);
            for (int i = 0; i < trackLockKeysBuffer.Count; i++)
            {
                int t = trackLockKeysBuffer[i];
                var locks = trackNotesBeingJudged[t];
                if (locks == null || locks.Count == 0) continue;

                noteIdsToRemoveBuffer.Clear();

                foreach (var noteId in locks)
                {
                    bool hasExpiry = noteIdExpiry.TryGetValue(noteId, out var expiry);
                    bool expired = hasExpiry && now > expiry;

                    // If expiry missing or expired, consider the id stale and remove lock
                    if (!hasExpiry || expired)
                    {
                        noteIdsToRemoveBuffer.Add(noteId);
                        continue;
                    }

                    // Also if the queues do not contain this id anywhere, it is likely stale
                    if (!allQueuedNoteIdsBuffer.Contains(noteId))
                    {
                        noteIdsToRemoveBuffer.Add(noteId);
                    }
                }

                for (int r = 0; r < noteIdsToRemoveBuffer.Count; r++)
                {
                    var id = noteIdsToRemoveBuffer[r];
                    locks.Remove(id);
                    noteIdExpiry.Remove(id);
                    if (JudgeManager.IsDebugEnabled) Debug.LogWarning($"[TrackKeyManager] Cleanup: removed stale lock id={id} on track {t}");
                }

                // If locks become empty, remove dictionary entry to keep structure clean
                if (locks.Count == 0)
                {
                    trackNotesBeingJudged.Remove(t);
                }
            }

            // 3) Defensive: if a track has no queued ids but has consumed-frame or locks, clear them
            allTracksBuffer.Clear();
            foreach (var k in trackKeyMappings.Keys) allTracksBuffer.Add(k);
            foreach (var k in trackNotesBeingJudged.Keys) allTracksBuffer.Add(k);
            foreach (var k in trackConsumedFrame.Keys) allTracksBuffer.Add(k);

            foreach (var track in allTracksBuffer)
            {
                bool hasQueue = trackKeyMappings.ContainsKey(track) && trackKeyMappings[track].Count > 0;
                bool hasLocks = trackNotesBeingJudged.ContainsKey(track) && trackNotesBeingJudged[track].Count > 0;
                if (!hasQueue && hasLocks)
                {
                    // clear locks for this track
                    trackNotesBeingJudged.Remove(track);
                    if (JudgeManager.IsDebugEnabled) Debug.LogWarning($"[TrackKeyManager] Cleanup: cleared locks for empty track {track}");
                }
            }
        }
    }

    private void Update()
    {
        // Clear per-frame consumption dictionary when focus is regained or at frame boundaries
        int currentFrame = Time.frameCount;
        bool isFocused = Application.isFocused;

        // If we've moved to a new frame or focus state changed, clear consumption tracking
        if (lastCheckedFrame != currentFrame)
        {
            // Reset per-frame tracking for the new frame
            // In high-performance scenarios, we might avoid Clear() if the dict is small,
            // but for safety with multiple keys, we keep it.
        }
        
        if (isFocused && !lastFocusedState)
        {
            // On focus regain, clear the consumed frame dictionary to allow input again
            trackConsumedFrame.Clear();
            if (JudgeManager.IsDebugEnabled) Debug.Log("[TrackKeyManager] Focus regained, cleared track consumption");
        }

        lastCheckedFrame = currentFrame;
        lastFocusedState = isFocused;
    }

    private void PruneStaleHeads(int trackIndex)
    {
        if (!trackKeyMappings.TryGetValue(trackIndex, out var q) || q == null || q.Count == 0)
            return;

        int pruned = 0;
        float now = Time.unscaledTime;

        while (q.Count > 0 && pruned < Mathf.Max(1, pruneBatchSize))
        {
            var head = q.Peek();
            if (string.IsNullOrEmpty(head.noteId))
            {
                q.Dequeue();
                pruned++;
                continue;
            }

            // never prune if this id is currently locked for judgement
            if (trackNotesBeingJudged.TryGetValue(trackIndex, out var locks) && locks != null && locks.Contains(head.noteId))
                break;

            // if an expiry exists and has passed, prune
            if (noteIdExpiry.TryGetValue(head.noteId, out var expiry) && now > expiry)
            {
                q.Dequeue();
                noteIdExpiry.Remove(head.noteId);
                noteHitTimes.Remove(head.noteId);
                if (trackActiveNotes.ContainsKey(trackIndex))
                    trackActiveNotes[trackIndex].Remove(head.noteId);
                pruned++;
                if (JudgeManager.IsDebugEnabled)
                    Debug.LogWarning($"[TrackKeyManager] Pruned stale head on track {trackIndex}: id={head.noteId} type={head.noteType} now={now:F2} expiry={expiry:F2}");
                continue;
            }

            // head is still within TTL
            break;
        }
    }

    /// <summary>
    /// Documentation text normalized.
    /// Documentation text normalized.
    /// </summary>
    public void RegisterKey(int trackIndex, string noteInstanceId, string noteType = "tap", float hitTime = float.NaN)
    {
        if (!trackKeyMappings.ContainsKey(trackIndex))
            trackKeyMappings[trackIndex] = new Queue<(string, string)>();

        trackKeyMappings[trackIndex].Enqueue((noteInstanceId, noteType));

        if (!trackActiveNotes.ContainsKey(trackIndex))
            trackActiveNotes[trackIndex] = new List<string>();
        if (!string.IsNullOrEmpty(noteInstanceId) && !trackActiveNotes[trackIndex].Contains(noteInstanceId))
            trackActiveNotes[trackIndex].Add(noteInstanceId);

        if (!string.IsNullOrEmpty(noteInstanceId) && !float.IsNaN(hitTime))
            noteHitTimes[noteInstanceId] = hitTime;

        // record/refresh TTL so the id can't block forever if exit/unregister is missed
        if (!string.IsNullOrEmpty(noteInstanceId))
        {
            noteIdExpiry[noteInstanceId] = Time.unscaledTime + Mathf.Max(0.25f, queuedIdTtlSeconds);
        }
    }

    /// <summary>
    /// Documentation text normalized.
    /// </summary>
    public void UnregisterKey(int trackIndex, string noteInstanceId)
    {
        if (!trackKeyMappings.ContainsKey(trackIndex) || trackKeyMappings[trackIndex].Count == 0)
        {
            // still clear TTL record
            if (!string.IsNullOrEmpty(noteInstanceId))
            {
                noteIdExpiry.Remove(noteInstanceId);
                noteHitTimes.Remove(noteInstanceId);
            }
            if (trackActiveNotes.ContainsKey(trackIndex))
                trackActiveNotes[trackIndex].Remove(noteInstanceId);
            return;
        }

        // If the head matches, dequeue. Otherwise try to remove from queue by rebuilding it.
        if (trackKeyMappings[trackIndex].Count > 0 && trackKeyMappings[trackIndex].Peek().noteId == noteInstanceId)
        {
            trackKeyMappings[trackIndex].Dequeue();
            if (!string.IsNullOrEmpty(noteInstanceId))
            {
                noteIdExpiry.Remove(noteInstanceId);
                noteHitTimes.Remove(noteInstanceId);
            }
            if (trackActiveNotes.ContainsKey(trackIndex))
                trackActiveNotes[trackIndex].Remove(noteInstanceId);
            return;
        }

        // Otherwise remove matching id if present (preserve order of others)
        var q = trackKeyMappings[trackIndex];
        var temp = new Queue<(string, string)>();
        while (q.Count > 0)
        {
            var item = q.Dequeue();
            if (item.noteId != noteInstanceId)
            {
                temp.Enqueue(item);
            }
        }
        trackKeyMappings[trackIndex] = temp;

        if (!string.IsNullOrEmpty(noteInstanceId))
        {
            noteIdExpiry.Remove(noteInstanceId);
            noteHitTimes.Remove(noteInstanceId);
        }
        if (trackActiveNotes.ContainsKey(trackIndex))
            trackActiveNotes[trackIndex].Remove(noteInstanceId);
    }

    /// <summary>
    /// Documentation text normalized.
    /// </summary>
    public string GetCurrentNoteId(int trackIndex)
    {
        PruneStaleHeads(trackIndex);

        if (trackKeyMappings.ContainsKey(trackIndex) && trackKeyMappings[trackIndex].Count > 0)
        {
            return trackKeyMappings[trackIndex].Peek().noteId;
        }
        return null; // Documentation text normalized.
    }

    /// <summary>
    /// Documentation text normalized.
    /// </summary>
    public string GetCurrentNoteType(int trackIndex)
    {
        PruneStaleHeads(trackIndex);

        if (trackKeyMappings.ContainsKey(trackIndex) && trackKeyMappings[trackIndex].Count > 0)
        {
            return trackKeyMappings[trackIndex].Peek().noteType;
        }
        return null;
    }

    /// <summary>
    /// Try to register a note as being judged on a track. Returns true if successful (no other note is currently being judged).
    /// This prevents multiple notes on the same track from being judged by a single key press.
    /// </summary>
    public bool TryLockTrackForJudge(int trackIndex, string noteId)
    {
        if (trackIndex < 0) return false;

        // Before locking, prune any expired heads so they can't block judgement.
        PruneStaleHeads(trackIndex);

        if (!trackNotesBeingJudged.ContainsKey(trackIndex))
        {
            trackNotesBeingJudged[trackIndex] = new HashSet<string>();
        }

        // If there are already notes being judged on this track, reject
        if (trackNotesBeingJudged[trackIndex].Count > 0)
        {
            return false;
        }

        // Lock this note
        trackNotesBeingJudged[trackIndex].Add(noteId);
        return true;
    }

    /// <summary>
    /// Release the lock for a note on a track. Call this after judgment is complete.
    /// </summary>
    public void UnlockTrackForJudge(int trackIndex, string noteId)
    {
        if (trackIndex < 0) return;

        if (trackNotesBeingJudged.ContainsKey(trackIndex))
        {
            trackNotesBeingJudged[trackIndex].Remove(noteId);
        }

        // also clear any expiry record for this note id (it has finished its lifecycle)
        if (!string.IsNullOrEmpty(noteId)) noteIdExpiry.Remove(noteId);
    }

    /// <summary>
    /// Try to mark this track as consumed for the current frame. Returns true if this call
    /// successfully consumes the track (no other note on the same track has consumed this frame).
    /// </summary>
    public bool TryConsumeTrackForFrame(int trackIndex)
    {
        if (trackIndex < 0) return false;
        int currentFrame = Time.frameCount;

        // If the recorded frame is old or invalid, allow consumption
        if (!trackConsumedFrame.ContainsKey(trackIndex))
        {
            trackConsumedFrame[trackIndex] = currentFrame;
            return true;
        }

        // If we're in a different frame, reset and allow
        int recordedFrame = trackConsumedFrame[trackIndex];
        if (recordedFrame != currentFrame)
        {
            trackConsumedFrame[trackIndex] = currentFrame;
            return true;
        }

        // Same frame, already consumed - reject
        return false;
    }

    /// <summary>
    /// Force reset the consumption state (useful for testing or specific scenarios)
    /// </summary>
    public void ResetConsumptionState()
    {
        trackConsumedFrame.Clear();
    }
    public bool IsBestCandidate(int trackIndex, string noteId, float pressTime, float maxWindow)
    {
        if (trackIndex < 0 || string.IsNullOrEmpty(noteId)) return true;
        if (!trackActiveNotes.TryGetValue(trackIndex, out var list) || list == null || list.Count == 0)
            return true;

        string bestId = null;
        float bestDelta = float.MaxValue;
        float bestHit = float.MaxValue;

        for (int i = 0; i < list.Count; i++)
        {
            var id = list[i];
            if (string.IsNullOrEmpty(id)) continue;
            if (!noteHitTimes.TryGetValue(id, out var ht)) continue;
            float delta = Mathf.Abs(pressTime - ht);
            if (delta > maxWindow) continue;

            if (delta < bestDelta - 0.0001f)
            {
                bestDelta = delta;
                bestId = id;
                bestHit = ht;
            }
            else if (Mathf.Abs(delta - bestDelta) <= 0.0001f)
            {
                // tie-breaker: earlier hit time wins
                if (ht < bestHit)
                {
                    bestId = id;
                    bestHit = ht;
                }
            }
        }

        if (bestId == null) return true;
        return bestId == noteId;
    }

    /// <summary>
    /// Force clear all track locks and per-frame consumption state.
    /// Used to simulate a global input release (e.g. when ending playback).
    /// </summary>
    public void ClearAllLocks()
    {
        trackNotesBeingJudged.Clear();
        trackConsumedFrame.Clear();
        noteIdExpiry.Clear();
        noteHitTimes.Clear();
        trackActiveNotes.Clear();
        if (JudgeManager.IsDebugEnabled) Debug.Log("[TrackKeyManager] ClearAllLocks called: cleared locks and consumption state");
    }
}