fix

LoadBalancer.py

@@ -2,7 +2,6 @@ import os
 import requests
 import json
 import urllib.request
-import shutil
 from tqdm import tqdm
 from indexer import indexer
 import re
@@ -37,41 +36,30 @@ class LoadBalancer:
         self.file_structure = None
 
         # Ensure CACHE_DIR exists
-        os.makedirs(self.CACHE_DIR, exist_ok=True)
+        if not os.path.exists(self.CACHE_DIR):
+            os.makedirs(self.CACHE_DIR)
 
-        # Initialize JSON files
-        self.initialize_json_file(self.FILM_STORE_JSON_PATH)
-        self.initialize_json_file(self.TV_STORE_JSON_PATH)
+        for path in [self.FILM_STORE_JSON_PATH, self.TV_STORE_JSON_PATH]:
+            if not os.path.exists(path):
+                with open(path, 'w') as json_file:
+                    json.dump({}, json_file)
 
         # Index the file structure
         indexer()
 
         # Load the file structure JSON
-        self.load_file_structure()
+        if not os.path.exists(self.INDEX_FILE):
+            raise FileNotFoundError(f"{self.INDEX_FILE} not found. Please make sure the file exists.")
+
+        with open(self.INDEX_FILE, 'r') as f:
+            self.file_structure = json.load(f)
 
-        # Start prefetching
         prefetch_thread = Thread(target=self.start_prefetching)
         prefetch_thread.daemon = True
         prefetch_thread.start()
 
-    def initialize_json_file(self, path):
-        """Initialize a JSON file with an empty object if it does not exist."""
-        if not os.path.exists(path):
-            with open(path, 'w') as json_file:
-                json.dump({}, json_file)
-            logging.info(f"Created JSON file: {path}")
-
-    def load_file_structure(self):
-        """Load the file structure from the index file."""
-        if not os.path.exists(self.INDEX_FILE):
-            raise FileNotFoundError(f"{self.INDEX_FILE} not found. Please make sure the file exists.")
-        
-        with open(self.INDEX_FILE, 'r') as f:
-            self.file_structure = json.load(f)
-        logging.info(f"Loaded file structure from {self.INDEX_FILE}")
 
     def register_instance(self, instance_url):
-        """Register a new instance."""
         if instance_url not in self.instances:
             self.instances.append(instance_url)
             logging.info(f"Registered instance {instance_url}")
@@ -79,7 +67,6 @@ class LoadBalancer:
             logging.info(f"Instance {instance_url} is already registered.")
 
     def remove_instance(self, instance_url):
-        """Remove an instance."""
         if instance_url in self.instances:
             self.instances.remove(instance_url)
             self.instances_health.pop(instance_url, None)
@@ -88,101 +75,45 @@ class LoadBalancer:
             logging.info(f"Removed instance {instance_url}")
         else:
             logging.info(f"Instance {instance_url} not found for removal.")
-
+        
     def get_reports(self):
-        """Fetch reports from instances and update JSON files."""
         reports = self.instances_api.fetch_reports()
-        temp_film_store_path = 'temp_film_store.json'
-        temp_tv_store_path = 'temp_tv_store.json'
-
-        # Initialize temporary JSON files
-        self.initialize_json_file(temp_film_store_path)
-        self.initialize_json_file(temp_tv_store_path)
-
-        for instance_url in self.instances[:]:
+        for instance_url in self.instances[:]:  # Copy list to avoid modification during iteration
             if instance_url in reports:
                 report = reports[instance_url]
                 logging.info(f"Report from {instance_url}: {report}")
-                self.process_report(instance_url, report, temp_film_store_path, temp_tv_store_path)
+                self.process_report(instance_url, report)
             else:
                 logging.error(f"Failed to get report from {instance_url}. Removing instance.")
                 self.remove_instance(instance_url)
 
-        # Replace actual JSON files with updated temporary files
-        self.replace_json_files(temp_film_store_path, self.FILM_STORE_JSON_PATH)
-        self.replace_json_files(temp_tv_store_path, self.TV_STORE_JSON_PATH)
-
-    def process_report(self, instance_url, report, temp_film_store_path, temp_tv_store_path):
-        """Process the report from an instance and update temporary JSON files."""
+    def process_report(self, instance_url, report):
         film_store = report.get('film_store', {})
         tv_store = report.get('tv_store', {})
         cache_size = report.get('cache_size')
 
         logging.info(f"Processing report from {instance_url}")
+        # logging.info(f"Film Store: {film_store}")
+        # logging.info(f"TV Store: {tv_store}")
+        # logging.info(f"Cache Size: {cache_size}")
 
-        # Update temporary film store JSON
+        # Process films
         for title, path in film_store.items():
             url = f"{instance_url}/api/film/{title.replace(' ', '%20')}"
