AppleHillsProduction/Assets/Scripts/PuzzleS/PuzzleManager.cs

using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using UnityEngine;
using UnityEngine.SceneManagement;
using AppleHills.Core.Settings;
using Core;
using Core.Lifecycle;
using UnityEngine.AddressableAssets;
using UnityEngine.ResourceManagement.AsyncOperations;
using Utils;

namespace PuzzleS
{
    /// <summary>
    /// Save data structure for puzzle progress
    /// </summary>
    [Serializable]
    public class PuzzleSaveData
    {
        public string levelId;
        public List<string> completedStepIds;
        public List<string> unlockedStepIds;
    }
    
    /// <summary>
    /// Manages puzzle step registration, dependency management, and step completion for the puzzle system.
    /// </summary>
    public class PuzzleManager : ManagedBehaviour
    {
        private static PuzzleManager _instance;

        [SerializeField] private float proximityCheckInterval = 0.02f;

        // Reference to player transform for proximity calculations
        private Transform _playerTransform;
        private Coroutine _proximityCheckCoroutine;

        // Settings reference
        private IInteractionSettings _interactionSettings;

        // Current level puzzle data
        private PuzzleLevelDataSO _currentLevelData;
        private AsyncOperationHandle<PuzzleLevelDataSO> _levelDataLoadOperation;
        private bool _isDataLoaded = false;

        // Store registered behaviors that are waiting for data to be loaded
        private List<ObjectiveStepBehaviour> _registeredBehaviours = new List<ObjectiveStepBehaviour>();

        // Save system configuration
        public override bool AutoRegisterForSave => true;
        
        /// <summary>
        /// SaveId uses CurrentGameplayScene instead of GetActiveScene() because PuzzleManager
        /// lives in DontDestroyOnLoad and needs to save/load data per-scene.
        /// </summary>
        public override string SaveId
        {
            get
            {
                string sceneName = SceneManagerService.Instance?.CurrentGameplayScene;
                if (string.IsNullOrEmpty(sceneName))
                {
                    // Fallback during early initialization
                    sceneName = SceneManager.GetActiveScene().name;
                }
                return $"{sceneName}/PuzzleManager";
            }
        }
        
        /// <summary>
        /// Singleton instance of the PuzzleManager.
        /// </summary>
        public static PuzzleManager Instance => _instance;

        // Events to notify about step lifecycle
        public event Action<PuzzleStepSO> OnStepCompleted;
        public event Action<PuzzleStepSO> OnStepUnlocked;
        public event Action<PuzzleLevelDataSO> OnLevelDataLoaded;
        public event Action<PuzzleLevelDataSO> OnAllPuzzlesComplete;

        // Save/Load state tracking - string-based for timing independence
        private HashSet<string> _completedSteps = new HashSet<string>();
        private HashSet<string> _unlockedSteps = new HashSet<string>();
        
        // Save/Load restoration tracking
        private bool _isDataRestored = false;
        private List<ObjectiveStepBehaviour> _pendingRegistrations = new List<ObjectiveStepBehaviour>();
        
        // Registration for ObjectiveStepBehaviour
        private Dictionary<PuzzleStepSO, ObjectiveStepBehaviour> _stepBehaviours = new Dictionary<PuzzleStepSO, ObjectiveStepBehaviour>();
        
        // Track pending unlocks for steps that were unlocked before their behavior registered
        private HashSet<string> _pendingUnlocks = new HashSet<string>();
        
        internal override void OnManagedAwake()
        {
            // Set instance immediately (early initialization)
            _instance = this;
        }
        
        internal override void OnManagedStart()
        {
            // Initialize settings reference
            _interactionSettings = GameManager.GetSettingsObject<IInteractionSettings>();
            
            // Find player transform
            _playerTransform = GameObject.FindGameObjectWithTag("Player")?.transform;
            
            // Start proximity check coroutine
            StartProximityChecks();
            
            // Load puzzle data for the current scene if not already loading
            if (_currentLevelData == null && !_isDataLoaded)
            {
                LoadPuzzleDataForCurrentScene();
            }
            
            // Subscribe to scene load events from SceneManagerService
            // This is necessary because PuzzleManager is in DontDestroyOnLoad and won't receive OnSceneReady() callbacks
            if (SceneManagerService.Instance != null)
            {
                SceneManagerService.Instance.SceneLoadCompleted += OnSceneLoadCompleted;
            }
            
            Logging.Debug("[PuzzleManager] Initialized");
        }
        
