#! /usr/bin/python
# simple python script to sort torrents
# modifications by xkonni
# - added folder sorting
# - added htdigest authorisation
# - fixed some minor glitches
#
# original version by jonassw from
# http://forum.xbmc.org/showthread.php?t=60749

# << DOCUMENTATION
#
# i) lighttpd configuration for htdigest
# server.modules += ( "mod_auth" )
# auth.backend = "htdigest"
# auth.backend.htdigest.userfile = "/etc/lighttpd/auth"
# auth.debug = 2
# auth.require = ( "/RPC2" =>
# (
# "method" => "digest",
# "realm" => "REALM",
# "require" => "valid-user"
# )
# )
# ii) i recommend starting this with a basic cron job you may find other
# suggestions like when rtorrent finishes hashing, but for me this caused
# lockups when starting rtorrent (as all downloads return hash_ok)
#
# */5 * * * * for i in /mnt/torrent/tv/*; do torrentSort.py $i > /dev/null; done
#
# DOCUMENTATION >>


import xmlrpclib, os, sys, re, shutil


class HTTPSDigestTransport(xmlrpclib.SafeTransport):
  """
Transport that uses urllib2 so that we can do Digest authentication.
Based upon code at http://bytes.com/topic/python/answers/509382-solution-xml-rpc-over-proxy
"""

  def __init__(self, username, pw, realm, verbose = None, use_datetime=0):
   self.__username = username
   self.__pw = pw
   self.__realm = realm
   self.verbose = verbose
   self._use_datetime = use_datetime

  def request(self, host, handler, request_body, verbose):
    import urllib2

    url='https://'+host+handler
    if verbose or self.verbose:
        print "ProxyTransport URL: [%s]"%url

    request = urllib2.Request(url)
    request.add_data(request_body)
    # Note: 'Host' and 'Content-Length' are added automatically
    request.add_header("User-Agent", self.user_agent)
    request.add_header("Content-Type", "text/xml") # Important

    # setup digest authentication
    authhandler = urllib2.HTTPDigestAuthHandler()
    authhandler.add_password(self.__realm, url, self.__username, self.__pw)
    opener = urllib2.build_opener(authhandler)

    #proxy_handler=urllib2.ProxyHandler()
    #opener=urllib2.build_opener(proxy_handler)
    f=opener.open(request)
    return(self.parse_response(f))


def adoptionCandidates(basedir, filename):
  dirs = filter(lambda x : os.path.isdir(os.path.join(basedir, x)), os.listdir(basedir))

  #set filename to lowercase for string comparisons
  filename=filename.lower()

  ignoredPhrases = ['-','_']

  candidates = []
  for dir in dirs:
    dirParts = dir.split()
    score = 0
    requiredScore = 0

    for part in dirParts:
      if ignoredPhrases.count(part) > 0:
        continue
      requiredScore = requiredScore + 1

      #force lower case for string comparison.
      part=part.lower()
      # replace "'" with "" and add word to list
      repPart = part.replace('\'','')
      if repPart != part:
        dirParts.append(repPart)
        requiredScore -= 1
      if filename.find(part) >= 0:
        score = score + 1
    if score == requiredScore:
      candidates.append( (os.path.join(basedir, dir), score) )

  return candidates

def getSeasonNumber(filename):
  patterns = [
    '.*S(\d+)E(\d+).*',
    '.*S(\d+)(\.)?E(\d+).*', # hopefully matches series.s01.e05.avi
    '(\d+)x(\d+).*',
    '(\d+)(\d+)(\d+).*'
    # commented out regex thought below Season was '4' not 14. Def better way of doing that
    #top_gear.14x04.720p_hdtv_x264-fov.mkv
    #'.*(\d+)x(\d+).*'
  ]

  for pattern in patterns:
    p = re.compile(pattern, re.I)
    g = p.findall(orphanFile)
    if len(g) > 0:
      season = int(g[0][0])
      return season
  return None


def getRtorrentId(filename):
  downloads = rtorrent.download_list('')
  for dl in downloads:
    rfile = rtorrent.d.get_base_filename(dl)
    if rfile == filename:
      return dl

# << CONFIGURATION
# i) FOLDER SETTINGS
showLocations = ['/mnt/media/tv']
allowedSourceLocation = '/mnt/media/torrent/complete/tv'

# ii) CONNECTION SETTINGS
# - using htdigest authorisation
digestTransport = HTTPSDigestTransport("username", "password", "realm")
rtorrent = xmlrpclib.ServerProxy('http://localhost',transport=digestTransport)
# - not using authorisation
# rtorrent = xmlrpclib.ServerProxy('http://localhost')

# CONFIGURATION >>

print '--------------- BEGIN ---------------'
orphanFile = sys.argv[1]
if os.path.isdir(orphanFile):
  (fpath, fname) = os.path.split(orphanFile)
  fname = os.path.relpath(orphanFile, allowedSourceLocation)

else:
  (fpath, fname) = os.path.split(orphanFile)
print 'File path: %s' % fpath
print 'File name: %s' % fname

candidates = []

if not orphanFile.startswith(allowedSourceLocation):
        print 'STOP! This file is not located in %s' % allowedSourceLocation
        exit()

print 'Attempting to find a home for file %s' % orphanFile

for location in showLocations:
        candidates.extend(adoptionCandidates(location, fname))

candidates.sort(lambda (da, sa), (db, sb): sb-sa)

if len(candidates) <= 0:
        print 'No one wanted this file :('
        exit()

for (dir, score) in candidates:
       print 'Candidate: %s with score %i' % (dir, score)

print 'Winner is %s with score %i' % candidates[0]

if os.path.isdir(orphanFile):
  finaldir = candidates[0][0]

else:
  # Determine Season and Episode number
  season = getSeasonNumber(fname)
  if not season:
    print 'STOP! Season could not be determined.'
    exit()

  print 'Season was determined to be %i' % season
  finaldir = os.path.join(candidates[0][0], 'Season %s' % season)

  # Check if season folder is present
  if not os.path.isdir(finaldir):
    print 'Season dir doesn\'t exist. Creating now'
    os.mkdir(finaldir)

  if os.path.isfile(os.path.join(finaldir, fname)):
    print 'error: file already exists, exiting'
    sys.exit(1)

  print 'Will move file to %s' % finaldir


print 'Requesting id from rtorrent'
rid = getRtorrentId(fname)
print '%s was resolved to rtorrent id: %s' % (fname, rid)

print 'Pausing rtorrent'
rtorrent.d.pause(rid)

print 'Updating rtorrent'
rtorrent.d.set_directory(rid, finaldir)

print 'Moving file'
shutil.move(orphanFile, finaldir)

print 'Resuming rtorrent'
rtorrent.d.resume(rid)