bb.py

# -*- coding: utf-8 -*-
import os
import sys
import shutil
import re
import subprocess

from textwrap import dedent
from collections import namedtuple, abc

import pymediainfo
import guessit
from unidecode import unidecode
from requests.exceptions import HTTPError

from .config import config
from .logging import log
from .torrent import make_torrent
from . import tvdb
from . import imdb
from .imgur import ImgurUploader
from .ffmpeg import FFMpeg
from . import templating as bb
from .submission import (Submission, form_field, finalize,
                         SubmissionAttributeError)
from .tracker import Tracker
from . import scene


def format_tag(tag):
    tag = unidecode(tag)
    if '/' in tag:
        # Multiple actors can be listed as a single actor like this:
        # "Thierry Kazazian / Max Mittleman"
        # (e.g. for "Miraculous: Tales of Ladybug & Cat Noir")
        tag = tag[:tag.index('/')].strip()
    return tag.replace(' ', '.').replace('-', '.').replace('\'', '.').lower()


def format_choices(choices):
    return ", ".join([
        str(num) + ": " + value
        for num, value in enumerate(choices)
    ])

def prompt_yesno(question, default):
    while True:
        choices = '[Y/n]' if default else '[y/N]'
        choice = input('%s %s ' % (question, choices))
        if not choice:
            return default
        elif choice.casefold() == 'y':
            return True
        elif choice.casefold() == 'n':
            return False


class BbSubmission(Submission):
    default_fields = ("form_title", "tags", "cover")

    def show_fields(self, fields):
        return super(BbSubmission, self).show_fields(
            fields or self.default_fields)

    def confirm_finalization(self, fields):
        return super(BbSubmission, self).confirm_finalization(
            fields or self.default_fields)

    def subcategory(self):
        # only video for now
        return VideoSubmission

    def subcategorise(self):
        log.debug('Attempting to narrow category')
        SubCategory = self.subcategory()
        if type(self) == SubCategory:
            return self

        log.info("Narrowing category from {} to {}",
                 type(self).__name__, SubCategory.__name__)
        sub = SubCategory(**self.fields)
        sub.depends_on = self.depends_on
        return sub

    @staticmethod
    def submit(payload):
        t = Tracker()
        return t.upload(**payload)

    # TODO: These should be detected as scene releases and produce an error message
    #       because they must be in a directory named after the release:

    # voa-the_omega_man_x264_bluray.mkv -> The.Omega.Man.1971.1080p.BluRay.x264-VOA/voa-the_omega_man_x264_bluray.mkv
    # hd1080-wtl.mkv -> Walk.the.Line.Extended.Cut.2005.1080p.BluRay.x264-HD1080/hd1080-wtl.mkv

    @form_field('scene', 'checkbox')
    def _render_scene(self):
        def ask_user():
            return prompt_yesno('Is this a scene release?', default=False)

        def handle_error(filepath, error):
            log.error(str(error))

        path = self['path']
        modified = scene.check_integrity(path, on_error=handle_error)
        # modified is True, False or None (unknown)
        if modified is False:
            return True
        elif modified is True:
            if prompt_yesno('Abort?', default=True):
                # The return statement is there because sys.exit() is mocked in tests.
                log.debug('Aborting')
                return sys.exit(1)
            else:
                log.debug('Asking user')
                return ask_user()

        # Check if release was renamed
        release_names = scene.release_names(path)
        if not release_names:
            return False
        elif release_names:
            print('Existing releases:')
            for release_name in release_names:
                print('  * {}'.format(release_name))
            log.error('This release was very likely modified and should not be uploaded like this.')
            if prompt_yesno('Abort?', default=True):
                # The return statement is there because sys.exit() is mocked in tests.
                return sys.exit(1)
            else:
                return ask_user()

    def data_method(self, source, target):
        def copy(source, target):
            if os.path.isfile(source):
                return shutil.copy(source, target)
            if os.path.isdir(source):
                return shutil.copytree(source, target)
            raise Exception('Source {} is neither '
                            'file nor directory'.format(source))

        cat_methods_map = {
            'movie': ['hard', 'sym', 'copy', 'move'],
            'tv': ['hard', 'sym', 'copy', 'move'],
            'music': ['copy', 'move'],
            }

        method_map = {'hard': os.link,
                      'sym': os.symlink,
                      'copy': copy,
                      'move': shutil.move}

        # use cmd line option if specified
        option_method = self['options'].get('data_method', 'auto')
        if option_method != 'auto':
            method = option_method
        else:
            pref_method = config.get('Torrent', 'data_method')
            if pref_method not in method_map:
                log.warning(
                    'Preferred method {} not valid. '
                    'Choices are {}'.format(pref_method,
                                            list(method_map.keys())))
            try:
                category = self._category
            except AttributeError:
                log.warning("{} does not have a category attribute",
                            type(self).__name__)
                category = 'movie'  # use movie data methods

            cat_methods = cat_methods_map[category]
            if pref_method in cat_methods:
                # use user preferred method if in category method list
                method = pref_method
            else:
                # otherwise use first category method
                method = cat_methods[0]

        log.notice('Copying data using \'{}\' method', method)
        return method_map[method](source, target)

    @finalize
    @form_field('file_input', 'file')
    def _render_torrentfile(self):
        return make_torrent(self['path'])

    def _finalize_torrentfile(self):
        # move data to upload directory
        up_dir = config.get('Torrent', 'upload_dir')
        path_dir, path_base = os.path.split(self['path'])
        if up_dir and not os.path.samefile(up_dir, path_dir):
            target = os.path.join(up_dir, path_base)
            print(target, os.path.exists(target))
            if not os.path.exists(target):
                self.data_method(self['path'], target)
            else:
                log.notice('Data method target already exists, skipping...')

        # black hole
        bh_dir = config.get('Torrent', 'black_hole')
        if bh_dir:
            fname = os.path.basename(self['torrentfile'])
            dest = os.path.join(bh_dir, fname)

            try:
                assert os.path.exists(bh_dir)
                assert not os.path.isfile(dest)
            except AssertionError as e:
                log.error(e)
            else:
                shutil.copy(self['torrentfile'], dest)
                log.notice("Torrent file copied to {}", dest)

        return self['torrentfile']

    @form_field('type')
    def _render_form_type(self):
        try:
            return self._form_type
        except AttributeError:
            raise SubmissionAttributeError(type(self).__name__ +
                                           ' has no _form_type attribute')

    @form_field('submit')
    def _render_form_submit(self):
        return 'true'


