try 0.0.2.6 V Beta

LoadBalancer.py

@@ -19,7 +19,7 @@ download_progress = {}
 
 class LoadBalancer:
     def __init__(self, cache_dir, index_file, token, repo, polling_interval=10, max_retries=3, initial_delay=1):
-        self.version = "0.0.2.4 V Beta"
+        self.version = "0.0.2.6 V Beta"
         self.instances = []
         self.instances_health = {}
         self.polling_interval = polling_interval
@@ -33,13 +33,15 @@ class LoadBalancer:
         self.REPO = repo
         self.FILM_STORE_JSON_PATH = os.path.join(cache_dir, "film_store.json")
         self.TV_STORE_JSON_PATH = os.path.join(cache_dir, "tv_store.json")
+        self.TEMP_FILM_STORE_JSON_PATH = os.path.join(cache_dir, "temp_film_store.json")
+        self.TEMP_TV_STORE_JSON_PATH = os.path.join(cache_dir, "temp_tv_store.json")
         self.file_structure = None
 
         # Ensure CACHE_DIR exists
         if not os.path.exists(self.CACHE_DIR):
             os.makedirs(self.CACHE_DIR)
 
-        for path in [self.FILM_STORE_JSON_PATH, self.TV_STORE_JSON_PATH]:
+        for path in [self.FILM_STORE_JSON_PATH, self.TV_STORE_JSON_PATH, self.TEMP_FILM_STORE_JSON_PATH, self.TEMP_TV_STORE_JSON_PATH]:
             if not os.path.exists(path):
                 with open(path, 'w') as json_file:
                     json.dump({}, json_file)
@@ -93,25 +95,22 @@ class LoadBalancer:
         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}")
 
         # 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)
+            self.update_film_store_json(title, url, temp=True)
 
- 
         # 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)
+                    self.update_tv_store_json(title, season, episode, url, temp=True)
 
         logging.info("Film and TV Stores processed successfully.")
         self.update_instances_health(instance=instance_url, cache_size=cache_size)
+        self.replace_store_files()
 
     def start_polling(self):
         logging.info("Starting polling.")
@@ -152,36 +151,44 @@ class LoadBalancer:
             print(f"Error getting system proxies: {e}")
             return {}
 
-    def update_film_store_json(self,title, url):
+    def update_film_store_json(self, title, url, temp=False):
         """
         Updates the film store JSON with the new file.
 
         Args:
             title (str): The title of the film.
             url (str): The url.
+            temp (bool): If True, update the temporary JSON file. Defaults to False.
         """
-        FILM_STORE_JSON_PATH = self.FILM_STORE_JSON_PATH
+        file_path = self.TEMP_FILM_STORE_JSON_PATH if temp else 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:
+        if os.path.exists(file_path):
+            with open(file_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:
+        with open(file_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):
+    def update_tv_store_json(self, title, season, episode, url, temp=False):
         """
         Updates the TV store JSON with the new file, organizing by title, season, and episode.
+
+        Args:
+            title (str): The title of the TV show.
+            season (str): The season of the TV show.
+            episode (str): The episode of the TV show.
+            url (str): The url.
+            temp (bool): If True, update the temporary JSON file. Defaults to False.
         """
-        TV_STORE_JSON_PATH = self.TV_STORE_JSON_PATH
+        file_path = self.TEMP_TV_STORE_JSON_PATH if temp else 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:
+        if os.path.exists(file_path):
+            with open(file_path, 'r') as json_file:
                 tv_store_data = json.load(json_file)
 
         if title not in tv_store_data:
@@ -192,7 +199,7 @@ class LoadBalancer:
 
         tv_store_data[title][season][episode] = url
 
-        with open(TV_STORE_JSON_PATH, 'w') as json_file:
+        with open(file_path, 'w') as json_file:
             json.dump(tv_store_data, json_file, indent=2)
 
         print(f'TV store updated with {title}, {season}, {episode}.')
@@ -230,12 +237,59 @@ class LoadBalancer:
             r'(?:/?|[/?]\S+)$', re.IGNORECASE)
         return re.match(regex, url) is not None
 
-#################################################################
+    @staticmethod
+    def validate_film_store_data(data):
+        """
+        Validate the film store data to ensure it contains valid URLs.
+
+        :param data: The film store data to validate
+        :return: True if valid, False otherwise
+        """
+        for key, value in data.items():
+            if not isinstance(key, str) or not LoadBalancer.is_valid_url(value):
+                return False
+        return True
+
+    @staticmethod
+    def validate_tv_store_data(data):
+        """
+        Validate the TV store data to ensure it contains valid URLs.
 
-    def update_instances_health(self, instance, cache_size):
-        self.instances_health[instance] = {"used":cache_size["cache_size"],
-                                           "total": "50 GB"}
-        logging.info(f"Updated instance {instance} with cache size {cache_size}")
+        :param data: The TV store data to validate
+        :return: True if valid, False otherwise
+        """
+        for show, seasons in data.items():
+            if not isinstance(show, str) or not isinstance(seasons, dict):
+                return False
+            for season, episodes in seasons.items():
+                if not isinstance(season, str) or not isinstance(episodes, dict):
+                    return False
+                for episode, url in episodes.items():
+                    if not isinstance(episode, str) or not LoadBalancer.is_valid_url(url):
+                        return False
+        return True
+
+    def replace_store_files(self):
+        """
+        Replace the actual store JSON files with the temporary ones if they contain valid data.
+        """
+        try:
+            # Load the temporary JSON files
+            with open(self.TEMP_FILM_STORE_JSON_PATH, 'r') as film_file, open(self.TEMP_TV_STORE_JSON_PATH, 'r') as tv_file:
+                temp_film_store_data = json.load(film_file)
+                temp_tv_store_data = json.load(tv_file)
+
+            # Validate the temporary JSON data
+            if self.validate_film_store_data(temp_film_store_data) and self.validate_tv_store_data(temp_tv_store_data):
+                # Replace the actual store files with the temporary files
+                with open(self.FILM_STORE_JSON_PATH, 'w') as film_file, open(self.TV_STORE_JSON_PATH, 'w') as tv_file:
+                    json.dump(temp_film_store_data, film_file)
+                    json.dump(temp_tv_store_data, tv_file)
+                print("Replaced actual store JSON files with temporary ones.")
+            else:
+                print("Temporary JSON data is invalid. Skipping replacement.")
+        except Exception as e:
+            print(f"An error occurred while replacing store files: {e}")
 
 
     def download_film_to_best_instance(self, title):