# -*- Mode: Python -*-
# vi:si:et:sw=4:sts=4:ts=4
#
import pygst
pygst.require('0.10')
import gst
import gst.interfaces
import gobject
gobject.threads_init()
import time
import gtk
import totem
import os
import string
channels = {}
transport_streams = {}

#TODO: Finish DVB-S handling (nit delivery system handling and channels output)

class DvbScanner:
    def __init__(self, adapter=0, frontend=0, scanning_complete_cb=None,
        frequency_scanned_cb=None, channel_added_cb=None):
        self.frequencies_to_scan = []
        self.channels = {} # service id -> transport stream id
        self.current_tuning_params = {}
        self.frequencies_scanned = []
        self.transport_streams = {} # ts id -> tuning_params
        self.pipeline = None
        self.locked = False
        self.nit_arrived = False
        self.sdt_arrived = False
        self.pat_arrived = False
        self.check_for_lock_event_id = None
        self.wait_for_tables_event_id = None
        self.scanning_complete_cb = scanning_complete_cb
        self.frequency_scanned_cb = frequency_scanned_cb
        self.channel_added_cb = channel_added_cb
        self.pipeline =  gst.parse_launch(
            "dvbsrc name=dvbsrc adapter=%d frontend=%d pids=0:16:17:18 " \
            "stats-reporting-interval=0 ! mpegtsparse ! " \
            "fakesink silent=true" % (adapter, frontend))
        bus = self.pipeline.get_bus()
        bus.add_signal_watch()
        bus.connect("message", self.bus_watch_func)

    def wait_for_tables(self):
        self.pipeline.set_state(gst.STATE_READY)
        self.pipeline.get_state()
        self.locked = False
        self.scanned = False
        if self.nit_arrived and self.sdt_arrived:
            self.frequencies_scanned.append(
                self.current_tuning_params["frequency"])
            if self.frequency_scanned_cb:
                self.frequency_scanned_cb(
                    self.current_tuning_params["frequency"])
        for f in self.frequencies_to_scan:
            if f["frequency"] in self.frequencies_scanned:
                self.frequencies_to_scan.remove(f)
            else:
                time.sleep(2)
                self.scan(f)
                self.scanned = True
                break
        if not self.scanned:
            if self.scanning_complete_cb:
                self.scanning_complete_cb()
        return False

    def check_for_lock(self):
        if not self.locked:
            self.pipeline.set_state(gst.STATE_READY)
        self.pipeline.get_state()
        self.scanned = False
        for f in self.frequencies_to_scan:
            if f["frequency"] in self.frequencies_scanned:
                self.frequencies_to_scan.remove(f)
            else:
                time.sleep(2)
                self.scan(f)
                self.scanned = True
                break
        if not self.scanned:
            self.scanning_complete_cb()
        return False

    def have_dvb_adapter_type(self, type):
        #print "The DVB Adapter we have is a %s" % type
        pass

    def bus_watch_func(self, bus, message):
        t = message.type
        #print "Bus watch function for message %r" % message
        if t == gst.MESSAGE_ELEMENT:
            if message.structure.get_name() == 'dvb-adapter':
                s = message.structure
                self.have_dvb_adapter_type(s["type"])
            elif message.structure.get_name() == 'dvb-frontend-stats':
                s = message.structure
                if s["lock"] and not self.locked:
                    self.locked = True
                    gobject.source_remove(self.check_for_lock_event_id)
                    self.check_for_lock_event_id = None
                    self.wait_for_tables_event_id = gobject.timeout_add(
                        10*1000, 
                        self.wait_for_tables)
            elif message.structure.get_name() == 'dvb-read-failure':
                if self.check_for_lock_event_id:
                    gobject.source_remove(self.check_for_lock_event_id)
                    self.check_for_lock_event_id = None
                if self.wait_for_tables_event_id:
                    gobject.source_remove(self.wait_for_tables_event_id)
                    self.wait_for_tables_event_id = None
                self.wait_for_tables()
            elif message.structure.get_name() == 'sdt':
                s = message.structure
                services = s["services"]
                tsid = s["transport-stream-id"]
                actual = s["actual-transport-stream"]
                if actual:
                    for service in services:
                        name = service.get_name()
                        sid = int(name[8:])
                        if service.has_key("name"):
                            name = service["name"]
                        if self.channels.has_key(sid):
                            self.channels[sid]["name"] = name
                            self.channels[sid]["transport-stream-id"] = tsid
                        else:
                            self.channels[sid] = { "name": name, \
                                "transport-stream-id": tsid }
                        if self.channel_added_cb:
                            self.channel_added_cb(sid, self.channels[sid])
                    self.sdt_arrived = True

            elif message.structure.get_name() == 'nit':
                s = message.structure
                name = s["network-id"]
                actual = s["actual-network"]
                if s.has_key("network-name"):
                    name = s["network-name"]
                transports = s["transports"]
                for transport in transports:
                    tsid = transport["transport-stream-id"]
                    if transport.has_key("delivery"):
                        delivery = transport["delivery"]
                        if delivery["frequency"] not in \
                            self.frequencies_scanned:
                            self.frequencies_to_scan.append(delivery)
                            self.transport_streams[tsid] = delivery

                        if transport.has_key("channels"):
                            chans = transport["channels"]
                            for chan in chans:
                                if self.channels.has_key(
                                    chan["service-id"]):
                                    self.channels[chan["service-id"]]["logical-channel-number"] = chan["logical-channel-number"]
                                else:
                                    self.channels[chan["service-id"]] = { "logical-channel-number" : chan["logical-channel-number"] }
                self.nit_arrived = True
            elif message.structure.get_name() == 'pat':
                programs = message.structure["programs"]
                for p in programs:
                    sid = p["program-number"]
                    pmt = p["pid"]
                    if self.channels.has_key(sid):
                        self.channels[sid]["pmt-pid"] = pmt
                    else:
                        self.channels[sid] = { "pmt-pid":pmt }
                self.pat_arrived = True

        if self.sdt_arrived and self.nit_arrived and self.pat_arrived:
            gobject.source_remove(self.wait_for_tables_event_id)
            self.wait_for_tables()

    def scan(self,tuning_params):
        self.frequencies_to_scan.remove(tuning_params)
        self.current_tuning_params = tuning_params
        self.sdt_arrived = False
        self.nit_arrived = False
        self.pat_arrived = False
    
        modulation=""
        if tuning_params["constellation"].startswith("QAM"):
            modulation="QAM %s" % tuning_params["constellation"][3:]
        elif tuning_params["constellation"] == "reserved":
            modulation = "QAM 64"
        else:
            modulation = tuning_params["constellation"]
        for param in ["code-rate-hp", "code-rate-lp"]:
            if tuning_params[param] == "reserved":
                tuning_params[param] = "NONE"
        if tuning_params["hierarchy"] ==  0:
            tuning_params["hierarchy"] = "NONE"
        dvbsrc = self.pipeline.get_by_name("dvbsrc")
        dvbsrc.set_property("frequency", tuning_params["frequency"])
        dvbsrc.set_property("bandwidth", str(tuning_params["bandwidth"]))
        dvbsrc.set_property("code-rate-hp", str(tuning_params["code-rate-hp"]))
        dvbsrc.set_property("code-rate-lp", str(tuning_params["code-rate-lp"]))
        dvbsrc.set_property("trans-mode", 
            str(tuning_params["transmission-mode"]))
        dvbsrc.set_property("guard", str(tuning_params["guard-interval"]))
        dvbsrc.set_property("hierarchy", str(tuning_params["hierarchy"]))
        dvbsrc.set_property("modulation", modulation)
        self.pipeline.set_state(gst.STATE_PLAYING)
        (statereturn, state, pending) = self.pipeline.get_state()
        if statereturn == gst.STATE_CHANGE_FAILURE:
            locked = False
            self.check_for_lock()
        else:
            self.locked = False
            # wait 10 seconds for lock 
            self.check_for_lock_event_id = gobject.timeout_add(10*1000, 
                self.check_for_lock)

    def add_brute_force_scan(self):
        for chan in range(5, 13):
            freq = 142500000 + chan * 7000000
            for transmode in ["2k", "8k"]:
                for guard in ["AUTO", "32", "16", "8", "4"]:
                    self.frequencies_to_scan.append(
                        { "frequency":freq, "bandwidth":7,
                            "code-rate-hp":"NONE", "code-rate-lp":"AUTO", 
                            "constellation":"QAM64",
                            "transmission-mode":transmode,
                            "guard-interval":guard,
                            "hierarchy":"AUTO"})

        for chan in range(21, 70):
            freq = 306000000 + chan* 8000000
            for transmode in ["2k", "8k"]:
                for guard in ["32", "16", "8", "4"]:
                    self.frequencies_to_scan.append(
                        { "frequency":freq, "bandwidth":8,
                            "code-rate-hp":"NONE", "code-rate-lp":"AUTO", 
                            "constellation":"QAM64",
                            "transmission-mode":transmode,
                            "guard-interval":guard,
                            "hierarchy":"AUTO"})
    
    def add_tuning_data(self, tuning_data):
        self.frequencies_to_scan.append(tuning_data)

    def start_scanning(self):
        f = self.frequencies_to_scan[0]
        gobject.idle_add(self.scan, f)

    def output_results(self):
        filename = os.getenv("GST_DVB_CHANNELS_CONF")
        if not filename:
            gstversion = gst.version()
            filename = os.path.expanduser(
                "~/.gstreamer-%d.%d/dvb-channels.conf" % (gstversion[0],
                    gstversion[1]))
        print "Filename: %s" % filename
        file = open(filename, "w")
        lines = []
        for sid in self.channels.keys():
            chan = self.channels[sid]
            tuningparams = None
            line = ""
            #print "%s:::::::::" % chan["name"]
            if chan.has_key("name"):
                line = "%s:" % chan["name"]
            else:
                line = "[%d]:" % sid
            if chan.has_key("transport-stream-id"):
                tuningparams = \
                    self.transport_streams[chan["transport-stream-id"]]
            else:
                continue
            line = "%s%s:INVERSION_AUTO:" % (line, tuningparams["frequency"])
            line = "%sBANDWIDTH_%s_MHZ:" % (line, tuningparams["bandwidth"])
            fecstr = ""
            if tuningparams["code-rate-hp"] == "1/2":
                fecstr = "FEC_1_2"
            elif tuningparams["code-rate-hp"] == "2/3":
                fecstr = "FEC_2_3"
            elif tuningparams["code-rate-hp"] == "3/4":
                fecstr = "FEC_3_4"
            elif tuningparams["code-rate-hp"] == "5/6":
                fecstr = "FEC_5_6"
            elif tuningparams["code-rate-hp"] == "7/8":
                fecstr = "FEC_7_8"
            else:
                fecstr = "FEC_AUTO"
            line = "%s%s:" % (line, fecstr)
            if tuningparams["code-rate-lp"] == "1/2":
                fecstr = "FEC_1_2"
            elif tuningparams["code-rate-lp"] == "2/3":
                fecstr = "FEC_2_3"
            elif tuningparams["code-rate-lp"] == "3/4":
                fecstr = "FEC_3_4"
            elif tuningparams["code-rate-lp"] == "5/6":
                fecstr = "FEC_5_6"
            elif tuningparams["code-rate-lp"] == "7/8":
                fecstr = "FEC_7_8"
            else:
                fecstr = "FEC_AUTO"
            line = "%s%s:" % (line, fecstr)
            constellation = ""
            if tuningparams["constellation"] == "QPSK":
                constellation = "QPSK"
            elif tuningparams["constellation"] == "QAM16":
                constellantion = "QAM_16"
            else:
                constellation = "QAM_64"
            line = "%s%s:" % (line, constellation)
            if tuningparams["transmission-mode"] == "2k":
                line = "%sTRANSMISSION_MODE_2K:" % line
            else:
                line = "%sTRANSMISSION_MODE_8K:" % line
            line = "%sGUARD_INTEVAL_1_%d:" % (line, 
                tuningparams["guard-interval"])
            if tuningparams["hierarchy"] == 0:
                line = "%sHIERARCHY_NONE:" % line
            else:
                line = "%sHIERARCHY_%d:" % (line, tuningparams["hierarchy"])
            line = "%s0:0:%d\n" % (line, sid) 
            lines.append(line)
            print line
        file.writelines(lines)
        file.close()

