1# -*- coding: utf-8 -*-
2from __future__ import absolute_import
3import logging
4
5from babelfish import Language, language_converters
6from requests import Session
7
8from . import Provider
9from .. import __short_version__
10from ..subtitle import Subtitle, fix_line_ending
11
12logger = logging.getLogger(__name__)
13
14language_converters.register('thesubdb = subliminal.converters.thesubdb:TheSubDBConverter')
15
16
17class TheSubDBSubtitle(Subtitle):
18    """TheSubDB Subtitle."""
19    provider_name = 'thesubdb'
20
21    def __init__(self, language, hash):
22        super(TheSubDBSubtitle, self).__init__(language)
23        self.hash = hash
24
25    @property
26    def id(self):
27        return self.hash + '-' + str(self.language)
28
29    def get_matches(self, video):
30        matches = set()
31
32        # hash
33        if 'thesubdb' in video.hashes and video.hashes['thesubdb'] == self.hash:
34            matches.add('hash')
35
36        return matches
37
38
39class TheSubDBProvider(Provider):
40    """TheSubDB Provider."""
41    languages = {Language.fromthesubdb(l) for l in language_converters['thesubdb'].codes}
42    required_hash = 'thesubdb'
43    server_url = 'http://api.thesubdb.com/'
44    subtitle_class = TheSubDBSubtitle
45
46    def __init__(self):
47        self.session = None
48
49    def initialize(self):
50        self.session = Session()
51        self.session.headers['User-Agent'] = ('SubDB/1.0 (subliminal/%s; https://github.com/Diaoul/subliminal)' %
52                                              __short_version__)
53
54    def terminate(self):
55        self.session.close()
56
57    def query(self, hash):
58        # make the query
59        params = {'action': 'search', 'hash': hash}
60        logger.info('Searching subtitles %r', params)
61        r = self.session.get(self.server_url, params=params, timeout=10)
62
63        # handle subtitles not found and errors
64        if r.status_code == 404:
65            logger.debug('No subtitles found')
66            return []
67        r.raise_for_status()
68
69        # loop over languages
70        subtitles = []
71        for language_code in r.text.split(','):
72            language = Language.fromthesubdb(language_code)
73
74            subtitle = self.subtitle_class(language, hash)
75            logger.debug('Found subtitle %r', subtitle)
76            subtitles.append(subtitle)
77
78        return subtitles
79
80    def list_subtitles(self, video, languages):
81        return [s for s in self.query(video.hashes['thesubdb']) if s.language in languages]
82
83    def download_subtitle(self, subtitle):
84        logger.info('Downloading subtitle %r', subtitle)
85        params = {'action': 'download', 'hash': subtitle.hash, 'language': subtitle.language.alpha2}
86        r = self.session.get(self.server_url, params=params, timeout=10)
87        r.raise_for_status()
88
89        subtitle.content = fix_line_ending(r.content)
90