ChandimaPrabath commited on
Commit
56a64e0
·
1 Parent(s): 151773c

0.0.2.8 V Beta

Browse files
Files changed (4) hide show
  1. LoadBalancer.py +51 -63
  2. app.py +20 -215
  3. main.py +219 -0
  4. requirements.txt +3 -2
LoadBalancer.py CHANGED
@@ -1,11 +1,10 @@
1
  import os
2
  import json
3
- from indexer import indexer
 
4
  import re
 
5
  from tvdb import fetch_and_cache_json
6
- from threading import Event, Thread
7
- import time
8
- import logging
9
  from utils import convert_to_gb
10
  from api import InstancesAPI
11
 
@@ -15,13 +14,13 @@ download_progress = {}
15
 
16
  class LoadBalancer:
17
  def __init__(self, cache_dir, index_file, token, repo, polling_interval=4, max_retries=3, initial_delay=1):
18
- self.version = "0.0.2.7 V Beta"
19
  self.instances = []
20
  self.instances_health = {}
21
  self.polling_interval = polling_interval
22
  self.max_retries = max_retries
23
  self.initial_delay = initial_delay
24
- self.stop_event = Event()
25
  self.instances_api = InstancesAPI(self.instances)
26
  self.CACHE_DIR = cache_dir
27
  self.INDEX_FILE = index_file
@@ -42,16 +41,11 @@ class LoadBalancer:
42
  # Load the file structure JSON
43
  self.load_file_structure()
44
 
45
- # Start polling and file checking in separate threads
46
- polling_thread = Thread(target=self.start_polling)
47
- polling_thread.daemon = True
48
- polling_thread.start()
49
-
50
- file_checking_thread = Thread(target=self.check_file_updates)
51
- file_checking_thread.daemon = True
52
- file_checking_thread.start()
53
 
54
- def load_file_structure(self):
55
  if not os.path.exists(self.INDEX_FILE):
56
  raise FileNotFoundError(f"{self.INDEX_FILE} not found. Please make sure the file exists.")
57
 
@@ -59,23 +53,21 @@ class LoadBalancer:
59
  self.file_structure = json.load(f)
60
  logging.info("File structure loaded successfully.")
61
 
62
- def check_file_updates(self):
63
  while not self.stop_event.is_set():
64
  if self.index_file_last_modified != os.path.getmtime(self.INDEX_FILE):
65
  logging.info(f"{self.INDEX_FILE} has been updated. Re-indexing...")
66
  indexer() # Re-run the indexer
67
- self.load_file_structure() # Reload the file structure
68
  self.index_file_last_modified = os.path.getmtime(self.INDEX_FILE)
69
 
70
- # Restart prefetching thread
71
- if hasattr(self, 'prefetch_thread') and self.prefetch_thread.is_alive():
72
- self.prefetch_thread.join()
73
 
74
- self.prefetch_thread = Thread(target=self.start_prefetching)
75
- self.prefetch_thread.daemon = True
76
- self.prefetch_thread.start()
77
 
78
- time.sleep(120) # Check every 2 minutes
79
 
80
  def register_instance(self, instance_url):
81
  if instance_url not in self.instances:
@@ -92,8 +84,8 @@ class LoadBalancer:
92
  else:
93
  logging.info(f"Instance {instance_url} not found for removal.")
94
 
95
- def get_reports(self):
96
- reports = self.instances_api.fetch_reports()
97
 
98
  # Initialize temporary JSON data holders
99
  temp_film_store = {}
@@ -137,30 +129,27 @@ class LoadBalancer:
137
  logging.info("Film and TV Stores processed successfully.")
138
  self.update_instances_health(instance=instance_url, cache_size=cache_size)
139
 
140
- def start_polling(self):
141
  logging.info("Starting polling.")
142
  while not self.stop_event.is_set():
143
- self.get_reports()
144
- time.sleep(self.polling_interval)
145
  logging.info("Polling stopped.")
146
 
147
- def stop_polling(self):
148
  logging.info("Stopping polling.")
149
  self.stop_event.set()
150
 
151
- def start_prefetching(self):
152
- """Start the metadata prefetching in a separate thread."""
153
- self.prefetch_metadata()
154
-
155
- #################################################################
156
 
157
  def update_instances_health(self, instance, cache_size):
158
- self.instances_health[instance] = {"used":cache_size["cache_size"],
159
  "total": "50 GB"}
160
  logging.info(f"Updated instance {instance} with cache size {cache_size}")
161
 
162
-
163
- def download_film_to_best_instance(self, title):
164
  """
165
  Downloads a film to the first instance that has more free space on the self.instance_health list variable.
166
  The instance_health looks like this:
@@ -187,14 +176,14 @@ class LoadBalancer:
187
  best_instance = instance_url
188
 
189
  if best_instance:
190
- result = self.instances_api.download_film(best_instance, title)
191
  film_id = result["film_id"]
192
  status = result["status"]
193
  progress_url = f'{best_instance}/api/progress/{film_id}'
194
  response = {
195
- "film_id":film_id,
196
- "status":status,
197
- "progress_url":progress_url
198
  }
199
 
200
  return response
@@ -202,9 +191,9 @@ class LoadBalancer:
202
  logging.error("No suitable instance found for downloading the film.")
203
  return {"error": "No suitable instance found for downloading the film."}
204
 
205
- def download_episode_to_best_instance(self, title, season, episode):
206
  """
207
- Downloads a episode to the first instance that has more free space on the self.instance_health list variable.
208
  The instance_health looks like this:
209
  {
210
  "https://unicone-studio-instance1.hf.space": {
@@ -213,9 +202,9 @@ class LoadBalancer:
213
  }
214
  }
215
  Args:
216
- title (str): The title of the Tv show.
217
- season (str): The season of the Tv show.
218
- episode (str): The title of the Tv show.
219
  """
220
  best_instance = None
221
  max_free_space = -1
@@ -231,23 +220,22 @@ class LoadBalancer:
231
  best_instance = instance_url
232
 
233
  if best_instance:
234
- result = self.instances_api.download_episode(best_instance, title, season, episode)
235
  episode_id = result["episode_id"]
236
  status = result["status"]
237
  progress_url = f'{best_instance}/api/progress/{episode_id}'
238
  response = {
239
- "episode_id":episode_id,
240
- "status":status,
241
- "progress_url":progress_url
242
  }
243
 
244
  return response
245
  else:
246
- logging.error("No suitable instance found for downloading the film.")
247
- return {"error": "No suitable instance found for downloading the film."}
248
 
249
- #################################################################
250
- def find_movie_path(self, title):
251
  """Find the path of the movie in the JSON data based on the title."""
252
  for directory in self.file_structure:
253
  if directory['type'] == 'directory' and directory['path'] == 'films':
@@ -258,7 +246,7 @@ class LoadBalancer:
258
  return item['path']
259
  return None
260
 
261
- def find_tv_path(self, title):
262
  """Find the path of the TV show in the JSON data based on the title."""
263
  for directory in self.file_structure:
264
  if directory['type'] == 'directory' and directory['path'] == 'tv':
@@ -267,7 +255,7 @@ class LoadBalancer:
267
  return sub_directory['path']
268
  return None
269
 
270
- def get_tv_structure(self, title):
271
  """Find the path of the TV show in the JSON data based on the title."""
272
  for directory in self.file_structure:
273
  if directory['type'] == 'directory' and directory['path'] == 'tv':
@@ -276,11 +264,11 @@ class LoadBalancer:
276
  return sub_directory
277
  return None
278
 
279
- def get_film_id(self, title):
280
  """Generate a film ID based on the title."""
281
  return title.replace(" ", "_").lower()
282
 
283
- def prefetch_metadata(self):
284
  """Prefetch metadata for all items in the file structure."""
285
  for item in self.file_structure:
286
  if 'contents' in item:
@@ -303,9 +291,9 @@ class LoadBalancer:
303
  title = parts[0].strip()
304
  year = int(parts[-1])
305
 
306
- fetch_and_cache_json(original_title, title, media_type, year)
307
 
308
- def get_all_tv_shows(self):
309
  """Get all TV shows from the indexed cache structure JSON file."""
310
  tv_shows = {}
311
  for directory in self.file_structure:
@@ -326,7 +314,7 @@ class LoadBalancer:
326
  })
327
  return tv_shows
328
 
329
- def get_all_films(self):
330
  """Get all films from the indexed cache structure JSON file."""
331
  films = []
332
  for directory in self.file_structure:
@@ -334,4 +322,4 @@ class LoadBalancer:
334
  for sub_directory in directory['contents']:
335
  if sub_directory['type'] == 'directory':
336
  films.append(sub_directory['path'])
337
- return films
 
1
  import os
2
  import json
3
+ import asyncio
4
+ import logging
5
  import re
6
+ from indexer import indexer
7
  from tvdb import fetch_and_cache_json
 
 
 
8
  from utils import convert_to_gb
9
  from api import InstancesAPI
10
 
 
14
 
15
  class LoadBalancer:
16
  def __init__(self, cache_dir, index_file, token, repo, polling_interval=4, max_retries=3, initial_delay=1):
17
+ self.version = "0.0.2.8 V Beta"
18
  self.instances = []
19
  self.instances_health = {}
20
  self.polling_interval = polling_interval
21
  self.max_retries = max_retries
22
  self.initial_delay = initial_delay
23
+ self.stop_event = asyncio.Event()
24
  self.instances_api = InstancesAPI(self.instances)
25
  self.CACHE_DIR = cache_dir
26
  self.INDEX_FILE = index_file
 
41
  # Load the file structure JSON
42
  self.load_file_structure()
43
 
44
+ # Start polling and file checking in separate tasks
45
+ asyncio.create_task(self.start_polling())
46
+ asyncio.create_task(self.check_file_updates())
 
 
 
 
 
47
 
48
+ async def load_file_structure(self):
49
  if not os.path.exists(self.INDEX_FILE):
50
  raise FileNotFoundError(f"{self.INDEX_FILE} not found. Please make sure the file exists.")
51
 
 
53
  self.file_structure = json.load(f)
54
  logging.info("File structure loaded successfully.")
55
 
56
+ async def check_file_updates(self):
57
  while not self.stop_event.is_set():
58
  if self.index_file_last_modified != os.path.getmtime(self.INDEX_FILE):
59
  logging.info(f"{self.INDEX_FILE} has been updated. Re-indexing...")
60
  indexer() # Re-run the indexer
61
+ await self.load_file_structure() # Reload the file structure
62
  self.index_file_last_modified = os.path.getmtime(self.INDEX_FILE)
63
 
64
+ # Restart prefetching task
65
+ if hasattr(self, 'prefetch_task') and not self.prefetch_task.done():
66
+ await self.prefetch_task
67
 
68
+ self.prefetch_task = asyncio.create_task(self.start_prefetching())
 
 
69
 
70
+ await asyncio.sleep(120) # Check every 2 minutes
71
 
72
  def register_instance(self, instance_url):
73
  if instance_url not in self.instances:
 
84
  else:
85
  logging.info(f"Instance {instance_url} not found for removal.")
86
 
87
+ async def get_reports(self):
88
+ reports = await self.instances_api.fetch_reports()
89
 
90
  # Initialize temporary JSON data holders
91
  temp_film_store = {}
 
129
  logging.info("Film and TV Stores processed successfully.")
130
  self.update_instances_health(instance=instance_url, cache_size=cache_size)
131
 
132
+ async def start_polling(self):
133
  logging.info("Starting polling.")
134
  while not self.stop_event.is_set():
135
+ await self.get_reports()
136
+ await asyncio.sleep(self.polling_interval)
137
  logging.info("Polling stopped.")
138
 
139
+ async def stop_polling(self):
140
  logging.info("Stopping polling.")
141
  self.stop_event.set()
142
 
143
+ async def start_prefetching(self):
144
+ """Start the metadata prefetching."""
145
+ await self.prefetch_metadata()
 
 
146
 
147
  def update_instances_health(self, instance, cache_size):
148
+ self.instances_health[instance] = {"used": cache_size["cache_size"],
149
  "total": "50 GB"}
150
  logging.info(f"Updated instance {instance} with cache size {cache_size}")
151
 
152
+ async def download_film_to_best_instance(self, title):
 
153
  """
154
  Downloads a film to the first instance that has more free space on the self.instance_health list variable.
155
  The instance_health looks like this:
 
176
  best_instance = instance_url
177
 
178
  if best_instance:
179
+ result = await self.instances_api.download_film(best_instance, title)
180
  film_id = result["film_id"]
181
  status = result["status"]
182
  progress_url = f'{best_instance}/api/progress/{film_id}'
183
  response = {
184
+ "film_id": film_id,
185
+ "status": status,
186
+ "progress_url": progress_url
187
  }
188
 
189
  return response
 
191
  logging.error("No suitable instance found for downloading the film.")
192
  return {"error": "No suitable instance found for downloading the film."}
193
 
194
+ async def download_episode_to_best_instance(self, title, season, episode):
195
  """
196
+ Downloads an episode to the first instance that has more free space on the self.instance_health list variable.
197
  The instance_health looks like this:
198
  {
199
  "https://unicone-studio-instance1.hf.space": {
 
202
  }
203
  }
204
  Args:
205
+ title (str): The title of the TV show.
206
+ season (str): The season of the TV show.
207
+ episode (str): The episode of the TV show.
208
  """
209
  best_instance = None
210
  max_free_space = -1
 
220
  best_instance = instance_url
221
 
222
  if best_instance:
223
+ result = await self.instances_api.download_episode(best_instance, title, season, episode)
224
  episode_id = result["episode_id"]