        /// <summary>
        /// Called when any scene finishes loading. Loads puzzles for the new scene.
        /// </summary>
        private void OnSceneLoadCompleted(string sceneName)
        {
            Logging.Debug($"[Puzzles] Scene loaded: {sceneName}, loading puzzle data");
            LoadPuzzlesForScene(sceneName);
        }
        
        internal override void OnManagedDestroy()
        {
            // Unsubscribe from SceneManagerService events
            if (SceneManagerService.Instance != null)
            {
                SceneManagerService.Instance.SceneLoadCompleted -= OnSceneLoadCompleted;
            }
        }

        /// <summary>
        /// Loads puzzle data for the specified scene
        /// </summary>
        private void LoadPuzzlesForScene(string sceneName)
        {
            // Skip for non-gameplay scenes
            if (sceneName == "BootstrapScene" || string.IsNullOrEmpty(sceneName))
            {
                return;
            }
            
            Logging.Debug($"[Puzzles] Scene loaded: {sceneName}, loading puzzle data");
            LoadPuzzleDataForCurrentScene(sceneName);
            
            // Find player transform again in case it changed with scene load
            _playerTransform = GameObject.FindGameObjectWithTag("Player")?.transform;
            
            // Restart proximity checks
            StartProximityChecks();
        }

        /// <summary>
        /// Load puzzle data for the current scene
        /// </summary>
        private void LoadPuzzleDataForCurrentScene(string sceneName = null)
        {
            string currentScene = sceneName ?? SceneManagerService.Instance.CurrentGameplayScene;
            if (string.IsNullOrEmpty(currentScene))
            {
                Logging.Warning("[Puzzles] Cannot load puzzle data: Current scene name is empty");
                return;
            }
            
            // Reset restoration flag when loading new scene data
            _isDataRestored = false;
            _isDataLoaded = false;
            string addressablePath = $"Puzzles/{currentScene}";
            
            Logging.Debug($"[Puzzles] Loading puzzle data from addressable: {addressablePath}");
            
            // Release previous handle if needed
            if (_levelDataLoadOperation.IsValid())
            {
                Addressables.Release(_levelDataLoadOperation);
            }

            // Check if the addressable exists before trying to load it
            if (!AppleHillsUtils.AddressableKeyExists(addressablePath))
            {
                Logging.Warning($"[Puzzles] Puzzle data does not exist for scene: {currentScene}");
                _isDataLoaded = true; // Mark as loaded but with no data
                _currentLevelData = null;
                return;
            }
            
            // Load the level data asset
            _levelDataLoadOperation = Addressables.LoadAssetAsync<PuzzleLevelDataSO>(addressablePath);
            _levelDataLoadOperation.Completed += handle => 
            {
                if (handle.Status == AsyncOperationStatus.Succeeded)
                {
                    _currentLevelData = handle.Result;
                    Logging.Debug($"[Puzzles] Loaded level data: {_currentLevelData.levelId} with {_currentLevelData.allSteps.Count} steps");
                    
                    // Don't clear steps here - SceneManagerService calls ClearPuzzleState() before scene transitions
                    // This allows save restoration to work properly without race conditions
                    
                    // Unlock initial steps (adds to existing unlocked steps from save restoration)
                    UnlockInitialSteps();
                    
                    // Update all registered behaviors now that data is loaded
                    UpdateAllRegisteredBehaviors();
                    
                    // Mark data as loaded
                    _isDataLoaded = true;
                    
                    // Notify listeners
                    OnLevelDataLoaded?.Invoke(_currentLevelData);
                }
                else
                {
                    Logging.Warning($"[Puzzles] Failed to load puzzle data for {currentScene}: {handle.OperationException?.Message}");
                    _isDataLoaded = true; // Mark as loaded but with error
                    _currentLevelData = null;
                }
            };
        }

        /// <summary>
        /// Start the proximity check coroutine.
        /// </summary>
        private void StartProximityChecks()
        {
            StopProximityChecks();
            _proximityCheckCoroutine = StartCoroutine(CheckProximityRoutine());
        }
        