-            self.update_film_store_json(title, url, temp_film_store_path)
+            self.update_film_store_json(title, url)
 
-        # Update temporary TV store JSON
+ 
+        # Process TV shows
         for title, seasons in tv_store.items():
             for season, episodes in seasons.items():
                 for episode, path in episodes.items():
                     url = f"{instance_url}/api/tv/{title.replace(' ', '%20')}/{season.replace(' ', '%20')}/{episode.replace(' ', '%20')}"
-                    self.update_tv_store_json(title, season, episode, url, temp_tv_store_path)
+                    self.update_tv_store_json(title, season, episode, url)
 
         logging.info("Film and TV Stores processed successfully.")
         self.update_instances_health(instance=instance_url, cache_size=cache_size)
 
-    def update_film_store_json(self, title, url, temp_film_store_path):
-        """Update the temporary film store JSON with the new file."""
-        film_store_data = {}
-        if os.path.exists(temp_film_store_path):
-            with open(temp_film_store_path, 'r') as json_file:
-                film_store_data = json.load(json_file)
-
-        film_store_data[title] = url
-
-        with open(temp_film_store_path, 'w') as json_file:
-            json.dump(film_store_data, json_file, indent=2)
-
-        logging.info(f'Film store updated with {title}. Content: {film_store_data}')
-
-    def update_tv_store_json(self, title, season, episode, url, temp_tv_store_path):
-        """Update the temporary TV store JSON with the new file."""
-        tv_store_data = {}
-        if os.path.exists(temp_tv_store_path):
-            with open(temp_tv_store_path, 'r') as json_file:
-                tv_store_data = json.load(json_file)
-
-        if title not in tv_store_data:
-            tv_store_data[title] = {}
-        
-        if season not in tv_store_data[title]:
-            tv_store_data[title][season] = {}
-
-        tv_store_data[title][season][episode] = url
-
-        with open(temp_tv_store_path, 'w') as json_file:
-            json.dump(tv_store_data, json_file, indent=2)
-
-        logging.info(f'TV store updated with {title}, {season}, {episode}. Content: {tv_store_data}')
-
-    def replace_json_files(self, temp_file_path, actual_file_path):
-        """Replace the actual JSON file with the temporary JSON file."""
-        try:
-            if os.path.exists(temp_file_path):
-                shutil.copy(temp_file_path, actual_file_path)
-                os.remove(temp_file_path)
-                logging.info(f'Updated {actual_file_path} with new data.')
-            else:
-                logging.error(f'Temporary file {temp_file_path} does not exist.')
-        except Exception as e:
-            logging.error(f'Error replacing JSON files: {e}')
-
     def start_polling(self):
-        """Start polling for reports."""
         logging.info("Starting polling.")
         while not self.stop_event.is_set():
             self.get_reports()
@@ -190,7 +121,6 @@ class LoadBalancer:
         logging.info("Polling stopped.")
 
     def stop_polling(self):
-        """Stop polling."""
         logging.info("Stopping polling.")
         self.stop_event.set()
 
@@ -222,6 +152,51 @@ class LoadBalancer:
             print(f"Error getting system proxies: {e}")
             return {}
 
+    def update_film_store_json(self,title, url):
+        """
+        Updates the film store JSON with the new file.
+
+        Args:
+            title (str): The title of the film.
+            url (str): The url.
+        """
+        FILM_STORE_JSON_PATH = self.FILM_STORE_JSON_PATH
+
+        film_store_data = {}
+        if os.path.exists(FILM_STORE_JSON_PATH):
+            with open(FILM_STORE_JSON_PATH, 'r') as json_file:
+                film_store_data = json.load(json_file)
+
+        film_store_data[title] = url
+
+        with open(FILM_STORE_JSON_PATH, 'w') as json_file:
+            json.dump(film_store_data, json_file, indent=2)
+        print(f'Film store updated with {title}.')
+
+    def update_tv_store_json(self, title, season, episode, url):
+        """
+        Updates the TV store JSON with the new file, organizing by title, season, and episode.
+        """
+        TV_STORE_JSON_PATH = self.TV_STORE_JSON_PATH
+
+        tv_store_data = {}
+        if os.path.exists(TV_STORE_JSON_PATH):
+            with open(TV_STORE_JSON_PATH, 'r') as json_file:
+                tv_store_data = json.load(json_file)
+
+        if title not in tv_store_data:
+            tv_store_data[title] = {}
+        
+        if season not in tv_store_data[title]:
+            tv_store_data[title][season] = {}
+
+        tv_store_data[title][season][episode] = url
+
+        with open(TV_STORE_JSON_PATH, 'w') as json_file:
+            json.dump(tv_store_data, json_file, indent=2)
+
+        print(f'TV store updated with {title}, {season}, {episode}.')
+
     @staticmethod
     def clear_json_file(file_path):
         """