From 68666be34063d4663d6e147762385fa1883d507a Mon Sep 17 00:00:00 2001
From: Goldy <153996346+g0ldyy@users.noreply.github.com>
Date: Sun, 30 Jun 2024 00:00:00 +0200
Subject: [PATCH] speed goes brrrrr

---
 .env-sample           |  2 +-
 comet/api/stream.py   | 49 ++++++++++++++++++++++++++++++++++++-------
 comet/utils/models.py |  5 ++---
 3 files changed, 45 insertions(+), 11 deletions(-)

diff --git a/.env-sample b/.env-sample
index 22bb2bd..cfa668e 100644
--- a/.env-sample
+++ b/.env-sample
@@ -7,7 +7,7 @@ DEBRID_PROXY_URL=http://127.0.0.1:1080 # https://github.com/cmj2002/warp-docker
 INDEXER_MANAGER_TYPE=jackett # or prowlarr
 INDEXER_MANAGER_URL=http://127.0.0.1:9117
 INDEXER_MANAGER_API_KEY=XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
-INDEXER_MANAGER_TIMEOUT=30 # maximum time to obtain search results from indexer manager in seconds
+INDEXER_MANAGER_TIMEOUT=60 # maximum time to obtain search results from indexer manager in seconds
 INDEXER_MANAGER_INDEXERS='["EXAMPLE1_CHANGETHIS", "EXAMPLE2_CHANGETHIS"]'
 GET_TORRENT_TIMEOUT=5 # maximum time to obtain the torrent info hash in seconds
 CUSTOM_HEADER_HTML=None # only set it if you know what it is
\ No newline at end of file
diff --git a/comet/api/stream.py b/comet/api/stream.py
index 66ed7ce..0718e86 100644
--- a/comet/api/stream.py
+++ b/comet/api/stream.py
@@ -135,13 +135,27 @@ async def stream(request: Request, b64config: str, type: str, id: str):
         if len(torrentHashes) == 0:
             return {"streams": []}
 
-        getAvailability = await session.get(f"https://api.real-debrid.com/rest/1.0/torrents/instantAvailability/{'/'.join(torrentHashes)}", headers={
-            "Authorization": f"Bearer {config['debridApiKey']}"
-        })
+        # hashChunks = [torrentHashes[i:i + 5] for i in range(0, len(torrentHashes), 5)]
 
-        files = {}
+        # tasks = []
+        # for chunk in hashChunks:
+        #     tasks.append(session.get(f"https://api.real-debrid.com/rest/1.0/torrents/instantAvailability/{'/'.join(chunk)}", headers={
+        #         "Authorization": f"Bearer {config['debridApiKey']}"
+        #     }))
+
+        tasks = []
+        for hash in torrentHashes:
+            tasks.append(session.get(f"https://api.real-debrid.com/rest/1.0/torrents/instantAvailability/{hash}", headers={
+                "Authorization": f"Bearer {config['debridApiKey']}"
+            }))
 
-        availability = await getAvailability.json()
+        responses = await asyncio.gather(*tasks)
+
+        availability = {}
+        for response in responses:
+            availability.update(await response.json())
+
+        files = {}
         for hash, details in availability.items():
             if not "rd" in details:
                 continue
@@ -217,8 +231,18 @@ async def stream(request: Request, b64config: str, type: str, id: str):
             "streams": results
         }
 
-@streams.route("/{b64config}/playback/{hash}/{index}", methods=["HEAD", "GET"])
-async def stream(b64config: str, hash: str, index: str):
+# @streams.route("/{b64config}/playback/{hash}/{index}", methods=["HEAD", "GET"])
+# async def playback(b64config: str, hash: str, index: str):
+#     config = configChecking(b64config)
+#     if not config:
+#         return
+
+#     downloadLink = await generateDownloadLink(config["debridApiKey"], hash, index)
+
+#     return RedirectResponse(downloadLink, status_code=302)
+
+@streams.head("/{b64config}/playback/{hash}/{index}")
+async def playback(b64config: str, hash: str, index: str):
     config = configChecking(b64config)
     if not config:
         return
@@ -226,3 +250,14 @@ async def stream(b64config: str, hash: str, index: str):
     downloadLink = await generateDownloadLink(config["debridApiKey"], hash, index)
 
     return RedirectResponse(downloadLink, status_code=302)
+
+
+@streams.get("/{b64config}/playback/{hash}/{index}")
+async def playback(b64config: str, hash: str, index: str):
+    config = configChecking(b64config)
+    if not config:
+        return
+
+    downloadLink = await generateDownloadLink(config["debridApiKey"], hash, index)
+
+    return RedirectResponse(downloadLink, status_code=302)
\ No newline at end of file
diff --git a/comet/utils/models.py b/comet/utils/models.py
index c9c0283..b0a281b 100644
--- a/comet/utils/models.py
+++ b/comet/utils/models.py
@@ -1,15 +1,14 @@
 import os
-from typing import List, Optional
 
+from typing import List, Optional
 from databases import Database
 from pydantic_settings import BaseSettings, SettingsConfigDict
 from RTN import RTN, BaseRankingModel, SettingsModel
 
-
 class AppSettings(BaseSettings):
     model_config = SettingsConfigDict(
         env_file=".env",
-        env_file_encoding="utf-8",
+        env_file_encoding="utf-8"
     )
 
     FASTAPI_HOST: str = "0.0.0.0"