#!/usr/bin/python
# coding=UTF-8

import sys
import os
import os.path
import urllib
import urllib2
import cookielib
import time
import re
import getopt
import subprocess

from cgi import parse_qs

def js_unescape(s):
    s = unicode(s).replace("\\\"", "\"").replace("\\'", "'")
    s = re.sub(r"\\u([0-9a-fA-F]{4})", lambda mo: unichr(int(mo.group(1), 16)), s)
    s = s.replace("\\\\", "\\");
    return s

class error(Exception): pass

class FakeBrowser:
    def __init__(self, quiet=True, cookie_jar_type=cookielib.LWPCookieJar):
        self.quiet = quiet

        self.cookie_jar = cookie_jar_type()

        self.opener = urllib2.build_opener(urllib2.HTTPHandler(debuglevel=not self.quiet), urllib2.HTTPCookieProcessor(self.cookie_jar))
        urllib2.install_opener(self.opener)

    def request(self, url, data=None, headers={}, return_file=False):
        if data is not None and not isinstance(data, basestring):
            data = urllib.urlencode(data)

        #headers["User-Agent"] = "Mozilla/5.0 (Windows; U; Windows NT 5.1; ja-JP; rv:1.8.1.6) Gecko/20070725 Firefox/2.0.0.6"
        headers["User-Agent"] = "Mozilla/4.0 (compatible; MSIE 6.0; Windows NT 5.1)"

        req = urllib2.Request(url, data, headers)

        try:
            f = urllib2.urlopen(req)
        except urllib2.URLError, e:
            raise Exception("Error %s while requesting '%s'" % (e.code, url), e.read())

        if return_file:
            return f

        response = f.read()
        f.close()

        return response

    def load_cookies(self, cookie_file):
        self.cookie_jar.load(cookie_file)

class NicoFetch:
    VIDEO_ID_RE = re.compile(r"(?:/|^)([a-z]{2}\d+)")
    VIDEO_TITLE_RE = re.compile(r"title:\s*'([^']*)'")
    KEY_RE = re.compile(r"thumbPlayKey:\s*'([^']*)'")
    MOVIE_TYPE_RE = re.compile(r"movieType:\s*'([^']*)'")
    #NOT_LOGGED_IN_RE = re.compile(r"<form[^>]*?id=\"login\"")
    
    LEVEL_QUIET = 0
    LEVEL_INFO = 1
    LEVEL_DEBUG = 2

    def __init__(self, msg_level=1):
        self.msg_level = msg_level
        self.browser = FakeBrowser(quiet=not (msg_level >= self.LEVEL_DEBUG))

    def fetch(self, video_id, output_dir="", output_file="", comments=True, audio=False, audio_only=False, ass=False):
        if audio_only:
            audio = True

        mo = self.VIDEO_ID_RE.search(video_id)
        if mo is None:
            raise error("No video ID found from video_id")
        
        video_id = mo.group(1)
        self.video_id = video_id

        self.log(self.LEVEL_INFO, "Downloading %s" % video_id)

        # See if we can find a login cookie

        thumb_data = self.request("http://ext.nicovideo.jp/thumb_watch/" + video_id, message="Requesting thumb_watch page",
            headers={"Referer": "http://dic.hatena.ne.jp/"})

        mo = self.KEY_RE.search(thumb_data)
        if mo is None:
            raise error("Error parsing thumb_watch result")
        key = mo.group(1)

        mo = self.MOVIE_TYPE_RE.search(thumb_data)
        if mo is None:
            raise error("Error parsing thumb_watch result")
        ext = "." + mo.group(1)
        self.ext = ext

        if output_file == "":
            mo = self.VIDEO_TITLE_RE.search(thumb_data)
            if mo is None:
                self.log(self.LEVEL_INFO, "Couldn't find video title (outdated version)? Using ID as filename...")
                output_file = video_id
                self.video_title = video_id
            else:
                self.video_title = js_unescape(mo.group(1))
                output_file = self.video_title

        self.output_file = output_file

        video_data = self.request("http://ext.nicovideo.jp/thumb_watch/" + video_id + "/" + key, message="Requesting thumb_watch/key page")

        try:
            getflv_values = parse_qs(video_data)

            thread_id = getflv_values["thread_id"][0]
            flv_url = getflv_values["url"][0]
            message_url = getflv_values["ms"][0]
        except:
            raise error("Something went wrong parsing the video_data result")

        if flv_url[-3:] == "low":
            self.log(self.LEVEL_INFO, "NOTICE: Economy mode is in use. Quality will not be as good as possible.")

        output_path = os.path.expanduser(os.path.join(output_dir, output_file))

        self.log(self.LEVEL_INFO, "Saving to %s, etc." % (output_path + ext))
        
        if comments:
            self.log(self.LEVEL_INFO, "Downloading comments...")

            comments_data = "<thread no_compress=\"0\" user_id=\"0\" when=\"0\" waybackkey=\"0\" res_from=\"-1000\" version=\"20061206\" thread=\"%s\" />" % (thread_id)
            comments_result = self.request(message_url, data=comments_data, message="Fetching comments XML...")

            f = open(output_path + ".xml", "w")
            f.write(comments_result)
            f.close()

        self.log(self.LEVEL_INFO, "Downloading video...")

        flv_file = self.request(flv_url, message="Requesting video...", return_file=True)

        if flv_file is None:
            raise error("Couldn't download video file")

        f = open(output_path + ext, "w")

        contentLength = int(flv_file.info().get("Content-Length", 1))

        self.download_to_file(flv_file, f, contentLength)

        f.close()

        self.log(self.LEVEL_INFO, "Download done.")

        if audio:
            self.log(self.LEVEL_INFO, "Extracting audio")
            self.extract_audio(output_path + ext, output_path + ".mp3")

        if audio_only:
            self.log(self.LEVEL_INFO, "Audio only, removing video file")
            os.remove(output_path + ext)

        self.log(self.LEVEL_INFO, "All done (〓ω〓・)")

        return True
                
    def request(self, url, data=None, level=LEVEL_DEBUG, message=None, headers={}, return_file=False):
        self.log(level, message)
        return self.browser.request(url, data=data, headers=headers, return_file=return_file)

    def log(self, level, s, newline=True):
        if level <= self.msg_level:
            if newline:
                print s
            else:
                print s,

    def extract_audio(self, in_path, out_path):
        rv = subprocess.call(["mplayer", "-really-quiet", "-dumpaudio", "-dumpfile", out_path, in_path])
        #rv = subprocess.call(["ffmpeg", "-i", "-dumpaudio", "-dumpfile", out_path, in_path])
        
        return rv == 0

    def download_to_file(self, in_file, out_file, size):
        start_time = time.clock()
        total_read = 0
        next_update = start_time
        format = "Downloading: [%s] %s/%s kB"
        bar_length = 40

        while True:
            data = in_file.read(10 * 1024)
            total_read += len(data)
            ctime = time.time()
            if ctime >= next_update or total_read == size:
                next_update = ctime + 0.5
                bar_fill = int((float(total_read) / float(size)) * float(bar_length))
                self.log(self.LEVEL_INFO,
                    "\x1B[2K\x1B[1000D" + format % ("=" * bar_fill + "-" * (bar_length - bar_fill), int(total_read / 1024), int(size / 1024)), False)
                sys.stdout.flush()
            if len(data) == 0:
                break
            out_file.write(data)