        /// <summary>
        /// Stop the proximity check coroutine.
        /// </summary>
        private void StopProximityChecks()
        {
            if (_proximityCheckCoroutine != null)
            {
                StopCoroutine(_proximityCheckCoroutine);
                _proximityCheckCoroutine = null;
            }
        }
        
        /// <summary>
        /// Coroutine that periodically checks player proximity to all puzzle steps.
        /// </summary>
        private IEnumerator CheckProximityRoutine()
        {
            WaitForSeconds wait = new WaitForSeconds(proximityCheckInterval);
            
            while (true)
            {
                if (_playerTransform != null)
                {
                    // Get the proximity threshold from settings directly using our settings reference
                    float proximityThreshold = _interactionSettings?.DefaultPuzzlePromptRange ?? 3.0f;
                    
                    // Check distance to each step behavior
                    foreach (var kvp in _stepBehaviours)
                    {
                        if (kvp.Value == null) continue;
                        if (IsPuzzleStepCompleted(kvp.Key.stepId)) continue;
                        
                        float distance = Vector3.Distance(_playerTransform.position, kvp.Value.transform.position);
                        
                        // Determine the proximity state - only Close or Far now
                        ObjectiveStepBehaviour.ProximityState state = 
                            (distance <= proximityThreshold) 
                            ? ObjectiveStepBehaviour.ProximityState.Close 
                            : ObjectiveStepBehaviour.ProximityState.Far;
                        
                        // Update the step's proximity state
                        kvp.Value.UpdateProximityState(state);
                    }
                }
                
                yield return wait;
            }
        }

        /// <summary>
        /// Update all registered behaviors with their current state
        /// </summary>
        private void UpdateAllRegisteredBehaviors()
        {
            foreach (var behaviour in _registeredBehaviours)
            {
                if (behaviour == null) continue;
                
                // Only update if the step is in our dictionary
                bool stepUnlocked = IsStepUnlocked(behaviour.stepData);
                if (stepUnlocked)
                {
                    UpdateStepState(behaviour);
                }
            }
        }

        /// <summary>
        /// Registers a step behaviour with the manager.
        /// </summary>
        /// <param name="behaviour">The step behaviour to register.</param>
        public void RegisterStepBehaviour(ObjectiveStepBehaviour behaviour)
        {
            if (behaviour?.stepData == null) return;
            
            // Always add to our registered behaviors list
            if (!_registeredBehaviours.Contains(behaviour))
            {
                _registeredBehaviours.Add(behaviour);
            }
            
            // Add to the step behaviours dictionary if not already there
            if (!_stepBehaviours.ContainsKey(behaviour.stepData))
            {
                _stepBehaviours.Add(behaviour.stepData, behaviour);
                Logging.Debug($"[Puzzles] Registered step: {behaviour.stepData.stepId} on {behaviour.gameObject.name}");
                
                // Check if this step has a pending unlock
                if (_pendingUnlocks.Contains(behaviour.stepData.stepId))
                {
                    // Step was unlocked before behavior registered - unlock it now!
                    behaviour.UnlockStep();
                    _pendingUnlocks.Remove(behaviour.stepData.stepId);
                    Logging.Debug($"[Puzzles] Fulfilled pending unlock for step: {behaviour.stepData.stepId}");
                }
                else if (_isDataRestored)
                {
                    // Data already restored - update immediately
                    UpdateStepState(behaviour);
                }
                else
                {
                    // Data not restored yet - add to pending queue
                    if (!_pendingRegistrations.Contains(behaviour))
                    {
                        _pendingRegistrations.Add(behaviour);
                    }
                }
            }
        }
        
        /// <summary>
        /// Updates a step's state based on the current puzzle state.
        /// </summary>
        private void UpdateStepState(ObjectiveStepBehaviour behaviour)
        {
            if (behaviour?.stepData == null) return;
            
            // If step is already completed, ignore
            if (_completedSteps.Contains(behaviour.stepData.stepId)) 
                return;
                
            // If step is already unlocked, update the behaviour
            if (_unlockedSteps.Contains(behaviour.stepData.stepId))
            {
                behaviour.UnlockStep();
            }
            else
            {
                // Make sure it's locked
                behaviour.LockStep();
            }
        }
        
