From 4af9addd7b05f69a16df317b1e5fbe2aa3e62d2b Mon Sep 17 00:00:00 2001 From: g0ldyy <153996346+g0ldyy@users.noreply.github.com> Date: Fri, 9 Jan 2026 01:48:20 +0100 Subject: [PATCH] feat: kitsu offsets --- comet/api/endpoints/stream.py | 24 ++++++- comet/core/database.py | 12 ++++ comet/services/anime.py | 123 ++++++++++++++++++++++++++++++-- comet/services/orchestration.py | 19 +++-- 4 files changed, 165 insertions(+), 13 deletions(-) diff --git a/comet/api/endpoints/stream.py b/comet/api/endpoints/stream.py index 49e3a40..a67dd8c 100644 --- a/comet/api/endpoints/stream.py +++ b/comet/api/endpoints/stream.py @@ -12,6 +12,7 @@ from comet.debrid.exceptions import DebridAuthError from comet.debrid.manager import get_debrid_extension from comet.metadata.filter import release_filter from comet.metadata.manager import MetadataScraper +from comet.services.anime import anime_mapper from comet.services.debrid import DebridService from comet.services.lock import DistributedLock, is_scrape_in_progress from comet.services.orchestration import TorrentManager @@ -318,6 +319,25 @@ async def stream( get_client_ip(request), ) + is_kitsu = media_id.startswith("kitsu:") + search_episode = episode + search_season = season + + if is_kitsu and episode is not None: + kitsu_mapping = anime_mapper.get_kitsu_episode_mapping(id) + if kitsu_mapping: + from_episode = kitsu_mapping.get("from_episode") + from_season = kitsu_mapping.get("from_season") + if from_episode: + search_episode = from_episode + episode - 1 + if from_season: + search_season = from_season + if from_episode or from_season: + logger.log( + "SCRAPER", + f"📺 Multi-part anime detected (kitsu:{id}): searching for S{search_season:02d}E{search_episode:02d} instead of S{season:02d}E{episode:02d}", + ) + torrent_manager = TorrentManager( debrid_service, config["debridApiKey"], @@ -332,7 +352,9 @@ async def stream( episode, aliases, settings.REMOVE_ADULT_CONTENT and config["removeTrash"], - is_kitsu=media_id.startswith("kitsu:"), + is_kitsu=is_kitsu, + search_episode=search_episode, + search_season=search_season, ) await torrent_manager.get_cached_torrents() diff --git a/comet/core/database.py b/comet/core/database.py index 5007294..306969e 100644 --- a/comet/core/database.py +++ b/comet/core/database.py @@ -370,6 +370,18 @@ async def setup_database(): """ ) + await database.execute( + """ + CREATE TABLE IF NOT EXISTS kitsu_imdb_mapping ( + kitsu_id TEXT PRIMARY KEY, + imdb_id TEXT, + title TEXT, + from_season INTEGER, + from_episode INTEGER + ) + """ + ) + await database.execute( """ CREATE TABLE IF NOT EXISTS digital_release_cache ( diff --git a/comet/services/anime.py b/comet/services/anime.py index b8cbf14..c2566e7 100644 --- a/comet/services/anime.py +++ b/comet/services/anime.py @@ -45,8 +45,11 @@ class AnimeMapper: self._refresh_task = None self.anime_imdb_ids = set() + self._kitsu_episode_cache = {} + self._aod_url = "https://github.com/manami-project/anime-offline-database/releases/latest/download/anime-offline-database-minified.json" self._fribb_url = "https://raw.githubusercontent.com/Fribb/anime-lists/refs/heads/master/anime-list-full.json" + self._kitsu_imdb_url = "https://raw.githubusercontent.com/TheBeastLT/stremio-kitsu-anime/master/static/data/imdb_mapping.json" async def load_anime_mapping(self, session: aiohttp.ClientSession | None = None): if not settings.ANIME_MAPPING_ENABLED: @@ -58,15 +61,24 @@ class AnimeMapper: count = await database.fetch_val("SELECT COUNT(*) FROM anime_entries") if count and count > 0: await self._load_provider_ids() + await self._load_kitsu_episode_cache() - if await self._is_cache_stale(): + kitsu_count = await database.fetch_val( + "SELECT COUNT(*) FROM kitsu_imdb_mapping" + ) + needs_kitsu_refresh = ( + kitsu_count == 0 or len(self._kitsu_episode_cache) == 0 + ) + + if await self._is_cache_stale() or needs_kitsu_refresh: self._refresh_task = asyncio.create_task( self._refresh_from_remote(background=True) ) self.loaded = True logger.log( - "COMET", f"✅ Anime mapping loaded from database: {count} entries" + "COMET", + f"✅ Anime mapping loaded from database: {count} entries, {len(self._kitsu_episode_cache)} Kitsu episode offsets", ) return True @@ -148,6 +160,12 @@ class AnimeMapper: return val + def get_kitsu_episode_mapping(self, kitsu_id: str | int): + if not self.loaded: + return None + + return self._kitsu_episode_cache.get(str(kitsu_id)) + def is_loaded(self): return self.loaded @@ -187,6 +205,29 @@ class AnimeMapper: except Exception as e: logger.error(f"Failed to load anime provider IDs: {e}") + async def _load_kitsu_episode_cache(self): + try: + rows = await database.fetch_all( + """ + SELECT kitsu_id, imdb_id, from_season, from_episode + FROM kitsu_imdb_mapping + WHERE (from_episode IS NOT NULL AND from_episode > 1) + OR from_season IS NOT NULL + """ + ) + + self._kitsu_episode_cache.clear() + for row in rows: + kitsu_id = row["kitsu_id"] + + self._kitsu_episode_cache[str(kitsu_id)] = { + "imdb_id": row["imdb_id"], + "from_season": row["from_season"], + "from_episode": row["from_episode"], + } + except Exception as e: + logger.warning(f"Failed to load Kitsu episode cache: {e}") + async def _refresh_from_remote( self, session: aiohttp.ClientSession | None = None, @@ -205,7 +246,7 @@ class AnimeMapper: try: logger.log( "COMET", - "Downloading anime mapping (Source 1/2: Anime Offline Database)...", + "Downloading anime mapping (Source 1/3: Anime Offline Database)...", ) async with session.get(self._aod_url) as response_aod: if response_aod.status != 200: @@ -215,7 +256,7 @@ class AnimeMapper: logger.log( "COMET", - "Downloading anime mapping (Source 2/2: Fribb Anime List)...", + "Downloading anime mapping (Source 2/3: Fribb Anime List)...", ) async with session.get(self._fribb_url) as response_fribb: if response_fribb.status != 200: @@ -225,13 +266,28 @@ class AnimeMapper: return False data_fribb = orjson.loads(await response_fribb.read()) + logger.log( + "COMET", + "Downloading anime mapping (Source 3/3: Kitsu-IMDB Episode Offsets)...", + ) + async with session.get(self._kitsu_imdb_url) as response_kitsu: + if response_kitsu.status != 200: + logger.warning( + f"Failed to load Kitsu-IMDB mapping: HTTP {response_kitsu.status}" + ) + return False + data_kitsu_imdb = orjson.loads(await response_kitsu.read()) + anime_list = data_aod.get("data", []) total_entries, total_fribb = await self._persist_mapping( anime_list, data_fribb ) + await self._persist_kitsu_imdb_mapping(data_kitsu_imdb) + del data_aod del data_fribb + del data_kitsu_imdb del anime_list gc.collect() @@ -249,11 +305,12 @@ class AnimeMapper: pass await self._load_provider_ids() + await self._load_kitsu_episode_cache() self.loaded = True logger.log( "COMET", - f"✅ Anime mapping loaded: {total_entries} entries", + f"✅ Anime mapping loaded: {total_entries} entries, {len(self._kitsu_episode_cache)} Kitsu episode offsets cached", ) return True @@ -389,5 +446,61 @@ class AnimeMapper: logger.error(f"Failed to persist anime mapping cache: {exc}") return 0, 0 + async def _persist_kitsu_imdb_mapping(self, kitsu_imdb_data: dict): + total_count = 0 + batch = [] + batch_size = 1000 + + try: + async with database.transaction(): + await database.execute("DELETE FROM kitsu_imdb_mapping") + + insert_query = """ + INSERT INTO kitsu_imdb_mapping + (kitsu_id, imdb_id, title, from_season, from_episode) + VALUES (:kitsu_id, :imdb_id, :title, :from_season, :from_episode) + ON CONFLICT (kitsu_id) DO UPDATE SET + imdb_id = :imdb_id, + title = :title, + from_season = :from_season, + from_episode = :from_episode + """ + + for kitsu_id, entry in kitsu_imdb_data.items(): + if not isinstance(entry, dict): + continue + + imdb_id = entry.get("imdb_id") + if not imdb_id: + continue + + from_season = entry.get("fromSeason") + from_episode = entry.get("fromEpisode") + + batch.append( + { + "kitsu_id": str(kitsu_id), + "imdb_id": imdb_id, + "title": entry.get("title"), + "from_season": from_season, + "from_episode": from_episode, + } + ) + + if len(batch) >= batch_size: + await database.execute_many(insert_query, batch) + total_count += len(batch) + batch.clear() + + if batch: + await database.execute_many(insert_query, batch) + total_count += len(batch) + batch.clear() + + return total_count + except Exception as exc: + logger.error(f"Failed to persist Kitsu-IMDB mapping: {exc}") + return 0 + anime_mapper = AnimeMapper() diff --git a/comet/services/orchestration.py b/comet/services/orchestration.py index 4fb079e..6c4bfe2 100644 --- a/comet/services/orchestration.py +++ b/comet/services/orchestration.py @@ -28,8 +28,10 @@ class TorrentManager: episode: int, aliases: dict, remove_adult_content: bool, - is_kitsu: bool = False, # Kitsu treats each season as separate anime - context: str = "live", # "live" or "background" + is_kitsu: bool = False, + context: str = "live", + search_episode: int | None = None, + search_season: int | None = None, ): self.debrid_service = debrid_service self.debrid_api_key = debrid_api_key @@ -42,6 +44,8 @@ class TorrentManager: self.year_end = year_end self.season = season self.episode = episode + self.search_episode = search_episode if search_episode is not None else episode + self.search_season = search_season if search_season is not None else season self.aliases = aliases self.remove_adult_content = remove_adult_content self.is_kitsu = is_kitsu @@ -64,8 +68,8 @@ class TorrentManager: title=self.title, year=self.year, year_end=self.year_end, - season=self.season, - episode=self.episode, + season=self.search_season, + episode=self.search_episode, context=self.context, ) @@ -81,7 +85,7 @@ class TorrentManager: ) if self.is_kitsu: - if episode is not None and episode != self.episode: + if episode is not None and episode != self.search_episode: continue else: if (season is not None and season != self.season) or ( @@ -111,7 +115,7 @@ class TorrentManager: """, { "media_id": self.media_only_id, - "episode": self.episode, + "episode": self.search_episode, }, ) else: @@ -136,7 +140,8 @@ class TorrentManager: parsed_data = ParsedData(**orjson.loads(row["parsed"])) if row["episode"] is None and parsed_data.episodes: - if self.episode not in parsed_data.episodes: + target_episode = self.search_episode if self.is_kitsu else self.episode + if target_episode not in parsed_data.episodes: continue info_hash = row["info_hash"]