225
  status = result["status"]
226
  progress_url = f'{best_instance}/api/progress/{episode_id}'
227
  response = {
228
+ "episode_id": episode_id,
229
+ "status": status,
230
+ "progress_url": progress_url
231
  }
232
 
233
  return response
234
  else:
235
+ logging.error("No suitable instance found for downloading the episode.")
236
+ return {"error": "No suitable instance found for downloading the episode."}
237
 
238
+ async def find_movie_path(self, title):
 
239
  """Find the path of the movie in the JSON data based on the title."""
240
  for directory in self.file_structure:
241
  if directory['type'] == 'directory' and directory['path'] == 'films':
 
246
  return item['path']
247
  return None
248
 
249
+ async def find_tv_path(self, title):
250
  """Find the path of the TV show in the JSON data based on the title."""
251
  for directory in self.file_structure:
252
  if directory['type'] == 'directory' and directory['path'] == 'tv':
 
255
  return sub_directory['path']
256
  return None
257
 
258
+ async def get_tv_structure(self, title):
259
  """Find the path of the TV show in the JSON data based on the title."""
260
  for directory in self.file_structure:
261
  if directory['type'] == 'directory' and directory['path'] == 'tv':
 
264
  return sub_directory
265
  return None
266
 
267
+ async def get_film_id(self, title):
268
  """Generate a film ID based on the title."""
269
  return title.replace(" ", "_").lower()
270
 
271
+ async def prefetch_metadata(self):
272
  """Prefetch metadata for all items in the file structure."""
273
  for item in self.file_structure:
274
  if 'contents' in item:
 
291
  title = parts[0].strip()
292
  year = int(parts[-1])
293
 
294
+ await fetch_and_cache_json(original_title, title, media_type, year)
295
 
296
+ async def get_all_tv_shows(self):
297
  """Get all TV shows from the indexed cache structure JSON file."""
298
  tv_shows = {}
299
  for directory in self.file_structure:
 
314
  })
315
  return tv_shows
316
 
317
+ async def get_all_films(self):
318
  """Get all films from the indexed cache structure JSON file."""
319
  films = []
320
  for directory in self.file_structure:
 
322
  for sub_directory in directory['contents']:
323
  if sub_directory['type'] == 'directory':
324
  films.append(sub_directory['path'])
325
+ return films
app.py CHANGED
@@ -1,217 +1,22 @@
1
- from flask import Flask, jsonify, request, send_from_directory
2
- from flask_cors import CORS
3
- from utils import is_valid_url, bytes_to_human_readable, encode_episodeid
4
- import os
5
- import json
6
- from threading import Thread
7
- import urllib.parse
8
- from LoadBalancer import LoadBalancer
9
- import logging
 
 
 
 
 
 
10
 
