diff --git a/comet/core/log_levels.py b/comet/core/log_levels.py index 9fe25dc..a7802ac 100644 --- a/comet/core/log_levels.py +++ b/comet/core/log_levels.py @@ -68,6 +68,12 @@ CUSTOM_LOG_LEVELS = { "loguru_color": "", "no": 35, }, + "PLAYBACK": { + "color": "#00FA9A", + "icon": "▶️", + "loguru_color": "", + "no": 35, + }, } ALL_LOG_LEVELS = {**STANDARD_LOG_LEVELS, **CUSTOM_LOG_LEVELS} diff --git a/comet/debrid/stremthru.py b/comet/debrid/stremthru.py index 561fac2..780ae18 100644 --- a/comet/debrid/stremthru.py +++ b/comet/debrid/stremthru.py @@ -2,13 +2,14 @@ import asyncio from urllib.parse import quote, unquote import aiohttp -from RTN import parse, title_match +from RTN import normalize_title, parse, title_match from comet.core.execution import get_executor from comet.core.logger import logger from comet.core.models import settings from comet.debrid.exceptions import DebridAuthError from comet.services.debrid_cache import cache_availability +from comet.services.filtering import quick_alias_match from comet.services.torrent_manager import torrent_update_queue from comet.utils.parsing import is_video @@ -217,6 +218,18 @@ class StremThru: sources: list = None, aliases: dict = None, ): + """ + Smart file selection algorithm with scoring system. + + Priority order (highest to lowest): + 1. Exact season + episode match with single episode file (+1000) + 2. Exact season + episode match with multi-episode file (+500) + 3. Episode match without season info (+200) + 4. Exact filename match with requested torrent_name (+100) + 5. Title alias match (+50) + 6. Index match from original selection (+25) + 7. Fallback: largest video file (+file_size as tiebreaker) + """ try: magnet_uri = f"magnet:?xt=urn:btih:{hash}&dn={quote(torrent_name)}" @@ -236,82 +249,170 @@ class StremThru: name = unquote(name) torrent_name = unquote(torrent_name) - name_parsed = parse(name) - target_file = None + ez_aliases = aliases.get("ez", []) + if ez_aliases: + ez_aliases_normalized = [normalize_title(a) for a in ez_aliases] debrid_files = magnet["data"]["files"] - debrid_files_parsed = [] - files = [] + + # Filter to video files only, excluding samples + video_files = [] + filenames_to_parse = [] for file in debrid_files: filename = file["name"] + filename_lower = filename.lower() - if "sample" in filename.lower(): + if "sample" in filename_lower: continue - if not is_video(filename): continue - filename_parsed = parse(filename) + video_files.append(file) + filenames_to_parse.append(filename) - file_season = ( - filename_parsed.seasons[0] if filename_parsed.seasons else None - ) - file_episode = ( - filename_parsed.episodes[0] if filename_parsed.episodes else None - ) - file_index = file["index"] if file["index"] != -1 else None - file_size = file["size"] if file["size"] != -1 else None + if not video_files: + logger.warning(f"No video files found in torrent {hash}") + return - file = { - "index": file_index, - "title": filename, - "size": file_size, - "season": file_season, - "episode": file_episode, - "link": file.get("link", None), - } + loop = asyncio.get_running_loop() + parsed_results = await loop.run_in_executor( + get_executor(), batch_parse, filenames_to_parse + ) - debrid_files_parsed.append(file) + scored_files = [] - if not filename_parsed.parsed_title or not title_match( - name_parsed.parsed_title, - filename_parsed.parsed_title, - aliases=aliases, - ): + for file, filename, parsed in zip( + video_files, filenames_to_parse, parsed_results + ): + file_index = file["index"] if file.get("index", -1) != -1 else None + file_size = file["size"] if file.get("size", -1) != -1 else 0 + file_link = file.get("link") + + if not file_link: continue - file["info_hash"] = hash - file["parsed"] = filename_parsed - files.append(file) + file_season = parsed.seasons[0] if parsed.seasons else None + file_episode = parsed.episodes[0] if parsed.episodes else None - for file in debrid_files_parsed: - if file["title"] == torrent_name: - target_file = file - break + # Calculate score + score = 0 + match_reason = [] - if season == file["season"] and episode == file["episode"]: - target_file = file - break + # Season + Episode matching (highest priority) + if season is not None and episode is not None: + season_matches = (not parsed.seasons) or (season in parsed.seasons) + episode_matches = parsed.episodes and episode in parsed.episodes - if not target_file: - for file in files: - if str(file["index"]) == index: - target_file = file - break + if season_matches and episode_matches: + if len(parsed.episodes) == 1: + score += 1000 # Perfect single episode match + match_reason.append("exact_episode") + else: + score += 500 # Multi-episode file containing our episode + match_reason.append("multi_episode") + elif episode_matches: + score += 200 # Episode matches but season doesn't + match_reason.append("episode_only") - if len(files) > 0: - asyncio.create_task(cache_availability(self.real_debrid_name, files)) + # Exact filename match + if filename == torrent_name: + score += 100 + match_reason.append("exact_name") - if not target_file and len(debrid_files) > 0: - files_with_link = [ - file for file in debrid_files if file.get("link", None) - ] - if len(files_with_link) > 0: - target_file = max(files_with_link, key=lambda x: x["size"]) + # Title/alias matching + if parsed.parsed_title: + # Quick alias match first + if ez_aliases and quick_alias_match( + normalize_title(filename), ez_aliases_normalized + ): + score += 50 + match_reason.append("alias") + elif title_match(name, parsed.parsed_title, aliases=aliases): + score += 50 + match_reason.append("title") - if not target_file: + # Index match from original selection + if file_index is not None and str(file_index) == str(index): + score += 25 + match_reason.append("index") + + # Use file size as tiebreaker (larger files preferred) + # Normalize to 0-10 range to not overwhelm other scores + size_score = min( + file_size / (10 * 1024 * 1024 * 1024), 10 + ) # Cap at 10GB + score += size_score + + enriched_file = { + "index": file_index, + "title": filename, + "size": file_size if file_size > 0 else None, + "season": file_season, + "episode": file_episode, + "link": file_link, + "parsed": parsed, + "score": score, + "match_reason": match_reason, + } + + scored_files.append(enriched_file) + + if not scored_files: + logger.log( + "PLAYBACK", + f"No valid video files with links found in torrent {hash}", + ) return + # Sort by score descending + scored_files.sort(key=lambda x: x["score"], reverse=True) + + # Select best file + target_file = scored_files[0] + + logger.log( + "PLAYBACK", + f"File selection for {hash}: selected '{target_file['title']}' " + f"(score={target_file['score']:.1f}, reasons={target_file['match_reason']}) " + f"from {len(scored_files)} candidates", + ) + + all_files_for_cache = [] + + for f in scored_files: + if f["season"] is not None or f["episode"] is not None: + all_files_for_cache.append( + { + "info_hash": hash, + "index": f["index"], + "title": f["title"], + "size": f["size"], + "season": f["season"] or season, + "episode": f["episode"], + "parsed": f["parsed"], + } + ) + + # Also ensure the selected file is cached with the REQUESTED season/episode + # This handles cases where filename doesn't contain S/E info but user requested it + if season is not None or episode is not None: + all_files_for_cache.append( + { + "info_hash": hash, + "index": target_file["index"], + "title": target_file["title"], + "size": target_file["size"], + "season": season, + "episode": episode, + "parsed": target_file["parsed"], + } + ) + + if all_files_for_cache: + asyncio.create_task( + cache_availability(self.real_debrid_name, all_files_for_cache) + ) + link = await self.session.post( f"{self.base_url}/link/generate?client_ip={self.client_ip}", json={"link": target_file["link"]}, diff --git a/comet/services/debrid.py b/comet/services/debrid.py index 7c04810..b460228 100644 --- a/comet/services/debrid.py +++ b/comet/services/debrid.py @@ -80,7 +80,7 @@ class DebridService: for hash in info_hashes: torrents[hash]["cached"] = False - if self.debrid_service == "torrent" or len(torrents) == 0: + if len(torrents) == 0: return rows = await get_cached_availability( @@ -91,13 +91,24 @@ class DebridService: info_hash = row["info_hash"] torrents[info_hash]["cached"] = True - if row["parsed"] is not None: - torrents[info_hash]["parsed"] = ParsedData( - **orjson.loads(row["parsed"]) - ) if row["file_index"] is not None: - torrents[info_hash]["fileIndex"] = row["file_index"] - if row["title"] is not None: - torrents[info_hash]["title"] = row["title"] + try: + torrents[info_hash]["fileIndex"] = int(row["file_index"]) + except ValueError: + pass + if row["size"] is not None: torrents[info_hash]["size"] = row["size"] + + # Only update title/parsed if the cached file has resolution info + # Otherwise keep the original torrent info which may have better quality data + # E.g. torrent "[Group] Show S01 1080p" vs file "Show - 02.mkv" + if row["parsed"] is not None: + cached_parsed = ParsedData(**orjson.loads(row["parsed"])) + if ( + cached_parsed.resolution != "unknown" + or torrents[info_hash]["parsed"].resolution == "unknown" + ): + torrents[info_hash]["parsed"] = cached_parsed + if row["title"] is not None: + torrents[info_hash]["title"] = row["title"] diff --git a/comet/services/debrid_cache.py b/comet/services/debrid_cache.py index 34fa235..1796cb6 100644 --- a/comet/services/debrid_cache.py +++ b/comet/services/debrid_cache.py @@ -116,31 +116,43 @@ async def cache_availability(debrid_service: str, availability: list): async def get_cached_availability( debrid_service: str, info_hashes: list, season: int = None, episode: int = None ): - base_query = f""" - SELECT info_hash, file_index, title, size, parsed + select_clause = "SELECT info_hash, file_index, title, size, parsed" + if debrid_service == "torrent" and settings.DATABASE_TYPE == "postgresql": + select_clause = ( + "SELECT DISTINCT ON (info_hash) info_hash, file_index, title, size, parsed" + ) + + base_from_where = f""" FROM debrid_availability WHERE info_hash IN (SELECT CAST(value as TEXT) FROM {"json_array_elements_text" if settings.DATABASE_TYPE == "postgresql" else "json_each"}(:info_hashes)) - AND debrid_service = :debrid_service AND timestamp + :cache_ttl >= :current_time """ params = { "info_hashes": orjson.dumps(info_hashes).decode("utf-8"), - "debrid_service": debrid_service, "cache_ttl": settings.DEBRID_CACHE_TTL, "current_time": time.time(), "season": season, "episode": episode, } + if debrid_service != "torrent": + base_from_where += " AND debrid_service = :debrid_service" + params["debrid_service"] = debrid_service + + group_by_clause = "" + if debrid_service == "torrent" and settings.DATABASE_TYPE == "sqlite": + group_by_clause = " GROUP BY info_hash" + if debrid_service == "offcloud": - query = ( - base_query - + """ + season_episode_filter = """ AND ((CAST(:season as INTEGER) IS NULL AND season IS NULL) OR season = CAST(:season as INTEGER)) AND ((CAST(:episode as INTEGER) IS NULL AND episode IS NULL) OR episode = CAST(:episode as INTEGER)) """ + query = ( + select_clause + base_from_where + season_episode_filter + group_by_clause ) + results = await database.fetch_all(query, params) found_hashes = {r["info_hash"] for r in results} @@ -149,20 +161,26 @@ async def get_cached_availability( if remaining_hashes: null_title_params = { "info_hashes": orjson.dumps(remaining_hashes).decode("utf-8"), - "debrid_service": debrid_service, "cache_ttl": settings.DEBRID_CACHE_TTL, "current_time": time.time(), } - null_title_query = base_query + " AND title IS NULL" + if debrid_service != "torrent": + null_title_params["debrid_service"] = debrid_service + + null_title_query = ( + select_clause + base_from_where + " AND title IS NULL" + group_by_clause + ) null_results = await database.fetch_all(null_title_query, null_title_params) results.extend(null_results) else: query = ( - base_query + select_clause + + base_from_where + """ AND ((CAST(:season as INTEGER) IS NULL AND season IS NULL) OR season = CAST(:season as INTEGER)) AND ((CAST(:episode as INTEGER) IS NULL AND episode IS NULL) OR episode = CAST(:episode as INTEGER)) """ + + group_by_clause ) results = await database.fetch_all(query, params) diff --git a/comet/services/filtering.py b/comet/services/filtering.py index 95f0737..7aca56c 100644 --- a/comet/services/filtering.py +++ b/comet/services/filtering.py @@ -1,4 +1,4 @@ -from RTN import parse, title_match +from RTN import normalize_title, parse, title_match from comet.core.logger import logger from comet.core.models import settings @@ -13,12 +13,22 @@ else: pass +def quick_alias_match(text_normalized: str, ez_aliases_normalized: list[str]): + return any(alias in text_normalized for alias in ez_aliases_normalized) + + def filter_worker(torrents, title, year, year_end, aliases, remove_adult_content): results = [] + ez_aliases = aliases.get("ez", []) + if ez_aliases: + ez_aliases_normalized = [normalize_title(a) for a in ez_aliases] + for torrent in torrents: torrent_title = torrent["title"] - if "sample" in torrent_title.lower() or torrent_title == "": + torrent_title_lower = torrent_title.lower() + + if "sample" in torrent_title_lower or torrent_title == "": _log_exclusion(f"🚫 Rejected (Sample/Empty) | {torrent_title}") continue @@ -28,14 +38,20 @@ def filter_worker(torrents, title, year, year_end, aliases, remove_adult_content _log_exclusion(f"🔞 Rejected (Adult) | {torrent_title}") continue - if not parsed.parsed_title or not title_match( - title, parsed.parsed_title, aliases=aliases - ): - _log_exclusion( - f"❌ Rejected (Title Mismatch) | {torrent_title} | Parsed: {parsed.parsed_title} | Expected: {title}" - ) + if not parsed.parsed_title: + _log_exclusion(f"❌ Rejected (No Parsed Title) | {torrent_title}") continue + alias_matched = ez_aliases and quick_alias_match( + normalize_title(torrent_title), ez_aliases_normalized + ) + if not alias_matched: + if not title_match(title, parsed.parsed_title, aliases=aliases): + _log_exclusion( + f"❌ Rejected (Title Mismatch) | {torrent_title} | Parsed: {parsed.parsed_title} | Expected: {title}" + ) + continue + if year and parsed.year: if year_end is not None: if not (year <= parsed.year <= year_end): diff --git a/comet/services/orchestration.py b/comet/services/orchestration.py index b342363..4fb079e 100644 --- a/comet/services/orchestration.py +++ b/comet/services/orchestration.py @@ -104,7 +104,7 @@ class TorrentManager: if self.is_kitsu: rows = await database.fetch_all( """ - SELECT info_hash, file_index, title, seeders, size, tracker, sources, parsed + SELECT info_hash, file_index, title, seeders, size, tracker, sources, parsed, episode FROM torrents WHERE media_id = :media_id AND (episode IS NULL OR episode = CAST(:episode as INTEGER)) @@ -117,7 +117,7 @@ class TorrentManager: else: rows = await database.fetch_all( """ - SELECT info_hash, file_index, title, seeders, size, tracker, sources, parsed + SELECT info_hash, file_index, title, seeders, size, tracker, sources, parsed, episode FROM torrents WHERE media_id = :media_id AND ((season IS NOT NULL AND season = CAST(:season as INTEGER)) OR (season IS NULL AND CAST(:season as INTEGER) IS NULL)) @@ -130,7 +130,15 @@ class TorrentManager: }, ) + rows = sorted(rows, key=lambda r: (r["episode"] is not None, r["episode"])) + for row in rows: + parsed_data = ParsedData(**orjson.loads(row["parsed"])) + + if row["episode"] is None and parsed_data.episodes: + if self.episode not in parsed_data.episodes: + continue + info_hash = row["info_hash"] self.torrents[info_hash] = { "fileIndex": row["file_index"], @@ -139,35 +147,45 @@ class TorrentManager: "size": row["size"], "tracker": row["tracker"], "sources": orjson.loads(row["sources"]), - "parsed": ParsedData(**orjson.loads(row["parsed"])), + "parsed": parsed_data, } async def cache_torrents(self): for torrent in self.ready_to_cache: if self.is_kitsu: - cache_season = self.season + cache_seasons = [self.season] else: - cache_season = ( - torrent["parsed"].seasons[0] + cache_seasons = ( + torrent["parsed"].seasons if torrent["parsed"].seasons - else self.season + else [self.season] ) - file_info = { - "info_hash": torrent["infoHash"], - "index": torrent["fileIndex"], - "title": torrent["title"], - "size": torrent["size"], - "season": cache_season, - "episode": torrent["parsed"].episodes[0] - if torrent["parsed"].episodes - else None, - "parsed": torrent["parsed"], - "seeders": torrent["seeders"], - "tracker": torrent["tracker"], - "sources": torrent["sources"], - } - await torrent_update_queue.add_torrent_info(file_info, self.media_only_id) + parsed_episodes = ( + torrent["parsed"].episodes if torrent["parsed"].episodes else [None] + ) + + if len(parsed_episodes) > 1: + episode = None + else: + episode = parsed_episodes[0] + + for season in cache_seasons: + file_info = { + "info_hash": torrent["infoHash"], + "index": torrent["fileIndex"], + "title": torrent["title"], + "size": torrent["size"], + "season": season, + "episode": episode, + "parsed": torrent["parsed"], + "seeders": torrent["seeders"], + "tracker": torrent["tracker"], + "sources": torrent["sources"], + } + await torrent_update_queue.add_torrent_info( + file_info, self.media_only_id + ) async def filter_manager(self, scraper_name: str, torrents: list): if len(torrents) == 0: diff --git a/comet/services/torrent_manager.py b/comet/services/torrent_manager.py index e934024..386b26e 100644 --- a/comet/services/torrent_manager.py +++ b/comet/services/torrent_manager.py @@ -125,48 +125,53 @@ async def add_torrent( parsed: ParsedData, ): try: - parsed_season = parsed.seasons[0] if parsed.seasons else search_season - parsed_episode = parsed.episodes[0] if parsed.episodes else None + seasons_to_process = parsed.seasons if parsed.seasons else [search_season] + parsed_episodes = parsed.episodes if parsed.episodes else [None] - if parsed_episode is not None: - await database.execute( - """ - DELETE FROM torrents - WHERE info_hash = :info_hash - AND season = :season - AND episode IS NULL - """, + episode_to_insert = parsed_episodes[0] if len(parsed_episodes) == 1 else None + + for season in seasons_to_process: + # Only delete season-only entry if we are inserting a SINGLE specific episode. + if episode_to_insert is not None: + await database.execute( + """ + DELETE FROM torrents + WHERE info_hash = :info_hash + AND season = :season + AND episode IS NULL + """, + { + "info_hash": info_hash, + "season": season, + }, + ) + logger.log( + "SCRAPER", + f"Deleted season-only entry for S{season:02d} of {info_hash} in favor of specific episode", + ) + + await _upsert_torrent_record( { + "media_id": media_id, "info_hash": info_hash, - "season": parsed_season, - }, + "file_index": file_index, + "season": season, + "episode": episode_to_insert, + "title": title, + "seeders": seeders, + "size": size, + "tracker": tracker, + "sources": orjson.dumps(sources).decode("utf-8"), + "parsed": orjson.dumps(parsed, default_dump).decode("utf-8"), + "timestamp": time.time(), + } ) - logger.log( - "SCRAPER", - f"Deleted season-only entry for S{parsed_season:02d} of {info_hash}", - ) - - await _upsert_torrent_record( - { - "media_id": media_id, - "info_hash": info_hash, - "file_index": file_index, - "season": parsed_season, - "episode": parsed_episode, - "title": title, - "seeders": seeders, - "size": size, - "tracker": tracker, - "sources": orjson.dumps(sources).decode("utf-8"), - "parsed": orjson.dumps(parsed, default_dump).decode("utf-8"), - "timestamp": time.time(), - } - ) additional = "" - if parsed_season: - additional += f" - S{parsed_season:02d}" - additional += f"E{parsed_episode:02d}" if parsed_episode else "" + if seasons_to_process: + additional += f" - S{seasons_to_process[0]:02d}" + if parsed_episodes and parsed_episodes[0] is not None: + additional += f"E{parsed_episodes[0]:02d}" logger.log("SCRAPER", f"Added torrent for {media_id} - {title}{additional}") except Exception as e: