dotfiles/config/scripts/tvshow

#!/usr/bin/env python
# pylint: disable=C0103,W0621

# pip install tmdbv3api

import os
import re
import subprocess
import sys
import typing

import tmdbv3api

# TODO Override files without warning
# TODO Dry run mode (just comment the last line ^^)

# Typing

Episode = typing.Any  # TODO

# Constants

API_KEY_PASS_PATH = 'http/themoviedb.org'
VIDEO_EXTENSIONS = {'mp4', 'mkv', 'avi', 'webm'}

# Functions


def get_pass_data(path: str) -> typing.Dict[str, str]:
    """
    Returns the data stored in the Unix password manager
    given its path.
    """
    run = subprocess.run(['pass', path], stdout=subprocess.PIPE, check=True)
    lines = run.stdout.decode().split('\n')
    data = dict()
    data['pass'] = lines[0]
    for line in lines[1:]:
        match = re.match(r'(\w+): ?(.+)', line)
        if match:
            data[match[1]] = match[2]
    return data


def confirm(text: str) -> bool:
    res = input(text + " [yn] ")
    while res not in ('y', 'n'):
        res = input("Please answer with y or n: ")
    return res == 'y'


def episode_identifier(episode: typing.Any) -> str:
    return f"S{episode['season_number']:02d}E" + \
        f"{episode['episode_number']:02d} {episode['name']}"

dryrun = '-n' in sys.argv
if dryrun:
    dryrun = True
    sys.argv.remove('-n')


# Connecting to TMBDB
tmdb = tmdbv3api.TMDb()
tmdb.api_key = get_pass_data(API_KEY_PASS_PATH)['api']
tmdb.language = sys.argv[1]

# Searching the TV show name (by current directory name)
tv = tmdbv3api.TV()
season = tmdbv3api.Season()
if len(sys.argv) >= 3:
    show_name = sys.argv[2]
else:
    show_name = os.path.split(os.path.realpath(os.path.curdir))[1]
    if '(' in show_name:
        show_name = show_name.split('(')[0].strip()

search = tv.search(show_name)

# Asking the user to select the one
show = None

for res in search:
    print(f"#{res.id} {res.name} ({res.first_air_date[:4]}): {res.overview}")
    if confirm("Is this the show for this folder?"):
        show = tv.details(res.id)
        break

if not show:
    print("Could not find a matching " +
          f"show on TheMovieDatabase for {show_name}.")
    sys.exit(1)

# Retrieving all the episode of the show
episodes: typing.List[Episode] = list()

print(f"List of episodes for {show.name}:")
for season_number in range(0, show.number_of_seasons+1):
    season_details = season.details(show.id, season_number)

    try:
        season_details.episodes
    except AttributeError:
        continue

    for episode in season_details.episodes:
        episodes.append(episode)
        print(f"- {episode_identifier(episode)}")

# Finding movie files in the folder
print("List of video files in this folder")
videos: typing.List[typing.Tuple[str, str]] = list()
for root, dirs, files in os.walk(os.path.curdir):
    for filename in files:
        basename, ext = os.path.splitext(filename)
        real_ext = ext[1:].lower()
        if real_ext not in VIDEO_EXTENSIONS:
            continue
        videos.append((root, filename))
        print(f"- {filename}")


def get_episode(season_number: int, episode_number: int
                ) -> typing.Optional[Episode]:
    # TODO Make more efficient using indexing
    for episode in episodes:
        if episode['season_number'] == season_number \
                and episode['episode_number'] == episode_number:
            return episode
    return None


# Matching movie files to episode
associations: typing.List[typing.Tuple[typing.Tuple[str,
                                                    str], Episode]] = list()
for video in videos:
    root, filename = video
    match = re.search(r'S(\d+)E(\d+)', filename)
    print(f"Treating file: {root}/{filename}")
    episode = None
    season_number = 0
    episode_number = 0
    while not episode:
        if match:
            season_number = int(match[1])
            episode_number = int(match[2])
        else:
            try:
                season_number = int(input("Season number ?"))
                episode_number = int(input("Episode number ?"))
            except ValueError:
                continue
        if season_number < 0 and episode_number < 0:
            break
        match = None
        episode = get_episode(season_number, episode_number)
        if not episode:
            print(
                f"  could not find episode S{season_number:02d}E{episode_number:02d} in TMBD")

    # Skip
    if not episode:
        if season_number < -1 and episode_number < -1:
            # Skip all
            break
        # Skip one
        continue

    associations.append((video, episode))
    print(f"  associated to: {episode_identifier(episode)}")


# Rename video files
for association in associations:
    video, episode = association
    root, filename = video
    basename, ext = os.path.splitext(filename)
    new_name = f"{show.name} ({show.first_air_date[:4]}) {episode_identifier(episode)}"
    # Rename all file with the same base name as the original file so we
    # can rename nfo files and subtitles (only one though)
    for a_filename in os.listdir(root):
        a_basename, a_ext = os.path.splitext(a_filename)
        if a_basename == basename:
            old_path = os.path.join(root, a_filename)
            new_path = os.path.join(root, new_name + a_ext)
            if old_path == new_path:
                continue

            print(old_path, "->", new_path)
            if not dryrun:
                os.rename(old_path, new_path)