title_tv_re = (
    r"^(?P<title>.+)(?<!season) "
    r"(?P<season_marker>(s|season |))"
    r"(?P<season>((?<= s)[0-9]{2,})|(?<= )[0-9]+(?=x)|(?<=season )[0-9]+(?=$))"
    r"((?P<episode_marker>[ex])(?P<episode>[0-9]+))?$")

TvSpecifier = namedtuple('TvSpecifier', ['title', 'season', 'episode'])


class VideoSubmission(BbSubmission):
    default_fields = BbSubmission.default_fields

    def _render_guess(self):
        return dict(guessit.guessit(self['path']))

    def subcategory(self):
        if type(self) == VideoSubmission:
            if self['tv_specifier']:
                return TvSubmission
            else:
                return MovieSubmission
        return type(self)

    def _render_title(self):
        # Use format "<original title> AKA <english title>" where applicable
        title_original = self['summary']['title']
        title_english = self['summary']['titles'].get('XWW', None)
        if title_english is not None and title_original != title_english:
            return '{} AKA {}'.format(title_original, title_english)
        else:
            return title_original

    def _render_tv_specifier(self):
        # if title is specified, look if season/episode are set
        if self['title_arg']:
            match = re.match(title_tv_re, self['title_arg'],
                             re.IGNORECASE)
            if match:
                episode = match.group('episode')
                return TvSpecifier(
                    match.group('title'), int(match.group('season')),
                    episode and int(episode))  # if episode is None

            # todo: test tv show name from title_arg, but episode from filename

        guess = self['guess']
        if guess['type'] == 'episode':
            if self['title_arg']:
                title = self['title_arg']
            else:
                title = guess['title']
            return TvSpecifier(title, guess['season'],
                               guess.get('episode', None))

    @form_field('tags')
    def _render_tags(self):
        # todo: get episode-specific actors (from imdb?)

        n = self['options']['num_cast']
        tags = list(self['summary']['genres'])
        tags += [a['name']
                 for a in self['summary']['cast'][:n]
                 if a['name']]
        return ",".join(format_tag(tag) for tag in tags)

    def _render_mediainfo_path(self):
        assert os.path.exists(self['path'])
        if os.path.isfile(self['path']):
            return self['path']

        contained_files = []
        for dp, dns, fns in os.walk(self['path']):
            contained_files += [os.path.join(dp, fn) for fn in fns
                                if (os.path.getsize(os.path.join(dp, fn))
                                    > 10 * 2**20)]
        if len(contained_files) == 1:
            return contained_files[0]

        print("\nWhich file would you like to run mediainfo on? Choices are")
        contained_files.sort()
        for k, v in enumerate(contained_files):
            print("{}: {}".format(k, os.path.relpath(v, self['path'])))
        while True:
            try:
                choice = input(
                    "Enter [0-{}]: ".format(len(contained_files) - 1))
                return contained_files[int(choice)]
            except (ValueError, IndexError):
                pass

    @finalize
    def _render_screenshots(self):
        ns = self['options']['num_screenshots']
        ffmpeg = FFMpeg(self['mediainfo_path'])
        return ffmpeg.take_screenshots(ns)

    def _finalize_screenshots(self):
        return ImgurUploader().upload(self['screenshots'])

    def _render_mediainfo(self):
        try:
            path = self['mediainfo_path']
            if os.name == "nt":
                mi = subprocess.Popen([r"mediainfo", path], shell=True,
                                      stdout=subprocess.PIPE
                                      ).communicate()[0].decode('utf8')
            else:
                mi = subprocess.Popen([r"mediainfo", path],
                                      stdout=subprocess.PIPE
                                      ).communicate()[0].decode('utf8')
        except OSError:
            sys.stderr.write(
                "Error: Media Info not installed, refer to "
                "http://mediainfo.sourceforge.net/en for installation")
            exit(1)
        else:
            # Replace absolute path with file name
            mi_dir = os.path.dirname(self['mediainfo_path']) + '/'
            mi = mi.replace(mi_dir, '')

            # bB's mediainfo parser expects "Xbps" instead of "Xb/s"
            mi = mi.replace('Kb/s', 'Kbps') \
                   .replace('kb/s', 'Kbps') \
                   .replace('Mb/s', 'Mbps')
            return mi

    def _render_tracks(self):
        video_tracks = []
        audio_tracks = []
        text_tracks = []
        general = None

        mi = pymediainfo.MediaInfo.parse(self['mediainfo_path'])

        for track in mi.tracks:
            if track.track_type == 'General':
                general = track.to_data()
            elif track.track_type == 'Video':
                video_tracks.append(track.to_data())
            elif track.track_type == 'Audio':
                audio_tracks.append(track.to_data())
            elif track.track_type == 'Text':
                text_tracks.append(track.to_data())
            else:
                log.debug("Unknown track {}", track)

        assert general is not None
        assert len(video_tracks) == 1
        video_track = video_tracks[0]

        assert len(audio_tracks) >= 1

        return {'general': general,
                'video': video_track,
                'audio': audio_tracks,
                'text': text_tracks}

    def _render_source(self):
        sources = ('BluRay', 'BluRay 3D', 'WEB-DL',
                   'WebRip', 'HDTV', 'DVDRip', 'DVDSCR', 'CAM')
        # ignored: R5, TeleSync, PDTV, SDTV, BluRay RC, HDRip, VODRip,
        #          TC, SDTV, DVD5, DVD9, HD-DVD

        # todo: replace with guess from self['guess']
        regpath = self['path'].lower().strip('-')
        if 'bluray' in regpath:
            return 'BluRay'  # todo: 3d
        elif 'webdl' in regpath:
            return 'WEB-DL'
        elif 'webrip' in regpath:
            return 'WebRip'
        elif 'hdtv' in regpath:
            return 'HDTV'
        # elif 'dvdscr' in self['path'].lower():
        #    markers['source'] = 'DVDSCR'
        else:
            print("File:", self['path'])
            print("Choices:", format_choices(sources))
            while True:
                choice = input("Please specify a source by number: ")
                try:
                    return sources[int(choice)]
                except (ValueError, IndexError):
                    print("Please enter a valid choice")

    def _render_container(self):
        general = self['tracks']['general']
        if general['format'] == 'Matroska':
            return 'MKV'
        elif general['format'] == 'AVI':
            return 'AVI'
        elif general['format'] == 'MPEG-4':
            return 'MP4'
        elif general['format'] == 'BDAV':
            return 'm2ts'
        else:
            raise RuntimeError("Unknown or unsupported container '{}'".format(
                general.format))

    def _render_video_codec(self):
        video_track = self['tracks']['video']
        # norm_bitrate = (float(bit_rate) /
        #     (video_track.width*video_track.height))
        if (video_track['codec_id'] in ('V_MPEG4/ISO/AVC', 'AVC', 'avc1') or
                video_track['format'] == 'AVC'):
            if 'x264' in video_track.get('writing_library', ''):
                return 'x264'
            else:
                return 'H.264'
        elif video_track['codec_id'] in ('V_MPEGH/ISO/HEVC', 'HEVC'):
            if 'x265' in video_track.get('writing_library', ''):
                return 'x265'
            else:
                return 'H.265'
        elif video_track['codec_id'] in ('V_MS/VFW/FOURCC / WVC1',):
            return 'VC-1'
        elif 'VP9' in video_track['codec_id']:
            return 'VP9'
        elif video_track['codec_id'] == 'XVID':
            return 'XVid'
        elif video_track['format'] == 'MPEG Video':
            return 'MPEG-2'
        else:
            msg = "Unknown or unsupported video codec '{}' ({})".format(
                            video_track.get('codec_id'),
                            video_track.get('writing_library'))
            raise RuntimeError(msg)

    def _render_audio_codec(self):
        audio_track = self['tracks']['audio'][0]  # main audio track
        if (audio_track.get('codec_id_hint') == 'MP3' or
                audio_track['codec_id'] in ('MPA1L3', '55')):
            return 'MP3'
        elif audio_track['codec_id'].lower().startswith('mp4a'):
            return 'AAC'
        elif 'Dolby Atmos' in audio_track['commercial_name']:
            return 'Dolby Atmos'
        elif 'DTS-HD' in audio_track['commercial_name']:
            if audio_track.get('other_format', '') == 'DTS XLL X':
                return 'DTS:X'
            return 'DTS-HD'

        if audio_track['codec_id'].startswith('A_'):
            audio_track['codec_id'] = audio_track['codec_id'][2:]
        audio_codecs = ('AC3', 'EAC3', 'DTS', 'FLAC', 'AAC', 'MP3', 'TRUEHD',
                        'PCM')
        for c in audio_codecs:
            if audio_track['codec_id'].startswith(c):
                c = c.replace('EAC3', 'AC-3') \
                     .replace('AC3', 'AC-3') \
                     .replace('TRUEHD', 'True-HD')
                return c

        raise ValueError("Unknown or unsupported audio codec '{}'".format(
            audio_track['codec_id']))

    def _render_resolution(self):
        resolutions = ('2160p', '1080p', '720p', '1080i', '720i',
                       '480p', '480i', 'SD')

        # todo: replace with regex?
        # todo: compare result with mediainfo
        for res in resolutions:
            if res.lower() in self['path'].lower():
                # warning: 'sd' might match any ol' title, but it's last anyway
                return res
        else:
            print("File:", self['path'])
            print("Choices:", format_choices(resolutions))
            while True:
                choice = input("Please specify a resolution by number: ")
                try:
                    return resolutions[int(choice)]
                except (ValueError, IndexError):
                    print("Please enter a valid choice")
        # from mediainfo and filename

    def _render_additional(self):
        additional = []
        video_track = self['tracks']['video']
        audio_tracks = self['tracks']['audio']
        text_tracks = self['tracks']['text']

        # print [(track.title, track.language) for track in text_tracks]
        # todo: rule checking, e.g.
        # main_audio = audio_tracks[0]
        # if (main_audio.language and main_audio.language != 'en' and
        #         not self['tracks']['text']):
        #     raise BrokenRule("Missing subtitles")

        if 'remux' in os.path.basename(self['path']).lower():
            additional.append('REMUX')

        if self['guess'].get('proper_count') and self['scene']:
            additional.append('PROPER')

        edition = self['guess'].get('edition')
        if isinstance(edition, str):
            additional.append(edition)
        elif isinstance(edition, abc.Sequence):
            additional.extend(edition)

        if 'BT.2020' in video_track.get('color_primaries', ''):
            additional.append('HDR10')

        for track in audio_tracks[1:]:
            if 'title' in track and 'commentary' in track['title'].lower():
                additional.append('w. Commentary')
                break
        if text_tracks:
            additional.append('w. Subtitles')

        return additional

    def _render_form_release_info(self):
        return " / ".join(self['additional'])

    @finalize
    @form_field('image')
    def _render_cover(self):
        return self['summary']['cover']

    def _finalize_cover(self):
        return ImgurUploader().upload(self['cover'])


