# Copyright (c) 2010, Nathaniel Theis

import pexpect, re, fcntl, os, sys
 

#Define some handy things (regexes, strings, etc)

trackmatcher = re.compile(r"Track\s(?P<tracknum>\d+):\s(?P<tracktype>[A-Za-z]+)"
            + r"\W+(?P<tracksize>\d+\s[MmKkGg]B)\s*(?P<tracktime>\(\d+"
            + r":\d+\.\d+\))?")

totalmatcher = re.compile(r"Total size:\s*(?P<totsize>\d+\s[MmKkGg]B)\s+"
            + r"(?P<tottime>\(\d+:\d+\.\d+\))\s=\s(?P<totsectors>\d+)\ssectors")

writestartmatcher = re.compile(r"0\sseconds.\sOperation\sstarts.")

trackstartmatcher = re.compile(r"Starting\snew\strack\sat\ssector:\s+"
                               +r"(?P<startsector>\d+)")

trackprogressmatcher = re.compile(r"\rTrack\s(?P<tracknum>\d+):\s+("
                                  + r"?P<sofar>\d+)\sof\s+(?P<total>\d+)\s"
                                  + r"(?P<unit>[MmKkGg]B)\swritten"
                                  + r"(\s\(fifo\s+(?P<fifofill>\d+)+%\)\s+)?"
                                  + r"(\[buf\s+(?P<buffill>\d+)%\]\s+)?"
                                  + r"((?P<speed>\d+\.\d+)x)?")

trackdonematcher = re.compile(r"Track\s(?P<tracknum>\d+):\sTotal\sbytes"
                              + r"\sread/written:\s(?P<bytesread>\d+)/"
                              + r"(?P<byteswritten>\d+)\s\((?P<sectors>\d+)\s"
                              + r"sectors\)\.")

fixatestartmatcher = re.compile(r"Fixating...")

fixatedonematcher = re.compile(r"Fixating\stime:\s+(?P<time>\d+\.\d+)s")
                                  
carriagereturn = chr(13)

class USBStickWriter():
    def __init__(self, image, bs=512):
        self.image = image
        self.bs = bs
        
    def run(self, device):
        image_file = open(self.image, "r")
        image_fd = image_file.fileno()
        
        fcntl.fcntl(image_fd, fcntl.F_SETFL, os.O_NONBLOCK)
        #We need to use non-blocking read, so tell Linux that
        
        with open(device, "w") as out
            while True:
                block = image_file.read(self.bs)

                if block == "": break
                
                out.write(block)
            
class WodimWrapper():
    '''WodimWrapper is a class to burn CDs with Wodim'''
    def _get_track_data(self):
        trackinfo = {}
        totalline = ""
        
        while True:
            whichmatch = self.wodim_process.expect([trackmatcher, totalmatcher])
            matchdata = self.wodim_process.match.groupdict()
            if whichmatch == 1:
                #Total info line
                totalline = matchdata
                break
            else:
                #Track info line
                trackinfo[matchdata["tracknum"]] = matchdata

        return (trackinfo, totalline)

    def _wait_until_write_start(self):
        match = self.wodim_process.expect(writestartmatcher)

    def _do_burn(self):
        self.wodim_process.expect(trackstartmatcher)
        match = self.wodim_process.match.groupdict()
        
        startsector = int(match['startsector'])
        
        self.trackstart_cb(startsector)

        while True:
            linetype = self.wodim_process.expect([trackprogressmatcher,
                                       trackdonematcher])

            match = self.wodim_process.match.groupdict()

            if linetype == 0:
                self.trackprogress_cb(match)
            elif linetype == 1:
                self.trackdone_cb(match)
                break

        ##FIXME: Only supports 1-track CD burning
        ##I need example output for a multi-track burn

        ##fixate CD

        self.wodim_process.expect(fixatestartmatcher)
        self.fixatestart_cb()
        
        self.wodim_process.expect(fixatedonematcher)
        self.fixatedone_cb(self.wodim_process.match.groupdict())

        ##now just wait for Wodim to finish!

        if self.wodim_process.exitstatus is None: 
            self.wodim_process.wait()

        self.wodim_process.close()

    def __init__(self, args, trackstart_cb, trackprogress_cb,
                 trackdone_cb, fixatestart_cb, fixatedone_cb, 
                 burncomplete_cb, trackdataready_cb,
                 command_path="wodim"):
        '''Initialize stuff, mainly set variables'''
        #Yeah, that's a lot of callbacks
        
        
        self.trackstart_cb = trackstart_cb
        self.trackprogress_cb = trackprogress_cb
        self.trackdone_cb = trackdone_cb
        self.fixatestart_cb = fixatestart_cb
        self.fixatedone_cb = fixatedone_cb
        self.burncomplete_cb = burncomplete_cb
        self.trackdataready_cb = trackdataready_cb
        
        self.wodim_command = command_path + " -v "
        #Some info is only shown with -v
        
        self.wodim_command += " ".join(args)

    def run(self):
        '''Do the burn!'''
        self.wodim_process = pexpect.spawn(self.wodim_command)

        trackdata = self._get_track_data()
        self.trackdataready_cb(trackdata)
        
        self._wait_until_write_start()

        self._do_burn()

        self.burncomplete_cb()

        return self.wodim_process.exitstatus
        
        
        
if __name__ == "__main__":
    #run a few tests for development purposes
    #i should write unit tests later
    
    def printer(*args):
        for arg in args: print arg

    f = USBStickWriter("ToasterActivity.py", bs=512)
    f.run("/dev/mem")
    #ww = WodimWrapper([], printer, printer, printer,
    #                  printer, printer, printer, printer,
    #                 command_path="./wodim")
    #ww.run()