#!/usr/bin/env python

# -*- coding: utf-8 -*-

u"""

Reads an RSS/Atom feed and converts the enclosures to AVI.

Dependencies:

- ``flvstreamer`` for the RTMP streams

- ``mimms`` for the MMS streams

- ``file`` to get the sizes of the videos

- ``mencoder`` to do the conversion

- ``tedtalksubs.py`` to dowload ted talks subtitles (in this repo)

:Authors:

    Aurélien Bompard <aurelien@bompard.org> <http://aurelien.bompard.org>

:License:

    GNU GPL v3 or later

"""

import os

import sys

import urllib2

import httplib

import glob

import subprocess

import re

import tempfile

import atexit

from urlparse import urlparse

from optparse import OptionParser

from stat import S_IRUSR, S_IWUSR, S_IROTH, S_IRGRP

#from xml.etree import ElementTree as ET

from pprint import pprint

from lxml import etree as ET

import urlgrabber

import urlgrabber.progress

# Tags to skip

EXCLUDE_TAGS = ""

# Max size of the encoded video

WIDTH = 800

HEIGHT = 480

# Default MIME type

MIME_DEFAULT = "video/x-msvideo"

EXTENSION = "avi"

def get_options():

    usage = "usage: %prog -i input_feed -u URL -o output_feed [-d directory]"

    parser = OptionParser(usage=usage)

    parser.add_option("-i", "--input", dest="input",

                      help="Process this file")

    parser.add_option("-o", "--output", dest="output",

                      help="Write the RSS in this file")

    parser.add_option("-d", "--directory", dest="directory",

                      help="Write the converted videos in this directory")

    parser.add_option("-u", "--url", dest="url",

                      help="The external URL of the video folder")

    parser.add_option("-W", "--width", dest="width", type="int", default=WIDTH,

          help="Width of the converted video [default: %default]")

    parser.add_option("-H", "--height", dest="height", type="int", default=HEIGHT,

          help="Height of the converted video [default: %default]")

    parser.add_option("-m", "--max", dest="max", type="int", default=10,

          help="Only convert that many videos, drop the rest [default: %default]")

    parser.add_option("-q", "--quiet", dest="quiet", action="store_true",

                      default=False, help="Don't show progress bars")

    parser.add_option("-k", "--keep", dest="keep", action="store_true",

                      default=False, help="Don't remove original files")

    parser.add_option("--exclude-tags", dest="exclude_tags",

                      default=EXCLUDE_TAGS, help="Drop videos tagged with a "

                      "tag in this comma-sparated list [default: %default]")

    parser.add_option("--subtitles", dest="subtitles", metavar="LANG",

                      help="Download subtitles in this language")

    parser.add_option("--old-ffmpeg", dest="oldffmpeg", action="store_true",

                      help="FFMpeg is old (like on Debian Lenny)")

    options, args = parser.parse_args()

    if len(args) > 0:

        parser.error("illegal arguments: %s" % ", ".join(args))

    if not options.input:

        parser.error("I need a file to process")

    if not os.path.exists(options.input):

        parser.error("The file to process does not exist")

    if not options.output:

        parser.error("I need a file to write to")

    if not options.url:

        parser.error("I need an external URL")

    if not options.directory and options.input != "-":

        options.directory = os.path.abspath(os.path.dirname(options.input))

    if not options.directory:

        parser.error("I need a directory for the videos")

    if not options.quiet and "TERM" not in os.environ:

        options.quiet = True # Not in a terminal, be quiet anyway

    if isinstance(options.exclude_tags, basestring):

        options.exclude_tags = [ t.strip() for t in

                                 options.exclude_tags.split(",") ]

    return options, args

class PodcastError(Exception): pass

class NotAPodcastError(PodcastError): pass

class TranscodingError(PodcastError): pass

class DownloadingError(PodcastError): pass