class TvSubmission(VideoSubmission):
    default_fields = VideoSubmission.default_fields + ('form_description',)
    _form_type = 'TV'
    __form_fields__ = {
        'form_title': ('title', 'text'),
        'form_description': ('desc', 'text'),
        }

    def _render_guess(self):
        return dict(guessit.guessit(self['path'],
                                    options=('--type', 'episode')))

    def _render_search_title(self):
        return self['tv_specifier'].title

    @form_field('title')
    def _render_form_title(self):
        markers_list = [self['source'], self['video_codec'],
                        self['audio_codec'], self['container'],
                        self['resolution']] + self['additional']
        markers = " / ".join(markers_list)

        if self['tv_specifier'].episode is not None:
            return "{t} S{s:02d}E{e:02d} [{m}]".format(
                t=self['title'], s=self['tv_specifier'].season,
                e=self['tv_specifier'].episode,
                m=markers)
        else:
            return "{t} - Season {s} [{m}]".format(
                t=self['title'], s=self['tv_specifier'].season, m=markers)

    def _render_summary(self):
        t = tvdb.TVDB()
        result = t.search(self['tv_specifier'])
        return result.summary()

    def _render_section_description(self):
        summary = self['summary']
        if 'episodesummary' in summary:
            return summary['seriessummary'] + bb.spoiler(
                summary['episodesummary'], "Episode description")
        else:
            return summary['seriessummary']

    def _render_section_information(self):
        def imdb_link(r):
            return bb.link(r.name, "https://www.imdb.com"+r.id)

        s = self['summary']
        links = [('TVDB', s['url'])]

        if s['imdb_id']:
            links.append(('IMDb',
                          "https://www.imdb.com/title/" + s['imdb_id']))
        # todo unify rating_bb and episode_fmt
        # get ratings from imdb
        i = imdb.IMDB()
        if self['tv_specifier'].episode is not None:
            if s['imdb_id']:
                rating, votes = i.get_rating(s['imdb_id'])

                rating_bb = (bb.format_rating(rating[0], max=rating[1]) + " " +
                             bb.s1("({votes} votes)".format(
                                 votes=votes)))
            else:
                rating_bb = ""

            description = dedent("""\
            [b]Episode title[/b]: {title} ({links})
            [b]Aired[/b]: {air_date} on {network}
            [b]IMDb Rating[/b]: {rating}
            [b]Directors[/b]: {directors}
            [b]Writer(s)[/b]: {writers}
            [b]Content rating[/b]: {contentrating}""").format(
                title=s['episode_title'],
                links=", ".join(bb.link(*l) for l in links),  # noqa: E741
                air_date=s['air_date'],
                network=s['network'],
                rating=rating_bb,
                directors=' | '.join(s['directors']),
                writers=' | '.join(s['writers']),
                contentrating=s['contentrating']
            )
        else:
            description = dedent("""\
            [b]Network[/b]: {network}
            [b]Content rating[/b]: {contentrating}\n""").format(
                contentrating=s['contentrating'],
                network=s['network'],
            )

            def episode_fmt(e):
                if not e['imdb_id']:
                    return bb.link(e['title'], e['url']) + "\n"

                try:
                    rating, votes = i.get_rating(e['imdb_id'])
                except ValueError:
                    return ''
                else:
                    return (bb.link(e['title'], e['url']) + "\n" +
                            bb.s1(bb.format_rating(*rating)))

            description += "[b]Episodes[/b]:\n" + bb.list(
                map(episode_fmt, s['episodes']), style=1)

        return description

    def _render_description(self):
        sections = [("Description", self['section_description']),
                    ("Information", self['section_information'])]

        description = "\n".join(bb.section(*s) for s in sections)
        description += bb.release
        return description

    @form_field('desc')
    def _render_form_description(self):
        ss = "".join(map(bb.img, self['screenshots']))
        return (self['description'] + "\n" +
                bb.section("Screenshots", bb.center(ss)) +
                bb.mi(self['mediainfo']))


