FINALIZE THIS SHIEEEEEET

Assets/Scripts/PuzzleS/PuzzleManager.cs

@@ -33,10 +33,13 @@ namespace PuzzleS
         // Current level puzzle data
         private PuzzleLevelDataSO _currentLevelData;
         private AsyncOperationHandle<PuzzleLevelDataSO> _levelDataLoadOperation;
-        private bool _isLoadingLevelData = false;
+        private bool _isDataLoaded = false;
+
+        // Store registered behaviors that are waiting for data to be loaded
+        private List<ObjectiveStepBehaviour> _registeredBehaviours = new List<ObjectiveStepBehaviour>();
 
         /// <summary>
-        /// Singleton instance of the PuzzleManager. No longer creates an instance if one doesn't exist.
+        /// Singleton instance of the PuzzleManager.
         /// </summary>
         public static PuzzleManager Instance => _instance;
 
@@ -66,6 +69,7 @@ namespace PuzzleS
         {
             // Subscribe to SceneManagerService events after boot is complete
             SceneManagerService.Instance.SceneLoadCompleted += OnSceneLoadCompleted;
+            SceneManagerService.Instance.SceneLoadStarted += OnSceneLoadStarted;
             
             // Find player transform
             _playerTransform = GameObject.FindGameObjectWithTag("Player")?.transform;
@@ -74,7 +78,7 @@ namespace PuzzleS
             StartProximityChecks();
             
             // Load puzzle data for the current scene if not already loading
-            if (_currentLevelData == null && !_isLoadingLevelData)
+            if (_currentLevelData == null && !_isDataLoaded)
             {
                 LoadPuzzleDataForCurrentScene();
             }
@@ -90,6 +94,7 @@ namespace PuzzleS
             if (SceneManagerService.Instance != null)
             {
                 SceneManagerService.Instance.SceneLoadCompleted -= OnSceneLoadCompleted;
+                SceneManagerService.Instance.SceneLoadStarted -= OnSceneLoadStarted;
             }
             
             // Release addressable handle if needed
@@ -99,6 +104,16 @@ namespace PuzzleS
             }
         }
 
+        /// <summary>
+        /// Called when a scene is starting to load
+        /// </summary>
+        public void OnSceneLoadStarted(string sceneName)
+        {
+            // Reset data loaded state when changing scenes to avoid using stale data
+            _isDataLoaded = false;
+            Logging.Debug($"[Puzzles] Scene load started: {sceneName}, marked puzzle data as not loaded");
+        }
+
         /// <summary>
         /// Called when a scene is loaded
         /// </summary>
@@ -132,7 +147,7 @@ namespace PuzzleS
                 return;
             }
             
-            _isLoadingLevelData = true;
+            _isDataLoaded = false;
             string addressablePath = $"Puzzles/{currentScene}";
             
             Logging.Debug($"[Puzzles] Loading puzzle data from addressable: {addressablePath}");
@@ -143,9 +158,12 @@ namespace PuzzleS
                 Addressables.Release(_levelDataLoadOperation);
             }
 
+            // Check if the addressable exists before trying to load it
             if (!AppleHillsUtils.AddressableKeyExists(addressablePath))
             {
-                Logging.Warning($"[Puzzles] Puzzle key does not exist in Addressables: {addressablePath}. Returning early!");
+                Logging.Warning($"[Puzzles] Puzzle data does not exist for scene: {currentScene}");
+                _isDataLoaded = true; // Mark as loaded but with no data
+                _currentLevelData = null;
                 return;
             }
             
@@ -153,8 +171,6 @@ namespace PuzzleS
             _levelDataLoadOperation = Addressables.LoadAssetAsync<PuzzleLevelDataSO>(addressablePath);
             _levelDataLoadOperation.Completed += handle => 
             {
-                _isLoadingLevelData = false;
-                
                 if (handle.Status == AsyncOperationStatus.Succeeded)
                 {
                     _currentLevelData = handle.Result;
@@ -167,8 +183,11 @@ namespace PuzzleS
                     // Unlock initial steps
                     UnlockInitialSteps();
                     
-                    // Update existing ObjectiveStepBehaviours with the current state
-                    UpdateRegisteredStepStates();
+                    // Update all registered behaviors now that data is loaded
+                    UpdateAllRegisteredBehaviors();
+                    
+                    // Mark data as loaded
+                    _isDataLoaded = true;
                     
                     // Notify listeners
                     OnLevelDataLoaded?.Invoke(_currentLevelData);
@@ -176,6 +195,8 @@ namespace PuzzleS
                 else
                 {
                     Logging.Warning($"[Puzzles] Failed to load puzzle data for {currentScene}: {handle.OperationException?.Message}");
+                    _isDataLoaded = true; // Mark as loaded but with error
+                    _currentLevelData = null;
                 }
             };
         }
@@ -238,6 +259,24 @@ namespace PuzzleS
             }
         }
 
+        /// <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>
@@ -246,13 +285,24 @@ namespace PuzzleS
         {
             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}");
                 
-                // Immediately set the correct state based on current puzzle state
-                UpdateStepState(behaviour);
+                // Only update state if data is already loaded
+                if (_isDataLoaded && _currentLevelData != null)
+                {
+                    UpdateStepState(behaviour);
+                }
+                // Otherwise, the state will be updated when data loads in UpdateAllRegisteredBehaviors
             }
         }
         
@@ -279,17 +329,6 @@ namespace PuzzleS
             }
         }
         
-        /// <summary>
-        /// Updates the states of all registered step behaviours based on current puzzle state.
-        /// </summary>
-        private void UpdateRegisteredStepStates()
-        {
-            foreach (var kvp in _stepBehaviours)
-            {
-                UpdateStepState(kvp.Value);
-            }
-        }
-
         /// <summary>
         /// Unregisters a step behaviour from the manager.
         /// </summary>
@@ -297,7 +336,10 @@ namespace PuzzleS
         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}");
         }
 
@@ -446,6 +488,15 @@ namespace PuzzleS
             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;
+        }
+
         void OnApplicationQuit()
         {
             _isQuitting = true;