plex_trakt_trending_movies.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Automated Trending Movies Plex library script

Usage:
    $ python plex_trakt_trending_movies.py
    $ python plex_trakt_trending_movies.py --sort-only

Requirements:
    requests
    plexapi
    trakt

Disclaimer:
    Use at your own risk! I am not responsible for damages to your Plex server or libraries.

Author:
    /u/haeri

Credit:
    Originally based on https://gist.github.com/JonnyWong16/f5b9af386ea58e19bf18c09f2681df23
    by /u/SwiftPanda16
"""

import sys
import os
import codecs
import json
import subprocess
import time
import datetime

import requests
from plexapi.server import PlexServer
import trakt

# Config
# Plex server details
PLEX_URL = 'http://localhost:32400'
# See https://support.plex.tv/hc/en-us/articles/204059436-Finding-an-authentication-token-X-Plex-Token
PLEX_TOKEN = ''
PLEX_HOME_DIR = '/usr/lib/plexmediaserver'
PLEX_MEDIA_SCANNER_PATH = '/usr/lib/plexmediaserver/Plex Media Scanner'
# Windows: 'C:\Program Files (x86)\Plex\Plex Media Server\Plex Media Scanner.exe'
# Mac OS: '/Applications/Plex\ Media\ Server.app/Contents/MacOS/Plex\ Media\ Scanner'

# Trakt API details
# Required
# Create a Trakt account, then create an API app here:
# https://trakt.tv/oauth/applications/new
TRAKT_USERNAME = ''
TRAKT_CLIENT_ID = ''
TRAKT_CLIENT_SECRET = ''
# Experiment with the limits in the URLs below to get a different balance.
# Currently watched. This tends to bias towards lots of animated movies
# during certain hours of the day...
TRAKT_TRENDING_URL = 'https://api.trakt.tv/movies/trending?limit=5'
# Watch count for the past [week, month, year]
TRAKT_WATCHED_WEEKLY_URL = 'https://api.trakt.tv/movies/watched/weekly?limit=55'
TRAKT_WATCHED_MONTHLY_URL = 'https://api.trakt.tv/movies/watched/monthly?limit=125'
TRAKT_WATCHED_YEARLY_URL = 'https://api.trakt.tv/movies/watched/yearly?limit=500'

# Existing movie library details
MOVIE_LIBRARY_NAME = 'Movies'
MOVIE_LIBRARY_FOLDERS = ['/mnt/plex/sorted/Movies']  # List of movie folders

# New trending library details
TRENDING_LIBRARY_NAME = 'Movies - Trending'
# New folder to symlink existing movies to
TRENDING_FOLDER = '/mnt/plex/sorted/Movies Trending/'
SORT_TITLE_FORMAT = u"{number} {title}"
MAX_AGE = 4  # Limit the age (in years) of movies to be considered
             # (0 for no limit)
MAX_COUNT = 200  # Maximum number of movies to keep in the library

# The Movie Database details
# Enter your TMDb API key if your movie library is using
# "The Movie Database" agent.
# This will be used to convert the TMDb IDs to IMDB IDs.
# You can leave this blank '' if your movie library is using the 
# "Plex Movie" agent.
TMDB_API_KEY = ''

# End config

TMDB_REQUEST_COUNT = 0  # DO NOT CHANGE


def create_trending_library():
    headers = {"X-Plex-Token": PLEX_TOKEN}
    params = {
        'name': TRENDING_LIBRARY_NAME,
        'type': 'movie',
        'agent': 'com.plexapp.agents.imdb',
        'scanner': 'Plex Movie Scanner',
        'language': 'en',
        'location': TRENDING_FOLDER
    }

    url = '{base_url}/library/sections'.format(base_url=PLEX_URL)
    r = requests.post(url, headers=headers, params=params)


def add_sort_title(library_key, rating_key, number, title):
    headers = {'X-Plex-Token': PLEX_TOKEN}
    params = {
        'type': 1,
        'id': rating_key,
        'title.locked': 1,
        'titleSort.value': SORT_TITLE_FORMAT.format(
            number=str(number).zfill(6), title=title),
        'titleSort.locked': 1,
    }

    url = "{base_url}/library/sections/{library}/all".format(
            base_url=PLEX_URL, library=library_key)
    r = requests.put(url, headers=headers, params=params)


def get_imdb_id_from_tmdb(tmdb_id):
    global TMDB_REQUEST_COUNT

    if not TMDB_API_KEY:
        return None

    # Wait 10 seconds for the TMDb rate limit
    if TMDB_REQUEST_COUNT >= 40:
        time.sleep(10)
        TMDB_REQUEST_COUNT = 0

    params = {"api_key": TMDB_API_KEY}

    url = "https://api.themoviedb.org/3/movie/{tmdb_id}".format(
            tmdb_id=tmdb_id)
    r = requests.get(url, params=params)

    TMDB_REQUEST_COUNT += 1

    if r.status_code == 200:
        movie = json.loads(r.text)
        return movie['imdb_id']
    else:
        return None


def refresh_library(library, path):
    section = library.key
    e = dict(os.environ)
    e['LC_ALL'] = "en_US.UTF-8"
    e['PLEX_MEDIA_SERVER_MAX_PLUGIN_PROCS'] = "6"
    e['PLEX_MEDIA_SERVER_TMPDIR'] = "/tmp"
    e['PLEX_MEDIA_SERVER_HOME'] = PLEX_HOME_DIR
    e['LD_LIBRARY_PATH'] = PLEX_HOME_DIR
    print(u"Scanning {path}".format(path=path))
    subprocess.call([PLEX_MEDIA_SCANNER_PATH, '--scan',
                     '--refresh', '--section', section, '--directory', path],
                    env=e)


def run_trakt_watched_sort_only():
    try:
        plex = PlexServer(PLEX_URL, PLEX_TOKEN)
    except:
        print(u"No Plex server found at: {base_url}".format(base_url=PLEX_URL))
        print(u"Exiting script.")
        return 0

    trakt.init(TRAKT_USERNAME, client_id=TRAKT_CLIENT_ID,
               client_secret=TRAKT_CLIENT_SECRET)
    core = trakt.core.Core()
    watched_movies = []
    watched_ids = []
    curyear = datetime.datetime.now().year

    def _add_from_list(url):
        print(u"Retrieving the trakt list: {}".format(url))
        movie_data = core._handle_request('get', url)
        for m in movie_data:
            # Skip already added movies
            if m['movie']['ids']['imdb'] in watched_ids:
                continue
            # Skip old movies
            if MAX_AGE != 0 \
                    and (curyear - (MAX_AGE - 1)) > int(m['movie']['year']):
                continue
            watched_movies.append({
                'id': m['movie']['ids']['imdb'],
                'title': m['movie']['title'].encode('utf8'),
                'year': m['movie']['year'],
            })
            watched_ids.append(m['movie']['ids']['imdb'])
            print(u"{} {} {}".format(
                len(watched_movies), m['movie']['title'], m['movie']['year']))

    # Get the trakt trending list
    _add_from_list(TRAKT_TRENDING_URL)

    # Get the trakt watched lists
    _add_from_list(TRAKT_WATCHED_WEEKLY_URL)
    _add_from_list(TRAKT_WATCHED_MONTHLY_URL)
    _add_from_list(TRAKT_WATCHED_YEARLY_URL)

    trending_library = plex.library.section(TRENDING_LIBRARY_NAME)
    trending_library_key = trending_library.key
    all_trending_movies = trending_library.all()

    # Create a dictionary of {imdb_id: movie}
    imdb_map = {}
    for m in all_trending_movies:
        if m.guid != None and 'imdb://' in m.guid:
            imdb_id = m.guid.split('imdb://')[1].split('?')[0]
        elif m.guid != None and 'themoviedb://' in m.guid:
            tmdb_id = m.guid.split('themoviedb://')[1].split('?')[0]
            imdb_id = get_imdb_id_from_tmdb(tmdb_id)
        else:
            imdb_id = None

        if imdb_id and imdb_id in watched_ids:
            imdb_map[imdb_id] = m
        else:
            imdb_map[m.ratingKey] = m

    # Modify the sort title to match the trakt watched order
    print(u"Setting the sort titles for the '{}' library...".format(
        TRENDING_LIBRARY_NAME))
    in_library_idx = []
    i = 0
    for m in watched_movies:
        movie = imdb_map.pop(m['id'], None)
        if movie:
            i += 1
            add_sort_title(trending_library_key, movie.ratingKey, i,
                           m['title'])
            in_library_idx.append(i)


def run_trakt_watched():
    try:
        plex = PlexServer(PLEX_URL, PLEX_TOKEN)
    except:
        print(u"No Plex server found at: {base_url}".format(base_url=PLEX_URL))
        print(u"Exiting script.")
        return 0

    trakt.init(TRAKT_USERNAME, client_id=TRAKT_CLIENT_ID,
               client_secret=TRAKT_CLIENT_SECRET)
    core = trakt.core.Core()
    watched_movies = []
    watched_ids = []
    curyear = datetime.datetime.now().year

    def _add_from_list(url):
        print(u"Retrieving the trakt list: {}".format(url))
        movie_data = core._handle_request('get', url)
        for m in movie_data:
            # Skip already added movies
            if m['movie']['ids']['imdb'] in watched_ids:
                continue
            # Skip old movies
            if MAX_AGE != 0 and \
                    (curyear - (MAX_AGE - 1)) > int(m['movie']['year']):
                continue
            watched_movies.append({
                'id': m['movie']['ids']['imdb'],
                'title': m['movie']['title'].encode('utf8'),
                'year': m['movie']['year'],
            })
            watched_ids.append(m['movie']['ids']['imdb'])

    # Get the trakt trending list
    _add_from_list(TRAKT_TRENDING_URL)

    # Get the trakt watched lists
    _add_from_list(TRAKT_WATCHED_WEEKLY_URL)
    _add_from_list(TRAKT_WATCHED_MONTHLY_URL)
    _add_from_list(TRAKT_WATCHED_YEARLY_URL)

    # Get list of movies from the Plex server
    print(u"Retrieving a list of movies from the '{library}' library in "
          u"Plex...".format(library=MOVIE_LIBRARY_NAME))
    try:
        movie_library = plex.library.section(MOVIE_LIBRARY_NAME)
        all_movies = movie_library.all()
    except:
        print(u"The '{library}' library does not exist in Plex.".format(
            library=MOVIE_LIBRARY_NAME))
        print(u"Exiting script.")
        return 0

    # Create a list of matching movies
    matching_movies = []
    nonmatching_idx = []
    for i, m in enumerate(watched_movies):
        if len(matching_movies) >= MAX_COUNT:
            nonmatching_idx.append(i)
            continue
        try:
            res = movie_library.search(title=m['title'], year=m['year'])
            if not res:
                res = movie_library.search(title=m['title'], year=int(m['year'])+1)
            if not res:
                res = movie_library.search(title=m['title'], year=int(m['year'])-1)
        except UnicodeDecodeError:
            res = None
        if res:
            for r in res:
                if r.guid != None and 'imdb://' in r.guid:
                    imdb_id = r.guid.split('imdb://')[1].split('?')[0]
                elif r.guid != None and 'themoviedb://' in r.guid:
                    tmdb_id = r.guid.split('themoviedb://')[1].split('?')[0]
                    imdb_id = get_imdb_id_from_tmdb(tmdb_id)
                else:
                    imdb_id = None

                if imdb_id and imdb_id == m['id']:
                    matching_movies.append(r)
                    print(u"{} {} {}".format(
                        len(matching_movies), m['title'], m['year']))
                    break
            else:
                nonmatching_idx.append(i)
        if not res:
            nonmatching_idx.append(i)

    for i in reversed(nonmatching_idx):
        del watched_movies[i]
        del watched_ids[i]

    # Create symlinks for all movies in your library on the trakt watched
    print(u"Creating symlinks for {count} matching movies in the "
          u"library...".format(count=len(matching_movies)))

    try:
        if not os.path.exists(TRENDING_FOLDER):
            os.mkdir(TRENDING_FOLDER)
    except:
        print(u"Unable to create the trending library folder "
              u"'{folder}'.".format(folder=TRENDING_FOLDER))
        print(u"Exiting script.")
        return 0

    count = 0
    updated_paths = []
    for movie in matching_movies:
        for part in movie.iterParts():
            old_path_file = part.file.encode('UTF-8')
            old_path, file_name = os.path.split(old_path_file)

            folder_name = ''
            for f in MOVIE_LIBRARY_FOLDERS:
                if old_path.lower().startswith(f.lower()):
                    folder_name = os.path.relpath(old_path, f)

            if folder_name == '.':
                new_path = os.path.join(TRENDING_FOLDER, file_name)
                dir = False
            else:
                new_path = os.path.join(TRENDING_FOLDER, folder_name)
                dir = True
                parent_path = os.path.dirname(os.path.abspath(new_path))
                if not os.path.exists(parent_path):
                    try:
                        os.makedirs(parent_path)
                    except OSError as e:
                        if e.errno == errno.EEXIST and \
                                os.path.isdir(parent_path):
                            pass
                        else:
                            raise
                # Clean up old, empty directories
                if os.path.exists(new_path) and not os.listdir(new_path):
                    os.rmdir(new_path)

            if (dir and not os.path.exists(new_path)) or \
                    (not dir and not os.path.isfile(new_path)):
                try:
                    if os.name == 'nt':
                        if dir:
                            subprocess.call(['mklink', '/D', new_path,
                                             old_path], shell=True)
                        else:
                            subprocess.call(['mklink', new_path,
                                             old_path_file], shell=True)
                    else:
                        if dir:
                            os.symlink(old_path, new_path)
                        else:
                            os.symlink(old_path_file, new_path)
                    count += 1
                    updated_paths.append(new_path)
                except Exception as e:
                    print(u"Symlink failed for {path}: {e}".format(
                        path=new_path, e=e))

    print(u"Created symlinks for {count} movies.".format(count=count))

    # Check if the trakt watched library exists in Plex
    print(u"Creating the '{}' library in Plex...".format(
        TRENDING_LIBRARY_NAME))
    try:
        trending_library = plex.library.section(TRENDING_LIBRARY_NAME)
        trending_library_key = trending_library.key
        print(u"Library already exists in Plex. Refreshing the library...")

        trending_library.update()
        #for path in updated_paths:
        #    refresh_library(trending_library, path)
    except:
        create_trending_library()
        trending_library = plex.library.section(TRENDING_LIBRARY_NAME)
        trending_library_key = trending_library.key

    # Wait for metadata to finish downloading before continuing
    print(u"Waiting for metadata to finish downloading...")
    trending_library = plex.library.section(TRENDING_LIBRARY_NAME)
    while trending_library.refreshing:
        time.sleep(5)
        trending_library = plex.library.section(TRENDING_LIBRARY_NAME)

    #time.sleep(5)

    # Retrieve a list of movies from the trakt watched library
    print(u"Retrieving a list of movies from the '{library}' library in "
          u"Plex...".format(library=TRENDING_LIBRARY_NAME))
    all_trending_movies = trending_library.all()

    # Create a dictionary of {imdb_id: movie}
    imdb_map = {}
    for m in all_trending_movies:
        if m.guid != None and 'imdb://' in m.guid:
            imdb_id = m.guid.split('imdb://')[1].split('?')[0]
        elif m.guid != None and 'themoviedb://' in m.guid:
            tmdb_id = m.guid.split('themoviedb://')[1].split('?')[0]
            imdb_id = get_imdb_id_from_tmdb(tmdb_id)
        else:
            imdb_id = None

        if imdb_id and imdb_id in watched_ids:
            imdb_map[imdb_id] = m
        else:
            imdb_map[m.ratingKey] = m

    # Modify the sort title to match the trakt watched order
    print(u"Setting the sort titles for the '{}' library...".format(
        TRENDING_LIBRARY_NAME))
    in_library_idx = []
    i = 0
    for m in watched_movies:
        movie = imdb_map.pop(m['id'], None)
        if movie:
            i += 1
            add_sort_title(trending_library_key, movie.ratingKey, i, m['title'])
            in_library_idx.append(i)

    # Remove movies from library with are no longer on the trakt watched list
    print(u"Removing symlinks for movies which are not on the trakt watched "
          u"list...".format(library=TRENDING_LIBRARY_NAME))
    count = 0
    updated_paths = []
    for movie in imdb_map.values():
        for part in movie.iterParts():
            old_path_file = part.file.encode('UTF-8')
            old_path, file_name = os.path.split(old_path_file)

            folder_name = os.path.relpath(old_path, TRENDING_FOLDER)

            if folder_name == '.':
                new_path = os.path.join(TRENDING_FOLDER, file_name)
                dir = False
            else:
                new_path = os.path.join(TRENDING_FOLDER, folder_name)
                dir = True

            if (dir and os.path.exists(new_path)) or \
                    (not dir and os.path.isfile(new_path)):
                try:
                    if os.name == 'nt':
                        if dir:
                            os.rmdir(new_path)
                        else:
                            os.remove(new_path)
                    else:
                        os.unlink(new_path)
                    count += 1
                    updated_paths.append(new_path)
                except Exception as e:
                    print(u"Remove symlink failed for {path}: {e}".format(
                        path=new_path, e=e))

    print(u"Removed symlinks for {count} movies.".format(count=count))

    # Refresh the library to clean up the deleted movies
    print(u"Refreshing the '{library}' library...".format(
        library=TRENDING_LIBRARY_NAME))
    trending_library.update()
    time.sleep(10)
    trending_library = plex.library.section(TRENDING_LIBRARY_NAME)
    while trending_library.refreshing:
        time.sleep(5)
        trending_library = plex.library.section(TRENDING_LIBRARY_NAME)
    trending_library.emptyTrash()

    return len(watched_ids)


if __name__ == "__main__":
    if '--sort-only' in sys.argv:
        run_trakt_watched_sort_only()
    else:
        list_count = run_trakt_watched()
        print(u"Number of movies in the library: {count}".format(
            count=list_count))

    print(u"Done!")