renamer

#!/usr/bin/env python3

from os.path import isdir, isfile
import requests
import re
import sys
import os


class Media:
    def __init__(self, tmdb_api_key, dest_path, files):
        self.dest_path = dest_path
        self.files = files
        self.tmdb_url = "https://api.themoviedb.org/3"
        self.tmdb_api_key = tmdb_api_key
        self.cache = {}

    def ask_to_select(self, orig_filename, selection_options):
        """Interface to select from search options"""

        print(orig_filename)
        print("\n".join(selection_options))

        try:
            selection = int(input("\nSelect option using num: "))
        except KeyboardInterrupt:
            exit(0)

        if selection < 1 or selection > len(selection_options):
            print("Index out of range...")
            exit(1)

        return int(selection) - 1

    def confirm_and_exec(self, rename_data):
        if not rename_data:
            return

        print("\n".join(f"{old} -> {new}" for old, new in rename_data))

        try:
            selection = input("\nDo you want to proceed?: ")
        except KeyboardInterrupt:
            exit(0)

        if selection == "y":
            for old, new in rename_data:
                os.makedirs(os.path.dirname(new), exist_ok=True)
                os.rename(old, new)

    def sanitize(self, name):
        blacklist = set("?:!/;'\",=")
        return "".join(c for c in name if c not in blacklist)

    def rename(self):
        pass


class Tv(Media):
    def __init__(self, tmdb_api_key, dest_path, files):
        super().__init__(tmdb_api_key, dest_path, files)

    def get_details(self, orig_filename, show_name):
        """Returns TMDB ID, Name of Show and Release Year"""
        url = f"{self.tmdb_url}/search/tv"
        res = requests.get(
            url, params={"api_key": self.tmdb_api_key, "query": show_name}
        )
        res.raise_for_status()
        res = res.json()["results"]

        if not res:
            new_show_name = " ".join(show_name.split()[:-1])
            if new_show_name:
                return self.get_details(orig_filename, new_show_name)

        if len(res) == 1:
            res = res[0]
        else:
            selection_options = [
                f"{i+1} {item.get('name')} ({item.get('first_air_date')}) [{item.get('id')}]"
                for i, item in enumerate(res)
            ]
            selection = self.ask_to_select(orig_filename, selection_options)
            res = res[selection]

        details = (
            res.get("id"),
            res.get("name"),
            res.get("first_air_date").split("-")[0],
        )

        return details

    def get_season_and_episode(self, file):
        """Returns season and episode number"""
        res = re.findall(r"s\d{1,2}e\d{1,2}", file.lower())
        if not res:
            return

        season, episode = re.findall(r"\d{1,2}", res[0])
        return season, episode

    def extract_show(self, path):
        """Extracts TV show name from file name"""
        path = re.sub(r"\W+", " ", path.lower())
        res = re.split(r"s\d{1,2}e\d{1,2}", path)

        return res[0].strip()

    def rename(self):
        rename_data = []
        for file in self.files:
            base_path = os.path.basename(file)

            show_name = self.extract_show(base_path)
            if not show_name:
                continue
            # print(f"{show_name=}")

            show_details = None
            if self.cache.get(show_name):
                show_details = self.cache.get(show_name)
            else:
                show_details = self.get_details(base_path, show_name)
                self.cache[show_name] = show_details

            if not show_details:
                continue
            tmdb_id, tmdb_name, tmdb_year = show_details
            # print(f"{tmdb_id=}, {tmdb_name=}, {tmdb_year=}")

            se = self.get_season_and_episode(file)
            if not se:
                continue

            season, episode = se
            # print(f"{season=}, {episode=}")

            new_path = os.path.join(
                self.dest_path,
                f"{self.sanitize(tmdb_name)} ({tmdb_year})",
                f"Season {season}",
                base_path,
            )
            rename_data.append((file, new_path))

        self.confirm_and_exec(rename_data)