class ScannerUI(gtk.Window):
    def __init__(self):
        super(ScannerUI, self).__init__()
        self.adapter = 1
        self.channels = {}
        self.adapters = []
        mainvbox = gtk.VBox(3)
        
        hbox = gtk.HBox(3)
        self.adapterlist = gtk.combo_box_new_text()
        self.adapterlist.connect("changed", self.adapter_changed)
        self.initialtuning = gtk.combo_box_new_text()
        self.find_adapters_and_fill_list()
        hbox.add(self.adapterlist)
        hbox.add(self.initialtuning)
        scanbutton = gtk.Button("Scan")
        scanbutton.connect("clicked", self.scan_button_clicked)
        hbox.add(scanbutton)
        mainvbox.add(hbox)

        scanvbox = gtk.VBox(2)
        self.label = gtk.Label("")
        self.progressbar = gtk.ProgressBar()
        scanvbox.add(self.label)
        scanvbox.add(self.progressbar)
        mainvbox.add(scanvbox)
        self.add(mainvbox)

    def adapter_changed(self, combobox):
         adapter, = self.adapterlist.get_model().get(
            self.adapterlist.get_active_iter(), 0)
         self.adapter = int(adapter[len(adapter)-1])
         self.fill_list_of_available_initial_tuning_info(self.adapters[self.adapter])
    def find_adapters_and_fill_list(self):
        for i in range (0,8):
            if os.path.exists('/dev/dvb/adapter%d/frontend0' % i):
                adaptertype = self.get_type_of_adapter(i)
                self.adapterlist.append_text("DVB (%s) Adapter %d" % (
                    adaptertype,i))
                self.adapters.append(adaptertype)
        if len(self.adapters) > 0:
            self.adapterlist.set_active(0)

    def fill_list_of_available_initial_tuning_info(self, adaptertype):
        # clear the current combobox
        liststore = gtk.ListStore(gobject.TYPE_STRING)
        self.initialtuning.set_model(liststore)
        if adaptertype == "DVB-T":
            self.initialtuning.append_text(
                "Unknown (will do a brute-force scan)")
        tuningfiles = []
        for path in ["/usr/share/dvb", "/usr/share/dvb-apps"]:
            if os.path.exists(path):
                for f in os.listdir(os.path.join(path, 
                    string.lower(adaptertype))):
                    tuningfiles.append(f)

        tuningfiles.sort()
        for f in tuningfiles:
            self.initialtuning.append_text(f)
        
        if len(liststore) > 0:
            self.initialtuning.set_active(0)
                
    def get_type_of_adapter(self, adapter):
        dvbelement = gst.element_factory_make ("dvbsrc", "test_dvbsrc")
        dvbelement.set_property("adapter", adapter)
        pipeline = gst.Pipeline("")
        pipeline.add(dvbelement)
        pipeline.set_state(gst.STATE_READY)
        pipeline.get_state()
        bus = pipeline.get_bus()
        adaptertype = None
        while bus.have_pending():
            msg = bus.pop()
            if msg.type == gst.MESSAGE_ELEMENT and msg.src == dvbelement:
                structure = msg.structure
                if structure.get_name() == "dvb-adapter":
                    adaptertype = structure["type"]
                    break
        pipeline.set_state(gst.STATE_NULL)
        return adaptertype

    def scan_button_clicked(self, button):
        adapter, = self.adapterlist.get_model().get(
            self.adapterlist.get_active_iter(), 0)
        self.adapter = int(adapter[len(adapter)-1])

        self.start_scan()

    def channel_added(self, sid, chan):
        # check if we already have this channel
        channum =  sid
        channame = ""
        if not chan.has_key("name"):
            pass
        else:
            channame = chan["name"]
        if chan.has_key("logical-channel-number"):
            channum = chan["logical-channel-number"]
        
        if self.channels.has_key("channame"):
            pass
        else:
            self.channels[channame] = "dvb://%d" % channum
            #self.label.set_text("Scanning: %d channels found" % 
            #    len(self.channels))
   
    def frequency_scanned(self, freq):
        self.label.set_text("Scanning: Frequency %d complete" % freq)
        self.progressbar.pulse()

    def scanning_complete(self):
        self.label.set_text("Scanning complete: %d channels found" % 
            len(self.channels))
        self.scanner.output_results()

    def start_scan(self):
        self.scanner = DvbScanner(adapter=self.adapter,frontend=0, 
            scanning_complete_cb=self.scanning_complete, 
            frequency_scanned_cb = self.frequency_scanned, 
            channel_added_cb = self.channel_added)

        initialtuning, = self.initialtuning.get_model().get(
            self.initialtuning.get_active_iter(), 0)
        if initialtuning == "Unknown (will do a brute-force scan)":
            self.scanner.add_brute_force_scan()
        else:
            # need to parse file to get tuning data
            itd = self.get_initial_tuning_data(self.adapters[self.adapter],
                initialtuning)
            for data in itd:
                self.scanner.add_tuning_data(data)
                print data
        
        self.label.set_text("Scanning:")
        self.scanner.start_scanning()

    def get_initial_tuning_data(self, adaptertype, initialtuningfile):
        ret = []
        for path in ["/usr/share/dvb", "/usr/share/dvb-apps"]:
            filename = os.path.join(path, string.lower(adaptertype), 
                    initialtuningfile)
            if os.path.exists(filename):
                f = open(filename, "r")
                while True:
                    line = f.readline()
                    if not line:
                        break
                    if line[0] != '#':
                        params = line[:-1].split(" ")
                        if params[0] == "T" and adaptertype == "DVB-T":
                            if len(params) == 9:
                                d = { "frequency":int(params[1]),
                                    "bandwidth":int(params[2][0]),
                                    "code-rate-hp":params[3],
                                    "code-rate-lp":params[4],
                                    "constellation":params[5],
                                    "transmission-mode":params[6],
                                    "guard-interval":int(params[7].split("/")[1]),
                                    "hierarchy":params[8]}
                                ret.append(d)
                                
                        elif params[0] == "S" and adaptertype == "DVB-S":
                            if len(params) == 5:
                                d = { "frequency":int(params[1]),
                                    "symbol-rate":int(params[3])/1000,
                                    "code-rate-hp":params[4] }
                                if params[2] == "V":
                                    d["polarization"] = "vertical"
                                else:
                                    d["polarization"] = "horizontal"
                                ret.append(d)
                f.close()
        return ret

class Scanner:
    def __init__(self):
        self.window = ScannerUI()
        self.window.show_all()
       
    def frequency_scanned(self, freq):
        pass

    def channel_added(self, sid, chan):
        self.window.channel_added(sid, chan)

    def scanning_complete(self, blah=None):
        global channels, transport_streams
        print "Scanning complete"
        channels = self.scanner.channels
        transport_streams = self.scanner.transport_streams
        self.scanner.output_results()
        self.window.scanning_complete()

mainloop = gobject.MainLoop()
s = Scanner()
mainloop.run()