class MovieSubmission(VideoSubmission):
    default_fields = (VideoSubmission.default_fields +
                      ("description", "mediainfo", "screenshots"))
    _form_type = 'Movies'
    __form_fields__ = {
        # field -> form field, type
        'source': ('source', 'text'),
        'video_codec': ('videoformat', 'text'),
        'audio_codec': ('audioformat', 'text'),
        'container': ('container', 'text'),
        'resolution': ('resolution', 'text'),
        'form_release_info': ('remaster_title', 'text'),
        'mediainfo': ('release_desc', 'text'),
        'screenshots': (lambda i, v: 'screenshot' + str(i + 1), 'text'),
        }

    def _render_guess(self):
        return dict(guessit.guessit(self['path'],
                                    options=('--type', 'movie')))

    def _render_search_title(self):
        if self['title_arg']:
            return self['title_arg']

        return self['guess']['title']

    @form_field('title')
    def _render_form_title(self):
        return self['title']

    @form_field('year')
    def _render_year(self):
        if 'summary' in self.fields:
            return self['summary']['year']

        elif 'year' in self['guess']:
            return self['guess']['year']

        else:
            while True:
                year = input('Please enter year: ')
                try:
                    year = int(year)
                except ValueError:
                    pass
                else:
                    return year

    def _render_summary(self):
        i = imdb.IMDB()
        movie = i.search(self['search_title'])
        return movie.summary()

    def _render_section_information(self):
        def imdb_link(r):
            return bb.link(r['name'], "https://www.imdb.com"+r['id'])

        # todo: synopsis/longer description
        n = self['options']['num_cast']
        summary = self['summary']
        links = [("IMDb", summary['url'])]

        return dedent("""\
        [b]Title[/b]: {name} ({links})
        [b]MPAA[/b]: {mpaa}
        [b]Rating[/b]: {rating} [size=1]({votes} votes)[/size]
        [b]Runtime[/b]: {runtime}
        [b]Director(s)[/b]: {directors}
        [b]Writer(s)[/b]: {writers}
        [b]Cast[/b]: {cast}""").format(
            links=", ".join(bb.link(*l) for l in links),  # noqa: E741
            name=summary['name'],
            mpaa=summary['mpaa'],
            rating=bb.format_rating(summary['rating'][0],
                                    max=summary['rating'][1]),
            votes=summary['votes'],
            runtime=summary['runtime'],
            directors=" | ".join(imdb_link(d) for d in summary['directors']),
            writers=" | ".join(imdb_link(w) for w in summary['writers']),
            cast=" | ".join(imdb_link(a) for a in summary['cast'][:n])
        )

    def _render_section_description(self):
        s = self['summary']
        return s['description']

    def _render_description(self):
        # todo: templating, rottentomatoes, ...

        sections = [("Description", self['section_description']),
                    ("Information", self['section_information'])]

        description = "\n".join(bb.section(*s) for s in sections)
        description += bb.release

        return description

    @form_field('desc')
    def _render_form_description(self):
        return self['description']