#!/usr/bin/env python2

# rtmpdump parameter generator for justin.tv/twitch.tv streams v6
# * Usage: jtvrtmp.py channelname [quality]
#   where channelname is the channel name ([twitch/justin].tv/channelname)
#   and quality is an optional parameter with a valid quality setting 
#   (360p, 480p, 720p, etc) If quality isn't present, it selects 'live'.
# * if quality is 'live', it picks the quality setting with a +, which
#   is a restream of the video being sent to jtv, no transcoding involved.
#
# Changelog:
# v6: Added documentation!
# v5: Changed to default to 'live' quality instead of 360p, as some streams
#     don't have 360p transcoding; made the code in main() slightly cleaner,
#     possibly shorter.
# v4: Add shebang line
# v3: Removed unicode copyright symbol
# v2: Added exception catching for when a stream is offline or the specified
#     quality setting is unavailable.
# v1: Initial release
#
# Copyright 2013 Steven Smith (blha303). All Rights Reserved.
# New BSD license
# http://www.opensource.org/licenses/BSD-3-Clause

from urllib2 import urlopen
import json
from sys import argv, exit, stderr
import traceback

def getswfaddr(channelname):
  """ Follows the redirect to get the channel player url, strips out the useragent """
  return urlopen("http://www.justin.tv/widgets/live_embed_player.swf?channel=" + channelname).geturl().split("&userAgent")[0]

def getstreaminfo(channelname):
  """ returns a dict of {quality: data}. Example output in streaminfooutput.txt """
  data = json.loads(urlopen("http://usher.justin.tv/find/%s.json?type=any" % channelname.lower()).read())
  newd = {}
  for i in data:
    newd[i['display']] = i
  return newd

def buildcmdline(channelname, data=None, quality="live"):
  """ Builds and returns rtmpdump parameter string for loading a twitch/justin.tv channel.
      Usage: buildcmdline(channelname as string, data dict from usher.justin.tv, formatted as {quality: data}
      Check getstreaminfo() for how to create this dict, or let buildcmdline call it by not setting data.     """
  if not data:
    data = getstreaminfo(channelname)
  if quality == "live":
    for i in data:
      if "Source" in i:
        quality = i
  try:
    data = data[quality]
  except KeyError:
    return '-q; echo "-------------"; echo "Couldn\'t find stream info for %s, maybe the stream is offline?"; echo "-------------" #' % quality
  try:
    if not "live-cdn" in data["connect"] and not "justintvlivefs" in data["connect"]:
      justinlegacyparams = '-j "%s" ' % data["token"].replace('"', r'\"')
    else:
      justinlegacyparams = ""
  except KeyError:
    qualities = []
    for i in data.keys():
      stderr.write(str(i))
      stderr.write(str(data))
      if "connect" in data[i]:
        qualities.append(i)
    return '-q; echo "-------------"; echo "Couldn\'t find stream info for %s, maybe you need to be subscribed to watch that resolution? Try one of these: ' + ", ".join(qualities) + '"; echo "-------------" #' % quality
  out = '-r "%s/%s" %s--swfVfy "%s" -v -o -' % (data["connect"], data["play"], justinlegacyparams, getswfaddr(channelname))
  return out

def main(args):
  """ Returns rtmpdump parameter string for channelname in specified list;
      argv is a list, either [channelname] or [channelname, quality]       """
  if len(args) < 2:
    return "Usage: %s channelname [quality]" % args[0]
  elif len(args) > 2:
    return buildcmdline(args[1], quality=args[2])
  else:
    return buildcmdline(args[1])

if __name__ == "__main__":
  try:
    out = main(argv)
    print out
    if "Usage:" in out:
      exit(2)
  except:
    print '-q; echo "-------------"; echo "Error:"#'
    traceback.print_exc()