From f90d2d2f0447d2af1e9f554bb9f30357e287b37a Mon Sep 17 00:00:00 2001
From: Zoe Roux <zoe.roux@zoriya.dev>
Date: Fri, 2 Feb 2024 16:54:41 +0100
Subject: [PATCH] Add multiple season rules

---
 scanner/scanner/parser/guess.py |  5 ++-
 scanner/scanner/parser/rules.py | 55 +++++++++++++++++++++++++--------
 2 files changed, 46 insertions(+), 14 deletions(-)

diff --git a/scanner/scanner/parser/guess.py b/scanner/scanner/parser/guess.py
index e13442bf..d594bc08 100644
--- a/scanner/scanner/parser/guess.py
+++ b/scanner/scanner/parser/guess.py
@@ -4,7 +4,10 @@ from guessit.api import default_api
 from typing import cast
 from rebulk import Rebulk
 
-import rules
+try:
+	from . import rules
+except:
+	import rules
 
 default_api.configure({})
 rblk = cast(Rebulk, default_api.rebulk)
diff --git a/scanner/scanner/parser/rules.py b/scanner/scanner/parser/rules.py
index 3f9eb94f..62212468 100644
--- a/scanner/scanner/parser/rules.py
+++ b/scanner/scanner/parser/rules.py
@@ -1,8 +1,9 @@
 # Read that for examples/rules: https://github.com/pymedusa/Medusa/blob/master/medusa/name_parser/rules/rules.py
 
-from typing import Any
+from typing import Any, List, Optional, cast
 from rebulk import Rule, RemoveMatch, AppendMatch, POST_PROCESS
-from rebulk.match import Matches
+from rebulk.match import Matches, Match
+from copy import copy
 
 
 class MultipleSeasonRule(Rule):
@@ -35,8 +36,8 @@ class MultipleSeasonRule(Rule):
 	Expected:
 	```json
 	{
-	        "title": "Spy x Family Season 2",
-	        "season": null,
+	        "title": "Spy x Family",
+	        "season": 2,
 	        "episode": 8,
 	        "screen_size": "1080p",
 	        "release_group": "Multiple Subtitle",
@@ -46,18 +47,46 @@ class MultipleSeasonRule(Rule):
 	        "type": "episode"
 	}
 	```
-
-	We want `Season 2 ` to be parsed as part of the title since this format is
-	often used for animes (where season often does not match, we use thexem for that)
 	"""
 
 	priority = POST_PROCESS
 	consequence = [RemoveMatch, AppendMatch]
 
 	def when(self, matches: Matches, context) -> Any:
-		seasons = matches.named("season")
-		print(seasons)
-		print(seasons[0])
-		print(vars(seasons[0]))
-		print(seasons[0].initiator)
-		return
+		seasons: List[Match] = matches.named("season")  # type: ignore
+
+		if not seasons or len(seasons) < 1:
+			return
+
+		# Only apply this rule if all seasons are due to the same match
+		initiator: Optional[Match] = seasons[0].initiator
+		if not initiator or any(
+			True for match in seasons if match.initiator != initiator
+		):
+			return
+
+		value: str = initiator.value  # type: ignore
+		if "-" not in value:
+			return
+
+		newSeason, *newEpisodes = (x.strip() for x in value.split("-"))
+		to_remove = [x for x in seasons if cast(Match, x.parent).value != newSeason]
+		to_add = []
+
+		try:
+			episodes = [int(x) for x in newEpisodes]
+			parents: List[Match] = [match.parent for match in to_remove]  # type: ignore
+			for episode in episodes:
+				smatch = next(
+					x
+					for x in parents
+					if int(str(x.value).replace("-", "").strip()) == episode
+				)
+				match = copy(smatch)
+				match.name = "episode"
+				match.value = episode
+				to_add.append(match)
+		except (ValueError, StopIteration):
+			return
+
+		return [to_remove, to_add]