        /// <summary>
        /// Unregisters a step behaviour from the manager.
        /// </summary>
        /// <param name="behaviour">The step behaviour to unregister.</param>
        public void UnregisterStepBehaviour(ObjectiveStepBehaviour behaviour)
        {
            if (behaviour?.stepData == null) return;
            
            _stepBehaviours.Remove(behaviour.stepData);
            _registeredBehaviours.Remove(behaviour);
            
            Logging.Debug($"[Puzzles] Unregistered step: {behaviour.stepData.stepId} on {behaviour.gameObject.name}");
        }

        /// <summary>
        /// Unlocks all initial steps (those with no dependencies)
        /// </summary>
        private void UnlockInitialSteps()
        {
            if (_currentLevelData == null) return;
            
            // Don't unlock initial steps if we've restored from save
            if (_isDataRestored)
            {
                Logging.Debug("[Puzzles] Skipping UnlockInitialSteps - data was restored from save");
                return;
            }
            
            // Unlock initial steps
            foreach (var step in _currentLevelData.initialSteps)
            {
                UnlockStep(step);
            }
            
            Logging.Debug($"[Puzzles] Unlocked {_unlockedSteps.Count} initial steps");
        }

        /// <summary>
        /// Called when a step is completed. Unlocks dependent steps if their dependencies are met.
        /// </summary>
        /// <param name="step">The completed step.</param>
        public void MarkPuzzleStepCompleted(PuzzleStepSO step)
        {
            if (_completedSteps.Contains(step.stepId)) return;
            if (_currentLevelData == null) return;
            
            _completedSteps.Add(step.stepId);
            Logging.Debug($"[Puzzles] Step completed: {step.stepId}");

            // Broadcast completion
            OnStepCompleted?.Invoke(step);

            // Unlock steps that are unlocked by this step
            foreach (var unlockStep in _currentLevelData.GetUnlockedSteps(step))
            {
                if (AreStepDependenciesMet(unlockStep))
                {
                    Logging.Debug($"[Puzzles] Unlocking step {unlockStep.stepId} after completing {step.stepId}");
                    UnlockStep(unlockStep);
                }
                else
                {
                    Logging.Debug($"[Puzzles] Step {unlockStep.stepId} not unlocked yet, waiting for other dependencies");
                }
            }
            
            // Check if all puzzle steps are now complete
            CheckPuzzleCompletion();
        }

        /// <summary>
        /// Checks if all dependencies for a step are met.
        /// </summary>
        /// <param name="step">The step to check.</param>
        /// <returns>True if all dependencies are met, false otherwise.</returns>
        private bool AreStepDependenciesMet(PuzzleStepSO step)
        {
            if (_currentLevelData == null || step == null) return false;
            
            // If it's an initial step, it has no dependencies
            if (_currentLevelData.IsInitialStep(step)) return true;
            
            // Check if dependencies are met using pre-processed data
            if (_currentLevelData.stepDependencies.TryGetValue(step.stepId, out string[] dependencies))
            {
                foreach (var depId in dependencies)
                {
                    // Check if dependency is in completed steps
                    if (!_completedSteps.Contains(depId))
                    {
                        return false;
                    }
                }
            }
            
            return true;
        }

        /// <summary>
        /// Unlocks a specific step and notifies its behaviour.
        /// </summary>
        /// <param name="step">The step to unlock.</param>
        private void UnlockStep(PuzzleStepSO step)
        {
            if (_unlockedSteps.Contains(step.stepId)) return;
            _unlockedSteps.Add(step.stepId);
            
            if (_stepBehaviours.TryGetValue(step, out var behaviour))
            {
                // Behavior exists - unlock it immediately
                behaviour.UnlockStep();
                Logging.Debug($"[Puzzles] Step unlocked: {step.stepId}");
            }
            else
            {
                // Behavior hasn't registered yet - add to pending unlocks
                _pendingUnlocks.Add(step.stepId);
                Logging.Debug($"[Puzzles] Step unlocked but behavior not registered yet, added to pending: {step.stepId}");
            }

            // Broadcast unlock
            OnStepUnlocked?.Invoke(step);
        }

        /// <summary>
        /// Checks if the puzzle is complete (all steps in level finished).
        /// </summary>
        private void CheckPuzzleCompletion()
        {
            if (_currentLevelData == null) return;
            
            // Check if all steps are completed
            bool allComplete = true;
            foreach (var step in _currentLevelData.allSteps)
            {
                if (step != null && !_completedSteps.Contains(step.stepId))
                {
                    allComplete = false;
                    break;
                }
            }
            
            if (allComplete)
            {
                Logging.Debug("[Puzzles] All puzzles complete! Level finished.");
                
                // Fire level complete event
                OnAllPuzzlesComplete?.Invoke(_currentLevelData);
            }
        }

        /// <summary>
        /// Returns whether a step is already unlocked.
        /// </summary>
        public bool IsStepUnlocked(PuzzleStepSO step)
        {
            return step != null && _unlockedSteps.Contains(step.stepId);
        }

        /// <summary>
        /// Checks if a puzzle step with the specified ID has been completed
        /// </summary>
        /// <param name="stepId">The ID of the puzzle step to check</param>
        /// <returns>True if the step has been completed, false otherwise</returns>
        public bool IsPuzzleStepCompleted(string stepId)
        {
            return _completedSteps.Contains(stepId);  // O(1) lookup!
        }
        
        /// <summary>
        /// Get the current level puzzle data
        /// </summary>
        public PuzzleLevelDataSO GetCurrentLevelData()
        {
            return _currentLevelData;
        }

        /// <summary>
        /// Checks if puzzle data is loaded
        /// </summary>
        /// <returns>True if data loading has completed (whether successful or not)</returns>
        public bool IsDataLoaded()
        {
            return _isDataLoaded;
        }
        
        /// <summary>
        /// Clears all puzzle state (completed steps, unlocked steps, registrations).
        /// Called by SceneManagerService before scene transitions to ensure clean state.
        /// </summary>
        public void ClearPuzzleState()
        {
            Logging.Debug("[PuzzleManager] Clearing puzzle state");
            
            _completedSteps.Clear();
            _unlockedSteps.Clear();
            _isDataRestored = false;
            
            // Clear any pending registrations from the old scene
            _pendingRegistrations.Clear();
            _pendingUnlocks.Clear();
            
            // Unregister all step behaviours from the old scene
            _stepBehaviours.Clear();
            _registeredBehaviours.Clear();
        }
        
        #region Save/Load Lifecycle Hooks
        
        internal override string OnSceneSaveRequested()
        {
            if (_currentLevelData == null)
            {
                Logging.Warning("[PuzzleManager] Cannot serialize state - no level data loaded");
                return "{}";
            }
            
            var saveData = new PuzzleSaveData
            {
                levelId = _currentLevelData.levelId,
                completedStepIds = _completedSteps.ToList(),
                unlockedStepIds = _unlockedSteps.ToList()
            };
            
            string json = JsonUtility.ToJson(saveData);
            Logging.Debug($"[PuzzleManager] Serialized puzzle state: {_completedSteps.Count} completed, {_unlockedSteps.Count} unlocked");
            return json;
        }
        
        internal override void OnSceneRestoreRequested(string data)
        {
            Logging.Debug("[XAXA] PuzzleManager loading with data: " + data);
            
            if (string.IsNullOrEmpty(data) || data == "{}")
            {
                Logging.Debug("[PuzzleManager] No puzzle save data to restore");
                _isDataRestored = true;
                return;
            }
            
            try
            {
                var saveData = JsonUtility.FromJson<PuzzleSaveData>(data);
                if (saveData == null)
                {
                    Logging.Warning("[PuzzleManager] Failed to deserialize puzzle save data");
                    _isDataRestored = true;
                    return;
                }
                
                // Restore step IDs directly - no timing dependency on level data!
                _completedSteps = new HashSet<string>(saveData.completedStepIds ?? new List<string>());
                _unlockedSteps = new HashSet<string>(saveData.unlockedStepIds ?? new List<string>());
                
                _isDataRestored = true;
                
                Logging.Debug($"[PuzzleManager] Restored puzzle state: {_completedSteps.Count} completed, {_unlockedSteps.Count} unlocked steps");
                
                // Update any behaviors that registered before RestoreState was called
                foreach (var behaviour in _pendingRegistrations)
                {
                    if(behaviour != null)
                        UpdateStepState(behaviour);
                }
                _pendingRegistrations.Clear();
            }
            catch (System.Exception e)
            {
                Debug.LogError($"[PuzzleManager] Error restoring puzzle state: {e.Message}");
                _isDataRestored = true;
            }
        }
        
        #endregion
    }
}