class Podcast(object):

    _mimetypes = {}

    def __init__(self, item):

        self.item = item

        self.enclosure = item.find("enclosure")

        if not ET.iselement(self.enclosure):

            raise NotAPodcastError()

        self.url = self._get_url()

        self.path_downloaded = self._get_downloaded_path()

        self.content_type = self._get_content_type()

        self.path_encoded = self._get_encoded_path()

        self.subs = self._get_subtitles()

        self.video_info = None

        self.size = None

    def _get_url(self):

        url = self.enclosure.get("url")

        if url.startswith(options.url):

            return url

        # Resolve redirects

        try:

            remote_file = urllib2.urlopen(url)

            url = remote_file.geturl()

            self.content_type = remote_file.info().get("Content-Type")

            remote_file.close()

        except (urllib2.HTTPError, httplib.HTTPException), e:

            print "Failed downloading %s" % url

            print e

        except urllib2.URLError, e:

            print "Probably RTMP or MMS: %s" % url

        return url

    def _get_content_type(self):

        if getattr(self, "content_type", None):

            return self.content_type # already set by _get_url()

        if "content-type" in self.enclosure.attrib:

            return self.enclosure.attrib.get("content-type")

        for extension, mimetype in self.mimetypes.iteritems():

            if self.path_downloaded.endswith("."+extension):

                return mimetype

        return MIME_DEFAULT

    def _get_downloaded_path(self):

        url_obj = urlparse(self.url)

        filename = os.path.basename(url_obj.path)

        if filename.count("?"):

            filename = filename[:filename.index("?")]

        return os.path.join(options.directory, filename)

    def _get_encoded_path(self):

        if self.content_type.startswith("audio/"):

            # the file won't be transcoded

            return self.path_downloaded

        filename_encoded = "%s.%s" % (

                os.path.splitext(self.path_downloaded)[0], EXTENSION)

        return os.path.join(options.directory, filename_encoded)

    def _get_subtitles(self):

        if not options.subtitles:

            return None

        subfile, subfile_path = tempfile.mkstemp(prefix="podcast-transcode-sub-",

                                                 suffix=".srt")

        os.close(subfile)

        atexit.register(os.remove, subfile_path)

        if (flux_xml.findtext("channel/title") == "TEDTalks (video)" or

                flux_xml.findtext("channel/title") == "TEDTalks (hd)"):

            talkid = self.item.findtext("guid").split(":")[1]

            subdl = subprocess.Popen(["tedtalksubs.py", "-l", options.subtitles,

                                      "-o", subfile_path, talkid],

                                     stdout=subprocess.PIPE,

                                     stderr=subprocess.STDOUT)

            out, err = subdl.communicate()

            if subdl.returncode != 0:

                print >>sys.stderr, "Failed to download subtitles. Message:"

                print >>sys.stderr, out

                return None

            return subfile_path

        return None

    def _get_mimetypes(self):

        if self._mimetypes:

            return self._mimetypes

        mimetypes_re = re.compile("\s*([^\s]+)\s+([^\s]+)\s*")

        mimetypes = open("/etc/mime.types")

        for line in mimetypes:

            line_mo = mimetypes_re.match(line)

            if not line_mo:

                continue

            mimetype = line_mo.group(1)

            extension = line_mo.group(2)

            if not mimetype.startswith("video/") \

                    and not mimetype.startswith("audio/"):

                continue

            self._mimetypes[line_mo.group(2)] = line_mo.group(1)

        mimetypes.close()

        return self._mimetypes

    mimetypes = property(_get_mimetypes)

    def is_already_transcoded(self):

        return self.url.startswith(options.url)

    def process(self):

        if not os.path.exists(self.path_encoded):

            self.download()

            self.encode_video()

        else:

            print "Already downloaded/encoded: %s" % self.path_encoded

        self.url = "%s/%s" % (options.url, os.path.basename(self.path_encoded))

        self.size = int(os.stat(self.path_encoded).st_size)

        self.update_item()

    def update_item(self):

        self.enclosure.set("url", self.url)

        self.enclosure.set("length", str(self.size))

        self.enclosure.set("type", self.content_type)

        fb = self.item.find("{http://rssnamespace.org/feedburner/ext/1.0}origEnclosureLink")

        if ET.iselement(fb):

            fb.text = self.url

        mediacontent = self.item.find("{http://search.yahoo.com/mrss/}content")

        if ET.iselement(mediacontent):

            mediacontent.set("url", self.url)

            mediacontent.set("fileSize", str(self.size))

            mediacontent.set("type", self.content_type)

    def download(self):

        if os.path.exists(self.path_downloaded):

            print "Already downloaded: %s" % self.path_downloaded

            return

        if self.url.startswith("rtmp://"):

            self.download_rtmp()

        elif self.url.startswith("mms://"):

            self.download_mms()

        elif self.content_type and self.content_type == "video/x-ms-asf":

            self.download_asf()

        else:

            if options.quiet:

                progress = urlgrabber.progress.BaseMeter()

            else:

                progress = urlgrabber.progress.TextMeter(fo=sys.stdout)

            print "Downloading %s to %s" % (self.url, self.path_downloaded)

            try:

                urlgrabber.urlgrab(self.url, filename=self.path_downloaded,

                                   reget='simple', progress_obj=progress)

            except urlgrabber.grabber.URLGrabError, e:

                raise DownloadingError("Error downloading %s: %s"

                                       % (self.url, e))

    def download_rtmp(self):

        MAX_TRIES = 10

        def download_rtmp_unit(url, path):

            command = ["flvstreamer", "-r", url, "-o", path]

            if options.quiet:

                command.append("-q")

            if os.path.exists(path):

                command.insert(1, "--resume")

            print "Streaming %s to %s" % (url, path)

            retcode = 0

            try:

                retcode = subprocess.call(command)

            except KeyboardInterrupt:

                retcode = 1

            return retcode

        retcode = download_rtmp_unit(self.url, self.path_downloaded)

        # flvstreamer returns 2 if the download is incomplete

        current_try = 1

        while retcode == 2:

            print "Trying again..."

            retcode = download_rtmp_unit(self.url, self.path_downloaded)

            current_try += 1

            if current_try > MAX_TRIES:

                print "Too many tries, aborting."

                break

        if retcode != 0:

            if os.path.exists(self.path_downloaded):

                os.remove(self.path_downloaded)

            raise DownloadingError("Error code: %s" % retcode)

    def download_mms(self):

        #command = ["mplayer", "-dumpstream", "-dumpfile", self.path_downloaded, self.url]

        command = ["mimms", self.url, self.path_downloaded]

        if options.quiet:

            command.append("-q")

        try:

            print "Streaming %s to %s" % (self.url, self.path_downloaded)

            retcode = 0

            retcode = subprocess.call(command)

        except KeyboardInterrupt:

            retcode = 1

        if retcode != 0 and os.path.exists(self.path_downloaded):

            os.remove(self.path_downloaded)

            raise DownloadingError("Error code: %s" % retcode)

    def download_asf(self):

        mms_xml = urllib2.urlopen(self.url).read()

        mms_match = re.search('"(mms://.*)"', mms_xml)

        mms_url = mms_match.group(1)

        return download_mms(mms_url.replace("&", "&"), self.path_downloaded)

    def encode_video(self):

        if self.path_encoded == self.path_downloaded:

            print "No transcoding required"

            return

        transcoded_video = self._transcode_video()

        os.rename(transcoded_video, self.path_encoded)

        os.chmod(self.path_encoded, S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH) # 644

    def _transcode_video(self):

        width, height = self.get_video_info()

        transcoded_video_file, transcoded_video = tempfile.mkstemp(

                prefix="podcast-transcode-", suffix=".avi", dir=options.directory)

        os.close(transcoded_video_file)

        def rm_if_exists(f):

            if os.path.exists(f):

                os.remove(f)

        if not options.keep:

            atexit.register(rm_if_exists, transcoded_video)

        command = ["mencoder", "-oac", "mp3lame",

                   "-ovc", "lavc", "-lavcopts", "vbitrate=600",

                   "-of", "avi", "-mc", "0", self.path_downloaded,

                   "-o", transcoded_video]

        if height and width:

            if int(height) > options.height:

                command.extend(["-vf", "scale=-3:%d" % options.height])

            elif int(width) > options.width:

                command.extend(["-vf", "scale=%d:-3" % options.width])

        if options.quiet:

            command.append("-quiet")

        if self.subs:

            command.extend(["-sub", self.subs, "-subfont-text-scale", "4"])

            if not os.path.exists(os.path.expanduser("~/.mplayer/subfont.ttf")):

                command.extend(["-fontconfig", "-font", "DejaVu Sans"])

        print " ".join(command)

        retcode = 0

        try:

            print "Encoding %s to %s" % (self.path_downloaded, transcoded_video)

            retcode = subprocess.call(command)

        except KeyboardInterrupt:

            retcode = 1

        if retcode != 0:

            if os.path.exists(transcoded_video):

                os.remove(transcoded_video)

            raise TranscodingError("Error code: %s" % retcode)

        self.content_type = "video/x-msvideo"

        return transcoded_video

    def get_video_info(self):

        if self.video_info is not None:

            return self.video_info

        ffmpeg_cmd = subprocess.Popen(["ffmpeg", "-i", self.path_downloaded],

                        stdout=subprocess.PIPE, stderr=subprocess.STDOUT)

        output = ffmpeg_cmd.stdout.read()

        info_match = re.search("Stream .*: Video: (\w+), \w+, (\d+)x(\d+)", output)

        if info_match:

            self.video_info = (info_match.group(2), info_match.group(3))

        else:

            self.video_info = (None, None)

        return self.video_info

