#!/usr/bin/python
#
#NAME:
#  rename_mkv_with_imdb.py
#
#PURPOSE:
#  To go through all of the files in the current directory 
#  and rename them with the correct imdb title so that Kodi 
#  will understand it 
#
#INPUTS:
#  None:
#    The script will operate on all the files with extensions 
#    .mkv, .mp4, .avi, .wmv, and .mov in the current directory.
#    It will also descend one level down into the subdirectories 
#    in the current directory.
#
#DEPENDENCIES:
#  This script requires the requests module, which can be 
#  downloaded using 'pip'. To get pip (and perhaps the latest 
#  version of python, follow the instructions here:
# 
#  http://python-guide-pt-br.readthedocs.io/en/latest/starting/install/osx/
#
#  Once pip is on the machine, follow these instructions:
#  
#  http://docs.python-requests.org/en/master/user/install/
#  
#  You also need to make the script executable, for instance using:
#    > cd [script directory]
#    > chmod a+x rename_mkv_with_imdb.py
#
#EXAMPLES:
#  > cd [MOVIE_DIRECTORY] 
#  /Users/lances/python_scripts/rename_mkv_with_imdb.py
#
import os
import re
import sys
import requests
rootdir = './'

print '------------------------------------------------------------------------------------'
print '------------------------------------------------------------------------------------'
print 'Looping through all files in current working directory...'
print '------------------------------------------------------------------------------------'
print '------------------------------------------------------------------------------------'

for root, directories, filenames in os.walk(rootdir):

  #Scan through the subdirectories of the rootdir 
  for filename in filenames:

    # Ask user to search on current file (if it is a valid video file)
    print ''
    print '------------------------------------------------------------------------------------'
    print 'Current file: ' + os.path.join(root, filename)
    if ((filename.find('mkv') != -1) or (filename.find('mp4') != -1) or \
        (filename.find('.avi') != -1) or (filename.find('.wmv') != -1) or \
        (filename.find('.mov') != -1)):
      print 'Search IMDB for a proper title for file : ' + filename + ' ?'
      print 'Press 1 to search; any other number to skip...'
      Choice = '0'
      while (type(Choice) != int):
        try:
          Choice = int(raw_input())
        except KeyboardInterrupt:
          sys.exit()
        except:
          print 'That is not a valid input. You must input a number between 0 and 9'
    else:
      Choice = 0

    # YES: Perform a search on the current file
    if (Choice == 1):

        # Get information on the current file 
        oldpath = os.path.join(root, filename)
        oldname = filename

        # Get the extension and the file extension
        movie_try, extension = os.path.splitext(oldname)
        print ''
        print 'Searching IMDB for the movie : ' + movie_try
 
        # Start a session
        r=requests.Session()

        # Form the URL and make the query
        url = 'http://www.imdb.com/find?ref_=nv_sr_fn&q='+movie_try+'&s=all'
        res=r.get(url)

        # Manipulate the text so that we can get the right name
        res=res.text.encode('utf8', 'replace')
        res=res.replace('</a>', '')

        # Now begin searching for titles, which come after href tags that have links  
        # ending in the sequence '?ref=fn_al_tt_#', where number is an integer index
        offset = 18
        total_titles   = 5
        this_title_num = 0
        newfilenames = []
        start = 0

        # Loop over all the matching titles
        while ((this_title_num < total_titles) and (start != -1)):

          # Look for the next instance of a title
          start = res.find('ref_=fn_al_tt_'+str(this_title_num+1), start+1)
          if (start == -1): break
          # If there was no match at all, exit
          if ((this_title_num == 0) and (start == -1)):
            print 'Could not find matching text! Skipping...'
            break
          # If there was a match, add the title to the list
          else:
            # Sometimes in IMDB, a </td> tag comes immediately after the title,
            # but sometimes it is a <br or <br/> tag
            start = res.find('ref_=fn_al_tt_'+str(this_title_num+1), start+1)
            end = res.find(' </td>', start)
            end2 = res.find(' <br', start)
            if ((end2 > start) and (end2 < end)): end = end2
            newfilenames.append(res[start+offset: end])
            this_title_num+=1

        # Allow the user to select which title he/she would like
        if (this_title_num != 0):
          print 'Top ' + str(this_title_num) + ' Matching movie titles ..........'
          for i in range(this_title_num):
            print '  Choice ' + str(i+1) + ' : ' + newfilenames[i]
          print 'Select the movie title you would like (enter ' + str(this_title_num+1) + \
                ' or above to select none of them)'
          Choice = '0'
          while (type(Choice) != int):
            try:
              Choice = int(raw_input())
            except KeyboardInterrupt:
              sys.exit()
            except:
              print 'That is not a valid input. You must input a number between 0 and 9'

          # Now that we have the proper choice, prompt the user one more time to make sure 
          # he or she wishes to make the move
          if ((Choice > 0) and (Choice <= this_title_num) and (this_title_num != 0)):
            newfilename = newfilenames[Choice-1]
            newfilename = newfilename.replace(': ', '-')
            newpath = os.path.join(root, newfilename+extension)
            print ''
            print 'Old file path is      : ' + oldpath
            print 'New file path will be : ' + newpath
            print ''
            print 'Would you like to rename the file? Press 1 for yes: any other number to skip...'
            Choice = '0'
            while (type(Choice) != int):
              try: 
                Choice = int(raw_input())
              except KeyboardInterrupt:
                sys.exit()
              except:
                print 'That is not a valid input. You must input a number between 0 and 9'
            if (Choice == 1):
              print ''
              print 'Renaming ' + oldpath + ' to ' + newpath
              os.rename(oldpath, newpath)
            else:
              print ''
              print 'Skipping...'
          else: 
            print ''
            print 'A valid choice was not selected. Skipping...'

        else:
          print 'No matching IMDB entry! Moving to next file...'
        
    # NO: we are skipping this file
    else:
      print 'Skipping file : ' + filename

# We made it all the way to the end
print ''
print 'Finished with all files and current subdirectories!'