class Movie(Media):
    def __init__(self, tmdb_api_key, dest_path, files):
        super().__init__(tmdb_api_key, dest_path, files)

    def get_name(self, file):
        """Extracts movie name from file name"""
        file = re.sub(r"\W+", " ", file.lower())
        res = re.split(r"\d{4}", file)

        return res[0].strip()

    def get_year(self, file):
        res = re.findall(r"\d{4}", file.lower())
        return res[0]

    def get_collection(self, tmdb_id):
        url = f"{self.tmdb_url}/movie/{tmdb_id}"
        res = requests.get(url, params={"api_key": self.tmdb_api_key})
        res.raise_for_status()
        res = res.json()

        # print(res)
        if res.get("belongs_to_collection"):
            return res.get("belongs_to_collection").get("name")

        return None

    def get_details(self, orig_filename, name, year):
        url = f"{self.tmdb_url}/search/movie"
        res = requests.get(
            url, params={"api_key": self.tmdb_api_key, "query": name, "year": year}
        )
        res.raise_for_status()
        res = res.json()["results"]

        if not res:
            new_name = " ".join(name.split()[:-1])
            if new_name:
                return self.get_details(orig_filename, new_name, year)

        # print(f"{name=} {year=}")
        if len(res) == 1:
            res = res[0]
        else:
            selection_options = [
                f"{i+1} {item.get('original_title')} ({item.get('release_date')}) [{item.get('id')}]"
                for i, item in enumerate(res)
            ]
            selection = self.ask_to_select(orig_filename, selection_options)
            res = res[selection]

        movie_collection = self.get_collection(res.get("id")) or ""

        return (
            res.get("id"),
            res.get("original_title"),
            res.get("release_date").split("-")[0],
            movie_collection,
        )

    def rename(self):
        rename_data = []
        for file in self.files:
            base_path = os.path.basename(file)
            # print(base_path)
            movie_name = self.get_name(base_path)
            if not movie_name:
                continue

            movie_year = self.get_year(base_path)

            movie_details = self.get_details(base_path, movie_name, movie_year)
            if not movie_details:
                continue

            tmdb_id, tmdb_name, tmdb_year, tmdb_col = movie_details

            new_path = os.path.join(
                self.dest_path, self.sanitize(tmdb_col), f"{self.sanitize(tmdb_name)} ({tmdb_year})", base_path
            )
            rename_data.append((file, new_path))

        self.confirm_and_exec(rename_data)


if __name__ == "__main__":
    tmdb_api_key = os.getenv("TMDB_API_KEY")
    if not tmdb_api_key:
        print("Please export TMDB_API_KEY")
        exit(1)

    if len(sys.argv) < 4:
        print("Usage: python3 renamer.py <tv|mov> <dest_path> <files and/or dirs>")
        exit(1)

    def is_media_file(file):
        _, ext = os.path.splitext(file)
        return ext in [".mkv", ".mp4", ".avi"]

    paths = []
    for path in sys.argv[3:]:
        if isdir(path):
            paths.extend(
                [
                    os.path.join(r, file)
                    for r, _, f in os.walk(path)
                    for file in f
                    if is_media_file(file)
                ]
            )
        elif isfile(path) and is_media_file(path):
            paths.append(path)

    if not paths:
        print("No media files found in given path")
        exit(1)

    media = None
    if sys.argv[1] == "tv":
        media = Tv(tmdb_api_key, sys.argv[2], paths)
    if sys.argv[1] == "mov":
        media = Movie(tmdb_api_key, sys.argv[2], paths)

    if not media:
        print("Invalid media type")
        exit(1)

    media.rename()

    # delete dir if its empty after rename
    print("deleting empty directories...")
    for path in sys.argv[3:]:
        if not isdir(path):
            continue

        dir_tree = [
            os.path.join(r, dir)
            for r, d, _ in os.walk(path, topdown=False)
            for dir in d
        ]
        dir_tree.append(path)

        for d in dir_tree:
            try:
                os.rmdir(d)
            except OSError as e:
                print(e)