def usage(status):
    print "Usage: nicofetch.py [OPTIONS] VIDEO_ID"
    print
    print "  -f, --output-file  set output file prefix (extensions will be appended)"
    print "                     default: use video title"
    print "  -d, --output-dir   set output directory (filename will be appended)"
    print "                     default: current working directory"
    print "  -a, --audio        extract audio"
    print "  -A, --audio-only   delete video file after extracting audio"
    print "  -C, --no-comments  don't fetch comments"
    print "  -s, --subtitles    generate .ass subtitles from comments (WIP)"
    print "  -v, --verbose      output lots of debug information"
    print "  -q, --quiet        quiet operation"
    print "  -h, --help         display this help and exit"
    print
    print "VIDEO_ID can be any string that contains a part that looks like"
    print "a video id, for example a video URL."
    print "Audio extraction requires an mplayer installation that can handle"
    print "FLV files, mplayer also has to be in $PATH."

    sys.exit(status)

if __name__ == "__main__":
    config = {
        "output-dir": "",
        "output-file": "",
        "audio": False,
        "audio-only": False,
        "comments": True,
        "subtitles": False,
        "msg-level": 1,
    }
    
    config_filename = os.path.expanduser("~/.nicofetchrc")

    if os.path.exists(config_filename):
        f = open(config_filename, "r")
        if f:
            config.update(eval(f.read()))
            f.close()

    try:
        opts, args = getopt.getopt(sys.argv[1:], "f:d:e:p:aACsvqh",
            ["output-file", "output-dir", "audio",
             "audio-only", "no-comments", "subtitles", "verbose",
             "quiet", "help"])
    except getopt.GetoptError, err:
        print str(err)
        usage(2)

    video_id = "".join(args)

    if video_id == "":
        print "no video ID given"
        usage(2)
    
    for opt, arg in opts:
        if opt in ("-f", "--output-file"):
            config["output-file"] = arg
        elif opt in ("-d", "--output-dir"):
            config["output-dir"] = arg
        elif opt in ("-a", "--audio"):
            config["audio"] = True
        elif opt in ("-A", "--audio-only"):
            config["audio-only"] = True
        elif opt in ("-C", "--no-comments"):
            config["comments"] = False
        elif opt in ("-s", "--subtitles"):
            config["ass"] = True
        elif opt in ("-v", "--verbose"):
            config["msg-level"] = 2
        elif opt in ("-q", "--quiet"):
            config["msg-level"] = 0
        elif opt in ("-h", "--help"):
            usage(0)
    
    nf = NicoFetch()
    
    try:
        nf.fetch(video_id, output_dir=config["output-dir"], output_file=config["output-file"],
            comments=config["comments"], audio=config["audio"], audio_only=config["audio-only"],
            ass=config["subtitles"])
    except error, e:
        print "error: %s" % str(e)
        sys.exit(1)

    sys.exit(0)

# fetch(video_id, output_dir="", output_filename="", comments=True, audio=False, audio_only=False, ass=False):