#!/usr/bin/python

# By hads@nice.net.nz
# Released under the BSD license

import sys
import pytz
import urllib
from time import strptime, strftime, mktime
from datetime import datetime
from xml.etree import cElementTree as ElementTree

NAME = 'epgsnoop-grabber'
VERSION = '0.2.1'
URL = 'http://freeviewnz.tv/epg_data.php'
IN_TIME_FORMAT = '%d/%m/%Y %H:%M:%S'
OUT_TIME_FORMAT = '%Y%m%d%H%M%S %z'
NZTIME = pytz.timezone('Pacific/Auckland')
XMLTVIDS = {
    'TV ONE': 'tv1.freeviewnz.tv',
    'TV One': 'tv1.freeviewnz.tv',
    'TV2': 'tv2.freeviewnz.tv',
    'TV 3': 'tv3.freeviewnz.tv',
    'TV3 PLUS 1': 'tv3-plus1.freeviewnz.tv',
    'PRIME': 'prime.freeviewnz.tv',
    'TV3': 'tv3.freeviewnz.tv',
    'C4': 'c4.freeviewnz.tv',
    'Maori Television': 'maori-tv.freeviewnz.tv',
    'TVNZ 6': 'tvnz6.freeviewnz.tv',
    'TVNZ 7': 'tvnz7.freeviewnz.tv',
    'TVNZ Sport Extra': 'sportextra.freeviewnz.tv',
    'Stratos': 'stratos.freeviewnz.tv',
    'Triangle Stratos': 'stratos.freeviewnz.tv',
    'Parliament TV': 'parliament.freeviewnz.tv',
    'CUE': 'cue.freeviewnz.tv',
    'Te Reo': 'te-reo.freeviewnz.tv',
    'CTV8': 'ctv8.freeviewnz.tv',
    'Chinese TV 8': 'ctv8.freeviewnz.tv',
    'tvCentral': 'tvcentral.freeviewnz.tv',
    'RNZ National': 'rnz-national.freeviewnz.tv',
    'Radio New Zealand National': 'rnz-national.freeviewnz.tv',
    'RNZ Concert': 'rnz-concert.freeviewnz.tv',
    'Radio New Zealand Concert': 'rnz-concert.freeviewnz.tv',
    'George FM': 'george-fm.freeviewnz.tv',
    'Base FM': 'base-fm.freeviewnz.tv',
    'Shine TV': 'shinetv.freeviewnz.tv',
}

class Opener(urllib.FancyURLopener):
    version = '%s/%s' % (NAME, VERSION)

urllib._urlopener = Opener()

def indent(elem, level=0):
    i = "\n" + level * "\t"
    if len(elem):
        if not elem.text or not elem.text.strip():
            elem.text = i + "\t"
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
        for elem in elem:
            indent(elem, level+1)
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
    else:
        if level and (not elem.tail or not elem.tail.strip()):
            elem.tail = i

chan_data = urllib.urlopen(URL).read()
chan_tree = ElementTree.XML(chan_data)
prog_trees = []

tv = ElementTree.Element('tv')
tv.set('generator-info-name', NAME)
tv.set('date', NZTIME.localize(datetime.now()).strftime(OUT_TIME_FORMAT))

for old_channel in chan_tree:
    url = old_channel.find('url').text
    name = old_channel.find('name').text
    try:
        xmltvid = XMLTVIDS[name]
    except:
        sys.stderr.write('No XMLTVID for channel: %s\n' % name)
        continue
    channel = ElementTree.SubElement(tv, 'channel')
    display_name = ElementTree.SubElement(channel, 'display-name')
    display_name.text = name
    icon = ElementTree.SubElement(channel, 'icon')
    icon.set('src', old_channel.find('image').text)
    channel.set('id', xmltvid)
    prog_data = urllib.urlopen(url).read()
    prog_trees.append((xmltvid, ElementTree.XML(prog_data)))

for xmltvid, prog_tree in prog_trees:
    for old_program in prog_tree.findall('.//programme'):
        program = ElementTree.SubElement(tv, 'programme')
        start = old_program.get('datetime_start')
        if start:
            start = datetime.fromtimestamp(
                mktime(
                    strptime(start, IN_TIME_FORMAT)
                )
            )
            start = NZTIME.localize(start)
            program.set('start', start.strftime(OUT_TIME_FORMAT))
        stop = old_program.get('datetime_end')
        if stop:
            stop = datetime.fromtimestamp(
                mktime(
                    strptime(stop, IN_TIME_FORMAT)
                )
            )
            stop = NZTIME.localize(stop)
            program.set('stop', stop.strftime(OUT_TIME_FORMAT))
        program.set('channel', xmltvid)
        title = ElementTree.SubElement(program, 'title')
        title.text = old_program.find('title').text
        synopsis = old_program.find('synopsis')
        if synopsis is not None:
            desc = ElementTree.SubElement(program, 'desc')
            desc.text = synopsis.text
        hd = old_program.get('hd')
        if hd and hd == 'Y':
            video = ElementTree.SubElement(program, 'video')
            present = ElementTree.SubElement(video, 'present')
            present.text = 'yes'
            aspect = ElementTree.SubElement(video, 'aspect')
            aspect.text = '16:9'
            quality = ElementTree.SubElement(video, 'quality')
            quality.text = 'HDTV'
        captioned = old_program.get('captioned')
        if captioned and captioned == 'Y':
            subtitles = ElementTree.SubElement(program, 'subtitles')
            subtitles.set('type', 'teletext')
        classification = old_program.get('classification')
        if classification:
            rating = ElementTree.SubElement(program, 'rating')
            rating.set('system', 'Freeview')
            value = ElementTree.SubElement(rating, 'value')
            value.text = classification

indent(tv)
ElementTree.dump(tv)