From 19fc235e838bcfb6d3d7e06f95309ff72c491fd5 Mon Sep 17 00:00:00 2001
From: Goldy <153996346+g0ldyy@users.noreply.github.com>
Date: Thu, 18 Jul 2024 15:19:42 +0200
Subject: [PATCH] fix: faster (we title match check before doing the rest)

---
 comet/api/stream.py    | 52 ++++++++++++++++++++++++++++++++----------
 comet/utils/general.py | 21 ++++++++++++++++-
 2 files changed, 60 insertions(+), 13 deletions(-)

diff --git a/comet/api/stream.py b/comet/api/stream.py
index aae9840..9065043 100644
--- a/comet/api/stream.py
+++ b/comet/api/stream.py
@@ -17,6 +17,7 @@
     get_debrid_extension,
     get_indexer_manager,
     get_zilean,
+    filter,
     get_torrent_hash,
     translate,
     get_balanced_hashes,
@@ -143,7 +144,9 @@ async def stream(request: Request, b64config: str, type: str, id: str):
                                     "torrentTitle": data["torrent_title"]
                                     if "torrent_title" in data
                                     else None,
-                                    "torrentSize": data["torrent_size"] if "torrent_size" in data else None,
+                                    "torrentSize": data["torrent_size"]
+                                    if "torrent_size" in data
+                                    else None,
                                     "url": f"{request.url.scheme}://{request.url.netloc}/{b64config}/playback/{hash}/{data['index']}",
                                 }
                             )
@@ -206,6 +209,31 @@ async def stream(request: Request, b64config: str, type: str, id: str):
             f"{len(torrents)} torrents found for {log_name} with {indexer_manager_type}{' and Zilean' if settings.ZILEAN_URL else ''}"
         )
 
+        if len(torrents) == 0:
+            return {"streams": []}
+
+        indexed_torrents = [(i, torrents[i]["Title"]) for i in range(len(torrents))]
+        chunk_size = 50
+        chunks = [
+            indexed_torrents[i : i + chunk_size]
+            for i in range(0, len(indexed_torrents), chunk_size)
+        ]
+
+        tasks = []
+        for chunk in chunks:
+            tasks.append(filter(chunk, name))
+
+        filtered_torrents = await asyncio.gather(*tasks)
+        index_less = 0
+        for result in filtered_torrents:
+            for filtered in result:
+                if not filtered[1]:
+                    del torrents[filtered[0] - index_less]
+                    index_less += 1
+                    continue
+
+        logger.info(f"{len(torrents)} torrents passed title match check for {log_name}")
+
         if len(torrents) == 0:
             return {"streams": []}
 
@@ -237,12 +265,13 @@ async def stream(request: Request, b64config: str, type: str, id: str):
 
         ranked_files = set()
         for hash in files:
-            try:
-                ranked_file = rtn.rank(
-                    files[hash]["title"], hash, correct_title=name, remove_trash=True
-                )
-            except:
-                continue
+            # try:
+            ranked_file = rtn.rank(
+                files[hash]["title"],
+                hash,  # , correct_title=name, remove_trash=True
+            )
+            # except:
+            #     continue
 
             ranked_files.add(ranked_file)
 
@@ -259,10 +288,7 @@ async def stream(request: Request, b64config: str, type: str, id: str):
             key: (value.model_dump() if isinstance(value, Torrent) else value)
             for key, value in sorted_ranked_files.items()
         }
-        torrents_by_hash = {
-            torrent["InfoHash"]: torrent
-            for torrent in torrents
-        }
+        torrents_by_hash = {torrent["InfoHash"]: torrent for torrent in torrents}
         for hash in sorted_ranked_files:  # needed for caching
             sorted_ranked_files[hash]["data"]["title"] = files[hash]["title"]
             sorted_ranked_files[hash]["data"]["torrent_title"] = torrents_by_hash[hash][
@@ -272,7 +298,9 @@ async def stream(request: Request, b64config: str, type: str, id: str):
                 "Tracker"
             ]
             sorted_ranked_files[hash]["data"]["size"] = files[hash]["size"]
-            sorted_ranked_files[hash]["data"]["torrent_size"] = torrents_by_hash[hash]["Size"]
+            sorted_ranked_files[hash]["data"]["torrent_size"] = torrents_by_hash[hash][
+                "Size"
+            ]
             sorted_ranked_files[hash]["data"]["index"] = files[hash]["index"]
 
         json_data = json.dumps(sorted_ranked_files).replace("'", "''")
diff --git a/comet/utils/general.py b/comet/utils/general.py
index 3e94b0c..5b23ab5 100644
--- a/comet/utils/general.py
+++ b/comet/utils/general.py
@@ -5,6 +5,8 @@
 import aiohttp
 import bencodepy
 
+from RTN import parse, title_match
+
 from comet.utils.logger import logger
 from comet.utils.models import settings, ConfigModel
 
@@ -247,7 +249,9 @@ async def get_indexer_manager(
             response = await response.json()
 
             for result in response:
-                result["InfoHash"] = result["infoHash"] if "infoHash" in result else None
+                result["InfoHash"] = (
+                    result["infoHash"] if "infoHash" in result else None
+                )
                 result["Title"] = result["title"]
                 result["Size"] = result["size"]
                 result["Link"] = result["downloadUrl"]
@@ -313,6 +317,21 @@ async def get_zilean(
     return results
 
 
+async def filter(torrents: list, name: str):
+    results = []
+    for torrent in torrents:
+        index = torrent[0]
+        torrent = torrent[1]
+
+        if title_match(name, parse(torrent).parsed_title):
+            results.append((index, True))
+            continue
+
+        results.append((index, False))
+
+    return results
+
+
 async def get_torrent_hash(session: aiohttp.ClientSession, torrent: tuple):
     index = torrent[0]
     torrent = torrent[1]