0.0.2 Alpha test

LoadBalancer.py

@@ -11,7 +11,7 @@ CACHE_DIR = os.getenv("CACHE_DIR")
 
 class LoadBalancer:
     def __init__(self, cache_dir, token, repo, polling_interval=4, max_retries=3, initial_delay=1):
-        self.version = "0.0.1 Alpha"
+        self.version = "0.0.2 Alpha"
         self.instances = []
         self.instances_health = {}
         self.polling_interval = polling_interval
@@ -24,13 +24,14 @@ class LoadBalancer:
         self.REPO = repo
         self.MUSIC_STORE = {}
         self.file_structure = None
+        self.category_files_map = {}
 
         # Ensure CACHE_DIR exists
         if not os.path.exists(self.CACHE_DIR):
             os.makedirs(self.CACHE_DIR)
 
         # Initialize file structure and start prefetching
-        self.file_structure = indexer()
+        self.update_file_structure()
 
         # Start polling and file checking in separate threads
         polling_thread = Thread(target=self.start_polling)
@@ -40,10 +41,22 @@ class LoadBalancer:
         # Start periodic tasks
         asyncio.create_task(self.run_periodic_tasks())
 
+    def update_file_structure(self):
+        """Update the file structure and the category-files map."""
+        self.file_structure = indexer()  # Assume this re-fetches the file structure
+        self.category_files_map = {}  # Reset the map
+
+        for directory in self.file_structure:
+            if directory['type'] == 'directory':
+                # Map category to its files
+                self.category_files_map[directory['path']] = [
+                    file['path'] for file in directory['contents'] if file['type'] == 'file'
+                ]
+
     async def run_periodic_tasks(self):
         """Run indexer and prefetch functions every 5 minutes."""
         while not self.stop_event.is_set():
-            self.file_structure = indexer()  # Re-run indexer
+            self.update_file_structure()  # Re-run indexer and update the map
             await asyncio.sleep(300)  # Sleep for 5 minutes
 
     def get_reports(self):
@@ -106,9 +119,7 @@ class LoadBalancer:
         logging.info(f"Updated instance {instance} with cache size {cache_size}")
 
     def download_music_to_best_instance(self, file_name):
-        """
-        Downloads a music file to the first instance that has more free space on the self.instance_health list variable.
-        """
+        """Downloads a music file to the first instance that has more free space on the self.instance_health list variable."""
         best_instance = None
         max_free_space = -1
 
@@ -159,3 +170,11 @@ class LoadBalancer:
                     if sub_directory['type'] == 'file':
                         music_files.append(sub_directory['path'])
         return music_files
+
+    def get_all_categories(self):
+        """Get a list of all category folders."""
+        return list(self.category_files_map.keys())
+
+    def get_files_from_category(self, category):
+        """Get all files from a specified category."""
+        return self.category_files_map.get(category, [])