def cleanup(items):

    feed_podcasts = set()

    for item in items:

        try:

            podcast = Podcast(item)

        except NotAPodcastError:

            continue

        feed_podcasts.add(os.path.basename(podcast.path_downloaded))

        feed_podcasts.add(os.path.basename(podcast.path_encoded))

    if options.keep:

        return

    for filepath in glob.glob(os.path.join(options.directory, "*")):

        if filepath.endswith(".xml"):

            continue # keep the RSS feed

        filename = os.path.basename(filepath)

        if filename not in feed_podcasts:

            print "Removing old file %s" % filename

            #print feed_podcasts

            os.remove(filepath)

def handle_item(item):

    try:

        podcast = Podcast(item)

    except NotAPodcastError:

        return

    if podcast.is_already_transcoded():

        print "Already converted: %s" % podcast.url

        return

    if options.subtitles and not podcast.subs:

        print "No subtitles for %s, skipping." % item.findtext("guid")

        flux_xml.find("channel").remove(item)

        return

    try:

        podcast.process()

    except DownloadingError, e:

        print e

        return

    except TranscodingError, e:

        print e

        return

def to_skip(item):

    tags = item.findall("category")

    for tag in tags:

        if tag.text in options.exclude_tags:

            return True

    return False

def main():

    global options, flux_xml

    options, args = get_options()

    if options.input == "-":

        options.input = sys.stdin

    flux_xml = ET.parse(options.input)

    items = flux_xml.findall("channel/item")

    # tag skipping

    for item in items[:]:

        if to_skip(item):

            flux_xml.find("channel").remove(item)

            items.remove(item)

    for i, item in enumerate(items[:]):

        if i < options.max:

            handle_item(item)

        else:

            flux_xml.find("channel").remove(item)

    #flux_xml.write(options.output, "utf-8")

    flux_xml.write(options.output)

    cleanup(items[:options.max])

if __name__ == "__main__":

    main()