#!/usr/bin/env python2
# coding: utf-8
# python3 works as well
import os
import sys
import re
import dbm
import json
import argparse
import math

try:
    import tmdbsimple as tmdb
except ImportError as e:
    print(u"Missing dependency: {0}".format(str(e)))
    print(u"Install using system package manager or `pip install --user <module>`")
    sys.exit(1)

def read_key():
    if u"TMDB_KEY" in os.environ.keys():
        return os.environ[u"TMDB_KEY"]
    if u"XDG_CONFIG_HOME" in os.environ.keys():
        cfg_home = os.environ[u"XDG_CONFIG_HOME"]
    else:
        cfg_home = os.path.join(os.path.expanduser(u"~"), ".config")
    if os.path.exists(os.path.join(cfg_home, u"tmdbkey")):
        return open(os.path.join(cfg_home, u"tmdbkey"), "r").read().strip()
    if os.path.exists(os.path.join(os.path.expanduser(u"~"), ".tmdbkey")):
        return open(os.path.join(os.path.expanduser(u"~"), ".tmdbkey")).read().strip()
    raise Exception(u"No TheMovieDB Key defined. Set Env. var. TMDB_KEY or .tmdbkey file")
        
class TMDBCache(object):
    def __enter__(self):
        self.db = dbm.open(self._get_db_filename(),"rw")
        return self

    def __exit__(self, type, value, traceback):
        self.db.close()

    def _get_db_filename(self):
        if u"XDG_CACHE_HOME" in os.environ.keys():
            cachedir = os.environ["XDG_CACHE_HOME"]
        else:
            cachedir = os.path.join(os.path.expanduser(u"~"), ".cache")
        return os.path.join(cachedir, "tmdbmovie.dbm")

    def _cache(self, key, callable_func):
        if not self.db.has_key(key):
            self.db[key] = json.dumps(callable_func())
        return json.loads(self.db[key])

    def infos(self, movie_id):
        try:
            return self._cache(movie_id + "movies.info", tmdb.Movies(movie_id).info)
        except Exception as e:
            raise Exception("Failed to query movie with id {id}: {reason}".format(id=movie_id, reason=str(e)))

    def alternative_title(self, movie_id, locale):
        """Returns the title in selected locale or original title otherwise"""
        try:
            alt_title = filter(lambda l: l["iso_3166_1"] == locale,
                    self._cache(movie_id + "movies.alt_titles", tmdb.Movies(movie_id).alternative_titles)["titles"])
            if alt_title:
                return alt_title[0]["title"]
            else:
                infos = self.infos(movie_id)
                return infos["title"] or infos["original_title"]
        except Exception as e:
            raise Exception("Failed to query movie with id {id}: {reason}".format(id=movie_id, reason=str(e)))
    
    def prune(self, movie_id):
        keys = [movie_id + "movies.info", movie_id + "movies.alt_titles"]
        for key in keys:
            if key in self.db:
                print "Remove {}".format(key)
                del self.db[key]
        
def do_aka(args, imdb_ids):
    with TMDBCache() as tmdbcache:
        for imdb_id in imdb_ids:
            print tmdbcache.alternative_title(imdb_id, locale=args.lang)

def do_data(args, imdb_ids):
    with TMDBCache() as tmdbcache:
        for imdb_id in imdb_ids:
            selected_properties = ["imdb_id", "revenue", "vote_average", "vote_count", "runtime", "budget", "vote_avarage", "release_date", "popularity", ]
            kv = map(lambda kv: u"{}={}".format(*kv), 
                    filter(lambda (k,v): k in selected_properties,
                            tmdbcache.infos(imdb_id).items()))
            print u" ".join(kv)

def do_year(args, imdb_ids):
    with TMDBCache() as tmdbcache:
        for imdb_id in imdb_ids:
            print tmdbcache.infos(imdb_id)["release_date"].split("-")[0]

def do_prune(args, imdb_ids):
    with TMDBCache() as tmdbcache:
        for imdb_id in imdb_ids:
            tmdbcache.prune(imdb_id)

def do_rating(args, imdb_ids):
    """Calculates a rating based on vote_average and vote_count.
    See http://blog.moertel.com/posts/2006-01-17-mining-gold-from-the-internet-movie-database-part-1.html
    for details"""
    with TMDBCache() as tmdbcache:
        infos = map(lambda imdb_id: tmdbcache.infos(imdb_id), imdb_ids)
        infos = filter(lambda i: "vote_average" in i and "vote_count" in i, infos)
        maxvotes = max(map(lambda i: i["vote_count"], infos))
        for info in infos:
            f = math.sin(math.pi * ( info["vote_average"]/10.0 ) )
            d = (float(info["vote_count"]) / maxvotes) - 0.5
            info['rating'] = info["vote_average"] + 2 * d * f
            print "{rating:.02f} {imdb_id} {title:30s} avg={vote_average:.1f} count={vote_count:.0f}".format(**info)

if __name__ == u"__main__":
    tmdb.API_KEY = read_key()

    parser = argparse.ArgumentParser(description="get movie data")
    subparsers = parser.add_subparsers()
    parser_aka = subparsers.add_parser("aka")
    parser_aka.add_argument("--lang", help="Language code (default 'DE')")
    parser_aka.set_defaults(func=do_aka)
    parser_aka.add_argument("files", action="append", nargs="+")

    parser_data = subparsers.add_parser("data")
    parser_data.set_defaults(func=do_data)
    parser_data.add_argument("files", action="append", nargs="+")

    parser_year = subparsers.add_parser("year")
    parser_year.set_defaults(func=do_year)
    parser_year.add_argument("files", action="append", nargs="+")

    parser_prune = subparsers.add_parser("prune")
    parser_prune.set_defaults(func=do_prune)
    parser_prune.add_argument("files", action="append", nargs="+")

    parser_rating = subparsers.add_parser("rating")
    parser_rating.set_defaults(func=do_rating)
    parser_rating.add_argument("files", action="append", nargs="+")

    args = parser.parse_args(sys.argv[1:])

    ids = map(lambda filename: (lambda x: x.groups()[0] if x else None)(re.match(".*#(tt[0-9]{7}).*", filename)),
                args.files[0])
    args.func(args, filter(lambda i: i is not None, ids))