#1859 - First pass at an implementation of the full reset method. Will now start testing...

src/primaite/simulator/file_system/file.py

@@ -73,6 +73,18 @@ class File(FileSystemItemABC):
 
         self.sys_log.info(f"Created file /{self.path} (id: {self.uuid})")
 
+        self.set_original_state()
+
+    def set_original_state(self):
+        """Sets the original state."""
+        super().set_original_state()
+        vals_to_include = {"folder_id", "folder_name", "file_type", "sim_size", "real", "sim_path", "sim_root"}
+        self._original_state.update(self.model_dump(include=vals_to_include))
+
+    def reset_component_for_episode(self, episode: int):
+        """Reset the original state of the SimComponent."""
+        super().reset_component_for_episode(episode)
+
     @property
     def path(self) -> str:
         """

src/primaite/simulator/file_system/file_system.py

@@ -35,6 +35,36 @@ class FileSystem(SimComponent):
         if not self.folders:
             self.create_folder("root")
 
+    def set_original_state(self):
+        """Sets the original state."""
+        for folder in self.folders.values():
+            folder.set_original_state()
+        super().set_original_state()
+        # Capture a list of all 'original' file uuids
+        self._original_state["original_folder_uuids"] = list(self.folders.keys())
+
+    def reset_component_for_episode(self, episode: int):
+        """Reset the original state of the SimComponent."""
+        # Move any 'original' folder that have been deleted back to folders
+        original_folder_uuids = self._original_state.pop("original_folder_uuids")
+        for uuid in original_folder_uuids:
+            if uuid in self.deleted_folders:
+                self.folders[uuid] = self.deleted_folders.pop(uuid)
+
+        # Clear any other deleted folders that aren't original (have been created by agent)
+        self.deleted_folders.clear()
+
+        # Now clear all non-original folders created by agent
+        current_folder_uuids = list(self.folders.keys())
+        for uuid in current_folder_uuids:
+            if uuid not in original_folder_uuids:
+                self.folders.pop(uuid)
+
+        # Now reset all remaining folders
+        for folder in self.folders.values():
+            folder.reset_component_for_episode(episode)
+        super().reset_component_for_episode(episode)
+
     def _init_request_manager(self) -> RequestManager:
         rm = super()._init_request_manager()
 

src/primaite/simulator/file_system/file_system_item_abc.py

@@ -85,6 +85,11 @@ class FileSystemItemABC(SimComponent):
     deleted: bool = False
     "If true, the FileSystemItem was deleted."
 
+    def set_original_state(self):
+        """Sets the original state."""
+        vals_to_keep = {"name", "health_status", "visible_health_status", "previous_hash", "revealed_to_red"}
+        self._original_state = self.model_dump(include=vals_to_keep)
+
     def describe_state(self) -> Dict:
         """
         Produce a dictionary describing the current state of this object.

src/primaite/simulator/file_system/folder.py

@@ -51,6 +51,44 @@ class Folder(FileSystemItemABC):
 
         self.sys_log.info(f"Created file /{self.name} (id: {self.uuid})")
 
+    def set_original_state(self):
+        """Sets the original state."""
+        for file in self.files.values():
+            file.set_original_state()
+        super().set_original_state()
+        vals_to_include = {
+            "scan_duration",
+            "scan_countdown",
+            "red_scan_duration",
+            "red_scan_countdown",
+            "restore_duration",
+            "restore_countdown",
+        }
+        self._original_state.update(self.model_dump(include=vals_to_include))
+        self._original_state["original_file_uuids"] = list(self.files.keys())
+
+    def reset_component_for_episode(self, episode: int):
+        """Reset the original state of the SimComponent."""
+        # Move any 'original' file that have been deleted back to files
+        original_file_uuids = self._original_state.pop("original_file_uuids")
+        for uuid in original_file_uuids:
+            if uuid in self.deleted_files:
+                self.files[uuid] = self.deleted_files.pop(uuid)
+
+        # Clear any other deleted files that aren't original (have been created by agent)
+        self.deleted_files.clear()
+
+        # Now clear all non-original files created by agent
+        current_file_uuids = list(self.files.keys())
+        for uuid in current_file_uuids:
+            if uuid not in original_file_uuids:
+                self.files.pop(uuid)
+
+        # Now reset all remaining files
+        for file in self.files.values():
+            file.reset_component_for_episode(episode)
+        super().reset_component_for_episode(episode)
+
     def _init_request_manager(self) -> RequestManager:
         rm = super()._init_request_manager()
         rm.add_request(