#!/usr/bin/python

## Small Python script to retrieve XMLTVID's from Zap2It's
## DataDirect service for US cable and broadcast channels
##
## Author: Mitko Haralanov (voidtrance@comcast.net)
##
## This package relies on the ClientCookie module by
## John J. Lee (http://wwwsearch.sourceforge.net/ClientCookie)
##
## This script was put together in a very short time without
## much though for robustness and "pretty coding", so it might
## crash or do the wrong thing from time to time.
## I did test it with several listing from several different
## areas in the US and it seem to work. YMMV!!

import sys
import urllib2
import urllib
import ClientCookie
import getopt
from HTMLParser import HTMLParser
from os.path import basename

SERVER = "labs.zap2it.com"
LOGIN_PAGE = "/ztvws/ztvws_login/1,1059,TMS01-1,00.html"

class ZapParser (HTMLParser):
    def __init__ (self):
        HTMLParser.__init__ (self)
        self.forms = []
        self.labels = []
        self.td1 = None
        self.td2 = None
        self.current_tag = None
        self.label_index = -1
        
    def handle_starttag (self, tag, attrs):
        self.current_tag = tag
        if tag == "form":
            self.forms.append (self.todict (attrs))
            self.forms[-1]["fields"] = []
            self.forms[-1]["name"] = self.td1
            self.forms[-1]["type"] = self.td2
            self.td1, self.td2 = None, None
        if tag == "input":
            attrs = self.todict (attrs)
            if attrs["type"] not in  ["submit", "reset", "button"]:
                self.forms[-1]["fields"].append (attrs)
        if tag == "label":
            attrs = self.todict (attrs)
            if attrs.has_key ("for"):
                self.labels.append ({"id":attrs["for"], "channel":None, "callsign":None})
                self.label_index = len (self.labels)-1
            else:
                self.label_index = -1
        
    def handle_endtag (self, tag):
        if self.current_tag == tag:
            self.current_tag = None
    
    def handle_data (self, data):
        if self.current_tag == "td":
            if self.label_index == -1:
                if not self.td1 or (self.td1 and self.td2): self.td1 = data.strip ()
                else: self.td2 = data.strip ()
            else:
                if not self.labels[self.label_index]["channel"]:
                    self.labels[self.label_index]["channel"] = data
                else:
                    self.labels[self.label_index]["callsign"] = data
            
    def todict (self, data):
        new = {}
        for elem in data:
            new[elem[0]] = elem[1]
        return new

    def get_form_data (self):
        return self.forms

    def get_label_data (self):
        return self.labels
    
def usage (msg=None):
    if msg: print msg+"\n"

    name = basename (sys.argv[0])
    print name, "[-u/--username <username>] [-p/--password <password>] [-i/--interactive]"
    print " "*len (name), "[-h/--help]"
    print
    print "  -u/--username    - The username with which to login to Zap2It"
    print "  -p/--password    - The password with which to login to Zap2It"
    print "  -i/--interactive - Enables interactive mode. In this mode, instead of printing"
    print "                     the information on all channels, the program will ask for a"
    print "                     specific channel callsign and do a look up on it."
    print
    
username = None
password = None
interactive = False

try: opts, args = getopt.getopt (sys.argv[1:], "u:p:ih", ["username=", "password=", "interactive",
                                                          "help"])
except getopt.GetoptError, e:
    usage (e)
    sys.exit (1)

for opt, value in opts:
    if opt in ("-u", "--username"):
        username = value
    if opt in ("-p", "--password"):
        password = value
    if opt in ("-i", "--interactive"):
        interactive = True
    if opt in ("-h", "--help"):
        usage ()
        sys.exit (0)
        
if not username:
    while not username:
        username = raw_input ("Service username: ")

if not password:
    import getpass
    while not password:
        password = getpass.getpass ("Service password: ")

print "Getting Line-Up list..."
params = urllib.urlencode({"username": username, "password": password, "action":"Login"})
headers = {"Content-type": "application/x-www-form-urlencoded",
           "Accept": "text/plain"}

response = ClientCookie.urlopen ("http://" + SERVER + LOGIN_PAGE, params)
data = response.read ()
p = ZapParser ()
p.feed (data)
lineups = p.get_form_data ()
print "Select from available lineups:"
for lineup in lineups:
    if not lineup["name"]: continue
    print "\t%s: %s (%s)"%(lineups.index (lineup)+1, lineup["name"], lineup["type"])
done = False    
while not done:
    sel = raw_input ("Select: ")
    try: sel = int (sel)
    except ValueError:
        print "ERROR: Please enter an integer"
        continue
    if sel > 0 and sel <= len (lineups):
        done = True
    else:
        print "ERROR: invalid selection"

lineup = lineups[sel-1]

print "Getting channel information..."
fields = {}
for field in lineup["fields"]:
    fields[field["name"]] = field["value"]
fields["submit"] = "Modify"
params = urllib.urlencode (fields)
response = ClientCookie.urlopen ("http://" + SERVER + lineup["action"], params)
data = response.read ()
p.feed (data)
id_data = p.get_label_data ()
header = "\n%-10s %-7s %-5s\n%s"%("Callsign", "Channel", "XMLTVID", "-"*30)

if not interactive:
    print "\n\nXMLTVID Data"
    print "="*30
    count = 0
    for id in id_data:
        if count % 20 == 0:
            print header
        print "%-10s %-7s %-5s"%(id["callsign"], id["channel"], id["id"][5:10])
        count += 1
else:
    done = False
    callsigns = [x["callsign"] for x in id_data]
    while not done:
        query = raw_input ("Enter callsign ('exit' when done): ")
        if query == "exit":
            done = True
            continue
        if query in callsigns:
            id = filter (lambda(x): x["callsign"] == query, id_data)
            if id:
                for i in id:
                    print header
                    print "%-10s %-7s %-5s\n"%(i["callsign"], i["channel"], i["id"][5:10])
        else: print "Callsign '%s' not in listing"%query