1# -*- coding: utf-8 -*-
2import io
3import logging
4import os
5import zipfile
6
7import rarfile
8from requests import Session
9from guessit import guessit
10from subliminal import Episode, Movie
11from subliminal.exceptions import ServiceUnavailable
12from subliminal.subtitle import SUBTITLE_EXTENSIONS, fix_line_ending
13from subliminal_patch.exceptions import APIThrottled
14from subliminal_patch.providers import Provider
15from subliminal_patch.subtitle import Subtitle, guess_matches
16from subzero.language import Language
17
18logger = logging.getLogger(__name__)
19
20SERVER_URL = "http://sapidb.caretas.club"
21PAGE_URL = "https://sucha.caretas.club"
22UNDESIRED_FILES = ("[eng]", ".en.", ".eng.", ".fr.", ".pt.")
23
24
25class SuchaSubtitle(Subtitle):
26    provider_name = "sucha"
27    hash_verifiable = False
28
29    def __init__(
30        self,
31        language,
32        release_info,
33        filename,
34        download_id,
35        download_type,
36        matches,
37    ):
38        super(SuchaSubtitle, self).__init__(
39            language, hearing_impaired=False, page_link=PAGE_URL
40        )
41        self.download_id = download_id
42        self.download_type = download_type
43        self.language = language
44        self.guessed_release_info = release_info
45        self.filename = filename
46        self.release_info = (
47            release_info if len(release_info) > len(filename) else filename
48        )
49        self.found_matches = matches
50
51    @property
52    def id(self):
53        return self.download_id
54
55    def get_matches(self, video):
56        type_ = "episode" if isinstance(video, Episode) else "movie"
57        self.found_matches |= guess_matches(
58            video,
59            guessit(self.filename, {"type": type_}),
60        )
61        self.found_matches |= guess_matches(
62            video,
63            guessit(self.guessed_release_info, {"type": type_}),
64        )
65        return self.found_matches
66
67
68class SuchaProvider(Provider):
69    """Sucha Provider"""
70
71    # This is temporary. Castilian spanish subtitles may exist, but are rare
72    # and currently impossible to guess from the API.
73    languages = {Language("spa", "MX")}
74    language_list = list(languages)
75    video_types = (Episode, Movie)
76
77    def initialize(self):
78        self.session = Session()
79        self.session.headers.update(
80            {"User-Agent": os.environ.get("SZ_USER_AGENT", "Sub-Zero/2")}
81        )
82
83    def terminate(self):
84        self.session.close()
85
86    def query(self, languages, video):
87        movie_year = video.year or "0"
88        is_episode = isinstance(video, Episode)
89        type_str = "episode" if is_episode else "movie"
90        language = self.language_list[0]
91
92        if is_episode:
93            q = {"query": f"{video.series} S{video.season:02}E{video.episode:02}"}
94        else:
95            q = {"query": video.title, "year": movie_year}
96
97        logger.debug(f"Searching subtitles: {q}")
98        result = self.session.get(f"{SERVER_URL}/{type_str}", params=q, timeout=10)
99        result.raise_for_status()
100
101        results = result.json()
102        if isinstance(results, dict):
103            logger.debug("No subtitles found")
104            return []
105
106        subtitles = []
107        for item in results:
108            matches = set()
109            title = item.get("title", "").lower()
110            alt_title = item.get("alt_title", title).lower()
111
112            if any(video.title.lower() in item for item in (title, alt_title)):
113                matches.add("title")
114
115            if str(item["year"]) == video.year:
116                matches.add("year")
117
118            if is_episode and any(
119                q["query"].lower() in item for item in (title, alt_title)
120            ):
121                matches.update(("title", "series", "season", "episode", "year"))
122
123            subtitles.append(
124                SuchaSubtitle(
125                    language,
126                    item["release"],
127                    item["filename"],
128                    str(item["id"]),
129                    type_str,
130                    matches,
131                )
132            )
133        return subtitles
134
135    def list_subtitles(self, video, languages):
136        return self.query(languages, video)
137
138    def _get_archive(self, content):
139        archive_stream = io.BytesIO(content)
140
141        if rarfile.is_rarfile(archive_stream):
142            logger.debug("Identified rar archive")
143            return rarfile.RarFile(archive_stream)
144
145        if zipfile.is_zipfile(archive_stream):
146            logger.debug("Identified zip archive")
147            return zipfile.ZipFile(archive_stream)
148
149        raise APIThrottled("Unsupported compressed format")
150
151    def get_file(self, archive):
152        for name in archive.namelist():
153            if os.path.split(name)[-1].startswith("."):
154                continue
155
156            if not name.lower().endswith(SUBTITLE_EXTENSIONS):
157                continue
158
159            if any(undesired in name.lower() for undesired in UNDESIRED_FILES):
160                continue
161
162            logger.debug(f"Returning from archive: {name}")
163            return archive.read(name)
164
165        raise APIThrottled("Can not find the subtitle in the compressed file")
166
167    def download_subtitle(self, subtitle):
168        logger.info("Downloading subtitle %r", subtitle)
169        response = self.session.get(
170            f"{SERVER_URL}/download",
171            params={"id": subtitle.download_id, "type": subtitle.download_type},
172            timeout=10,
173        )
174        response.raise_for_status()
175        archive = self._get_archive(response.content)
176        subtitle_file = self.get_file(archive)
177        subtitle.content = fix_line_ending(subtitle_file)
178