11
-
12
- app = Flask(__name__)
13
- CORS(app)
14
-
15
- logging.basicConfig(level=logging.INFO)
16
- # Constants and Configuration
17
- CACHE_DIR = os.getenv("CACHE_DIR")
18
- INDEX_FILE = os.getenv("INDEX_FILE")
19
- TOKEN = os.getenv("TOKEN")
20
- REPO = os.getenv("REPO")
21
-
22
- load_balancer = LoadBalancer(cache_dir=CACHE_DIR, index_file=INDEX_FILE, token=TOKEN, repo=REPO)
23
-
24
- # Start polling in a separate thread
25
- polling_thread = Thread(target=load_balancer.start_polling)
26
- polling_thread.start()
27
-
28
- # API Endpoints
29
- @app.route('/api/film/<title>', methods=['GET'])
30
- def get_movie_api(title):
31
- """Endpoint to get the movie by title."""
32
- if not title:
33
- return jsonify({"error": "Title parameter is required"}), 400
34
-
35
- # Check if the film is already cached
36
- if title in load_balancer.FILM_STORE:
37
- url = load_balancer.FILM_STORE[title]
38
- return jsonify({"url":url})
39
-
40
- movie_path = load_balancer.find_movie_path(title)
41
-
42
- if not movie_path:
43
- return jsonify({"error": "Movie not found"}), 404
44
-
45
- # Start the download in a instance
46
- response = load_balancer.download_film_to_best_instance(title=title)
47
- if response:
48
- return jsonify(response)
49
-
50
- @app.route('/api/tv/<title>/<season>/<episode>', methods=['GET'])
51
- def get_tv_show_api(title, season, episode):
52
- """Endpoint to get the TV show by title, season, and episode."""
53
- if not title or not season or not episode:
54
- return jsonify({"error": "Title, season, and episode parameters are required"}), 400
55
-
56
- # Check if the episode is already cached
57
- if title in load_balancer.TV_STORE and season in load_balancer.TV_STORE[title]:
58
- for ep in load_balancer.TV_STORE[title][season]:
59
- if episode in ep:
60
- url = load_balancer.TV_STORE[title][season][ep]
61
- return jsonify({"url":url})
62
-
63
- tv_path = load_balancer.find_tv_path(title)
64
-
65
- if not tv_path:
66
- return jsonify({"error": "TV show not found"}), 404
67
-
68
- episode_path = None
69
- for directory in load_balancer.file_structure:
70
- if directory['type'] == 'directory' and directory['path'] == 'tv':
71
- for sub_directory in directory['contents']:
72
- if sub_directory['type'] == 'directory' and title.lower() in sub_directory['path'].lower():
73
- for season_dir in sub_directory['contents']:
74
- if season_dir['type'] == 'directory' and season in season_dir['path']:
75
- for episode_file in season_dir['contents']:
76
- if episode_file['type'] == 'file' and episode in episode_file['path']:
77
- episode_path = episode_file['path']
78
- break
79
-
80
- if not episode_path:
81
- return jsonify({"error": "Episode not found"}), 404
82
-
83
- # Start the download in a instance
84
- response = load_balancer.download_episode_to_best_instance(title=title, season=season, episode=episode)
85
- if response:
86
- return jsonify(response)
87
-
88
- @app.route('/api/filmid/<title>', methods=['GET'])
89
- def get_film_id_by_title_api(title):
90
- """Endpoint to get the film ID by providing the movie title."""
91
- if not title:
92
- return jsonify({"error": "Title parameter is required"}), 400
93
- film_id = load_balancer.get_film_id(title)
94
- return jsonify({"film_id": film_id})
95
-
96
- @app.route('/api/episodeid/<title>/<season>/<episode>', methods=['GET'])
97
- def get_episode_id_api(title,season,episode):
98
- """Endpoint to get the episode ID by providing the TV show title, season, and episode."""
99
- if not title or not season or not episode:
100
- return jsonify({"error": "Title, season, and episode parameters are required"}), 400
101
- episode_id = encode_episodeid(title,season,episode)
102
- return jsonify({"episode_id": episode_id})
103
-
104
- @app.route('/api/cache/size', methods=['GET'])
105
- def get_cache_size_api():
106
- total_size = 0
107
- for dirpath, dirnames, filenames in os.walk(CACHE_DIR):
108
- for f in filenames:
109
- fp = os.path.join(dirpath, f)
110
- total_size += os.path.getsize(fp)
111
- readable_size = bytes_to_human_readable(total_size)
112
- return jsonify({"cache_size": readable_size})
113
-
114
- @app.route('/api/cache/clear', methods=['POST'])
115
- def clear_cache_api():
116
- for dirpath, dirnames, filenames in os.walk(CACHE_DIR):
117
- for f in filenames:
118
- fp = os.path.join(dirpath, f)
119
- os.remove(fp)
120
- return jsonify({"status": "Cache cleared"})
121
-
122
- @app.route('/api/tv/store', methods=['GET'])
123
- def get_tv_store_api():
124
- """Endpoint to get the TV store JSON."""
125
- return jsonify(load_balancer.TV_STORE)
126
-
127
-
128
- @app.route('/api/film/store', methods=['GET'])
129
- def get_film_store_api():
130
- """Endpoint to get the film store JSON."""
131
- return jsonify(load_balancer.FILM_STORE)
132
-
133
-
134
- @app.route('/api/film/metadata/<title>', methods=['GET'])
135
- def get_film_metadata_api(title):
136
- """Endpoint to get the film metadata by title."""
137
- if not title:
138
- return jsonify({'error': 'No title provided'}), 400
139
-
140
- json_cache_path = os.path.join(CACHE_DIR, f"{urllib.parse.quote(title)}.json")
141
-
142
- if os.path.exists(json_cache_path):
143
- with open(json_cache_path, 'r') as f:
144
- data = json.load(f)
145
- return jsonify(data)
146
-
147
- return jsonify({'error': 'Metadata not found'}), 404
148
-
149
- @app.route('/api/tv/metadata/<title>', methods=['GET'])
150
- def get_tv_metadata_api(title):
151
- """Endpoint to get the TV show metadata by title."""
152
- if not title:
153
- return jsonify({'error': 'No title provided'}), 400
154
-
155
- json_cache_path = os.path.join(CACHE_DIR, f"{urllib.parse.quote(title)}.json")
156
-
157
- if os.path.exists(json_cache_path):
158
- with open(json_cache_path, 'r') as f:
159
- data = json.load(f)
160
-
161
- # Add the file structure to the metadata
162
- tv_structure_data = load_balancer.get_tv_structure(title)
163
- if tv_structure_data:
164
- data['file_structure'] = tv_structure_data
165
-
166
- return jsonify(data)
167
-
168
- return jsonify({'error': 'Metadata not found'}), 404
169
-
170
-
171
- @app.route("/api/film/all")
172
- def get_all_films_api():
173
- return load_balancer.get_all_films()
174
-
175
- @app.route("/api/tv/all")
176
- def get_all_tvshows_api():
177
- return load_balancer.get_all_tv_shows()
178
-
179
- @app.route('/api/instances',methods=["GET"])
180
- def get_instances():
181
- return load_balancer.instances
182
-
183
- @app.route('/api/instances/health',methods=["GET"])
184
- def get_instances_health():
185
- return load_balancer.instances_health
186
- #############################################################
187
- # This API is only for instances
188
- @app.route('/api/register', methods=['POST'])
189
- def register_instance():
190
- try:
191
- data = request.json
192
- if not data or "url" not in data:
193
- return jsonify({"error": "No URL provided"}), 400
194
-
195
- url = data["url"]
196
- if not is_valid_url(url):
197
- return jsonify({"error": "Invalid URL"}), 400
198
-
199
- # Register the instance
200
- load_balancer.register_instance(url)
201
- logging.info(f"Instance registered: {url}")
202
-
203
- return jsonify({"message": f"Instance {url} registered successfully"}), 200
204
-
205
- except Exception as e:
206
- logging.error(f"Error registering instance: {e}")
207
- return jsonify({"error": "Failed to register instance"}), 500
208
-
209
- #############################################################
210
- # Routes
211
- @app.route('/')
212
- def index():
213
- return f"Load Balancer is Running {load_balancer.version}"
214
-
215
- # Main entry point
216
  if __name__ == "__main__":
217
- app.run(debug=True, host="0.0.0.0", port=7860)
 
 
 
 
 
1
+ import subprocess
2
+ import sys
3
+
4
+ def run_uvicorn():
5
+ # Command to run the FastAPI app with uvicorn
6
+ command = [
7
+ "uvicorn",
8
+ "main:app",
9
+ "--host", "0.0.0.0",
10
+ "--port", "7860",
11
+ "--reload"
12
+ ]
13
+
14
+ # Run the command
15
+ subprocess.run(command, check=True)
16
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
17
  if __name__ == "__main__":
18
+ try:
19
+ run_uvicorn()
20
+ except subprocess.CalledProcessError as e:
21
+ print(f"Error running uvicorn: {e}", file=sys.stderr)
22
+ sys.exit(1)
main.py ADDED
@@ -0,0 +1,219 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI, HTTPException, Request
2
+ from fastapi.responses import JSONResponse
3
+ from fastapi.middleware.cors import CORSMiddleware
4
+ from utils import is_valid_url, bytes_to_human_readable, encode_episodeid
5
+ import os
6
+ import json
7
+ import urllib.parse
8
+ from threading import Thread
9
+ from LoadBalancer import LoadBalancer
10
+ import logging
11
+
12
+ app = FastAPI()
13
+
14
+ # CORS
15
+ app.add_middleware(
16
+ CORSMiddleware,
17
+ allow_origins=["*"],
18
+ allow_credentials=True,
19
+ allow_methods=["*"],
20
+ allow_headers=["*"],
21
+ )
22
+
23
+ logging.basicConfig(level=logging.INFO)
24
+
25
+ # Constants and Configuration
26
+ CACHE_DIR = os.getenv("CACHE_DIR")
27
+ INDEX_FILE = os.getenv("INDEX_FILE")
28
+ TOKEN = os.getenv("TOKEN")
29
+ REPO = os.getenv("REPO")
30
+
31
+ load_balancer = LoadBalancer(cache_dir=CACHE_DIR, index_file=INDEX_FILE, token=TOKEN, repo=REPO)
32
+
33
+ # Start polling in a separate thread
34
+ polling_thread = Thread(target=load_balancer.start_polling)
35
+ polling_thread.start()
36
+
37
+ # API Endpoints
38
+ @app.get("/api/film/{title}")
39
+ async def get_movie_api(title: str):
40
+ """Endpoint to get the movie by title."""
41
+ if not title:
42
+ raise HTTPException(status_code=400, detail="Title parameter is required")
43
+
44
+ # Check if the film is already cached
45
+ if title in load_balancer.FILM_STORE:
46
+ url = load_balancer.FILM_STORE[title]
47
+ return JSONResponse(content={"url": url})
48
+
49
+ movie_path = await load_balancer.find_movie_path(title)
50
+
51
+ if not movie_path:
52
+ raise HTTPException(status_code=404, detail="Movie not found")
53
+
54
+ # Start the download in an instance
55
+ response = await load_balancer.download_film_to_best_instance(title=title)
56
+ if response:
57
+ return JSONResponse(content=response)
58
+
59
+ @app.get("/api/tv/{title}/{season}/{episode}")
60
+ async def get_tv_show_api(title: str, season: str, episode: str):
61
+ """Endpoint to get the TV show by title, season, and episode."""
62
+ if not title or not season or not episode:
63
+ raise HTTPException(status_code=400, detail="Title, season, and episode parameters are required")
64
+
65
+ # Check if the episode is already cached
66
+ if title in load_balancer.TV_STORE and season in load_balancer.TV_STORE[title]:
67
+ for ep in load_balancer.TV_STORE[title][season]:
68
+ if episode in ep:
69
+ url = load_balancer.TV_STORE[title][season][ep]
70
+ return JSONResponse(content={"url": url})
71
+
72
+ tv_path = await load_balancer.find_tv_path(title)
73
+
74
+ if not tv_path:
75
+ raise HTTPException(status_code=404, detail="TV show not found")
76
+
77
+ episode_path = None
78
+ for directory in load_balancer.file_structure:
79
+ if directory['type'] == 'directory' and directory['path'] == 'tv':
80
+ for sub_directory in directory['contents']:
81
+ if sub_directory['type'] == 'directory' and title.lower() in sub_directory['path'].lower():
82
+ for season_dir in sub_directory['contents']:
83
+ if season_dir['type'] == 'directory' and season in season_dir['path']:
84
+ for episode_file in season_dir['contents']:
85
+ if episode_file['type'] == 'file' and episode in episode_file['path']:
86
+ episode_path = episode_file['path']
87
+ break
88
+
89
+ if not episode_path:
90
+ raise HTTPException(status_code=404, detail="Episode not found")
91
+
92
+ # Start the download in an instance
93
+ response = await load_balancer.download_episode_to_best_instance(title=title, season=season, episode=episode)
94
+ if response:
95
+ return JSONResponse(content=response)
96
+
97
+ @app.get("/api/filmid/{title}")
98
+ async def get_film_id_by_title_api(title: str):
99
+ """Endpoint to get the film ID by providing the movie title."""
100
+ if not title:
101
+ raise HTTPException(status_code=400, detail="Title parameter is required")
102
+ film_id = await load_balancer.get_film_id(title)
103
+ return JSONResponse(content={"film_id": film_id})
104
+
105
+ @app.get("/api/episodeid/{title}/{season}/{episode}")
106
+ async def get_episode_id_api(title: str, season: str, episode: str):
107
+ """Endpoint to get the episode ID by providing the TV show title, season, and episode."""
108
+ if not title or not season or not episode:
109
+ raise HTTPException(status_code=400, detail="Title, season, and episode parameters are required")
110
+ episode_id = encode_episodeid(title, season, episode)
111
+ return JSONResponse(content={"episode_id": episode_id})
112
+
113
+ @app.get("/api/cache/size")
114
+ async def get_cache_size_api():
115
+ total_size = 0
116
+ for dirpath, dirnames, filenames in os.walk(CACHE_DIR):
117
+ for f in filenames:
118
+ fp = os.path.join(dirpath, f)
119
+ total_size += os.path.getsize(fp)
120
+ readable_size = bytes_to_human_readable(total_size)
121
+ return JSONResponse(content={"cache_size": readable_size})
122
+
123
+ @app.post("/api/cache/clear")
124
+ async def clear_cache_api():
125
+ for dirpath, dirnames, filenames in os.walk(CACHE_DIR):
126
+ for f in filenames:
127
+ fp = os.path.join(dirpath, f)
128
+ os.remove(fp)
129
+ return JSONResponse(content={"status": "Cache cleared"})
130
+
131
+ @app.get("/api/tv/store")
132
+ async def get_tv_store_api():
133
+ """Endpoint to get the TV store JSON."""
134
+ return JSONResponse(content=load_balancer.TV_STORE)
135
+
136
+
137
+ @app.get("/api/film/store")
138
+ async def get_film_store_api():
139
+ """Endpoint to get the film store JSON."""
140
+ return JSONResponse(content=load_balancer.FILM_STORE)
141
+
142
+ @app.get("/api/film/metadata/{title}")
143
+ async def get_film_metadata_api(title: str):
144
+ """Endpoint to get the film metadata by title."""
145
+ if not title:
146
+ raise HTTPException(status_code=400, detail="No title provided")
147
+
148
+ json_cache_path = os.path.join(CACHE_DIR, f"{urllib.parse.quote(title)}.json")
149
+
150
+ if os.path.exists(json_cache_path):
151
+ with open(json_cache_path, 'r') as f:
152
+ data = json.load(f)
153
+ return JSONResponse(content=data)
154
+
155
+ raise HTTPException(status_code=404, detail="Metadata not found")
156
+
157
+ @app.get("/api/tv/metadata/{title}")
158
+ async def get_tv_metadata_api(title: str):
159
+ """Endpoint to get the TV show metadata by title."""
160
+ if not title:
161
+ raise HTTPException(status_code=400, detail="No title provided")
162
+
163
+ json_cache_path = os.path.join(CACHE_DIR, f"{urllib.parse.quote(title)}.json")
164
+
165
+ if os.path.exists(json_cache_path):
166
+ with open(json_cache_path, 'r') as f:
167
+ data = json.load(f)
168
+
169
+ # Add the file structure to the metadata
170
+ tv_structure_data = await load_balancer.get_tv_structure(title)
171
+ if tv_structure_data:
172
+ data['file_structure'] = tv_structure_data
173
+
174
+ return JSONResponse(content=data)
175
+
176
+ raise HTTPException(status_code=404, detail="Metadata not found")
177
+
178
+ @app.get("/api/film/all")
179
+ async def get_all_films_api():
180
+ return JSONResponse(content=await load_balancer.get_all_films())
181
+
182
+ @app.get("/api/tv/all")
183
+ async def get_all_tvshows_api():
184
+ return JSONResponse(content=await load_balancer.get_all_tv_shows())
185
+
186
+ @app.get("/api/instances")
187
+ async def get_instances():
188
+ return JSONResponse(content=load_balancer.instances)
189
+
190
+ @app.get("/api/instances/health")
191
+ async def get_instances_health():
192
+ return JSONResponse(content=load_balancer.instances_health)
193
+
194
+ # This API is only for instances
195
+ @app.post("/api/register")
196
+ async def register_instance(request: Request):
197
+ try:
198
+ data = await request.json()
199
+ if not data or "url" not in data:
200
+ raise HTTPException(status_code=400, detail="No URL provided")
201
+
202
+ url = data["url"]
203
+ if not is_valid_url(url):
204
+ raise HTTPException(status_code=400, detail="Invalid URL")
205
+
206
+ # Register the instance
207
+ load_balancer.register_instance(url)
208
+ logging.info(f"Instance registered: {url}")
209
+
210
+ return JSONResponse(content={"message": f"Instance {url} registered successfully"}, status_code=200)
211
+
212
+ except Exception as e:
213
+ logging.error(f"Error registering instance: {e}")
214
+ raise HTTPException(status_code=500, detail="Failed to register instance")
215
+
216
+ # Routes
217
+ @app.get("/")
218
+ async def index():
219
+ return f"Load Balancer is Running {load_balancer.version}"
requirements.txt CHANGED
@@ -1,5 +1,6 @@
1
- Flask
2
- Flask-Cors
 
3
  requests
4
  python-dotenv
5
  tqdm
 
1
+ fastapi
2
+ aiofiles
3
+ uvicorn
4
  requests
5
  python-dotenv
6
  tqdm