pack

#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
import argparse,sys,os,selectors,subprocess,json,struct,tempfile

def creadline(i):
    line = "\n"
    while (len(line)>0 and len(line.strip())==0) or line.startswith('#'):
        line = i.readline()
        try: line = line.decode('utf8')
        except: pass
    return line

def isfloat(s):
    try: return not float(s) is None
    except: return None

def issubtitle(f):
    try:
        cmd = "ffprobe -hide_banner -loglevel quiet -of json -show_streams"
        cmd = cmd.split() + [f]
        o = subprocess.check_output(cmd)
        o = json.loads(o.decode('utf8'))['streams']
        return any(x['codec_type'] == 'subtitle' for x in o)
    except subprocess.CalledProcessError:
        return False

def peekline(f):
    # get first non-comment line and split into fields
    probe_buf = f.peek(2**12)
    line = [l for l in probe_buf.decode('utf8').split('\n') if\
            len(l.strip())!=0 and not l.startswith('#')][0]
    if len(line) == 0:
        raise Exception("unable to probe input file, does it conain data?")
    return line.strip().split()

class Data():
    def __init__(self, f, offset, nfields, rate, name):
        self.s   = struct.Struct('<%df'%nfields)
        self.out = tempfile.NamedTemporaryFile(delete=False)
        self.inf = f
        self.off = offset
        self.rate = rate
        self.name = name
        if nfields > 2:
            self.cmd  = "-f f32le -ac %d.0 -channel_layout %d.0 -ar %f -i"%(nfields,nfields,rate)
        elif nfields == 0:
            raise Exception("ERR: not enough input fields")
        else:
            self.cmd  = "-f f32le -ac %d.0 -ar %f -i"%(nfields,rate)
        self.streams = 1

    def __call__(self, fields=None):
        fields = fields or creadline(self.inf).split()
        fields = [float(f) for f in fields[self.off:]]
        self.out.write( self.s.pack(*fields) )

    def ffmpeg(self):
        self.out.close()
        return self.cmd.split() + [self.out.name]

    def ffmpeg_map(self,i):
        return ("-map %d:0"%i).split()

    def finalize(self):
        os.remove(self.out.name)

class Label():
    """ convert a label stream to an srt subtitle stream, extracting timecodes
        with a fixed sampling rate.
    """
    def __init__(self, f, rate, name):
        self.out = tempfile.NamedTemporaryFile(delete=False)
        self.inf = open(f,'r')
        self.rate = float(rate)
        self.name = name
        self.cmd  = "-f srt -i"
        self.n,self.b,self.l = 0,0,None
        self.count = 0
        self.streams = 1

    def totc(self, n):
        tc = n/self.rate
        return "%d:%02d:%02d,%03d"%\
               (tc/3600, (tc%3600)/60, tc%60, (tc*1000)%1000)

    def __call__(self, fields=None):
        fields = fields or creadline(self.inf).split()

        if self.l != fields[0]:
            srt = "%d\n%s --> %s\n%s\n\n"%\
                  (self.count, self.totc(self.b), self.totc(self.n), self.l)
            if self.l is not None:
                self.out.write(srt.encode('utf8'))
            self.b = self.n
            self.l = fields[0]
            self.count += 1

        self.n += 1

    def ffmpeg(self):
        srt = "%d\n%s --> %s\n%s\n\n"%\
              (self.count, self.totc(self.b), self.totc(self.n), self.l)
        if self.l is not None:
            self.out.write(srt.encode('utf8'))

        self.out.close()
        return self.cmd.split() + [self.out.name]

    def ffmpeg_map(self,i):
        return ("-map %d:0"%i).split()

    def finalize(self):
        os.remove(self.out.name)

class Media():
    def __init__(self, f):
        self.ifile = f
        self.cmd   = "-i"

        try:
            cmd = "ffprobe -hide_banner -loglevel quiet -of json -show_streams"
            cmd = cmd.split() + [f]
            o = subprocess.check_output(cmd)
            self.streams = json.loads(o.decode('utf8'))['streams']
        except CalledProcessError:
            raise Exception("%s is not an ffmpeg media file"%f)

    def __call__(self):
        pass

    def ffmpeg(self):
        cmd = self.cmd.split()
        cmd += [self.ifile]
        return cmd

    def ffmpeg_map(self,i):
        cmd = []
        for n in range(len(self.streams)):
            cmd += ("-map %d:%d"%(i,n)).split()
        return cmd

    def finalize(self):
        pass

class LabelMedia(Media):
    def __init__(self,f,name):
        Media.__init__(self,f)
        self.name = name

    def ffmpeg_map(self,i):
        n = [x['codec_type'] for x in self.streams].index('subtitle')
        return ("-map %d:%d"%(i,n)).split()

def pack(ins,labels,rate=[],name=[],media=[],verbose=False,veryverbose=False,output="-",seek=[],**rest):
    """ ffmpeg is used to pack everything into one .mkv file, to do this we need a
        pipe for each input file (that gets converted into binary) and a pipe for
        each label stream. There is special input file format (called grt), which
        the first field is non-float string that serve as a label stream. Decide
        for each file what is neccesary and generate the part of the neccesary
        ffmpeg command line.

        However, ffmpeg's ability to support sparse streaming input is limited,
        and when using it directly ends up in a dead-lock. Therefore we create
        temprorary files for each input, write them to the disk and the end
        aggregate them into an .mkv file after SIGTERM or EOF.
    """
    # for each input file we build a transformation
    filters = []

    for infile in ins:
        # these file can come in these formats:
        #  1. data-only, white space separated floats
        #  2. label and data, white space separated non-float and floats
        fields = peekline(infile)
        filters.append( Data(infile, not isfloat(fields[0]), len(fields)-1, rate.pop(0), name.pop(0)) )

    for infile in labels:
        # labels can either be supplied:
        #  1. subtitle files
        #  2. label per line
        if issubtitle(infile):
            filters.append( LabelMedia(infile, name.pop(0)) )
        else:
            filters.append( Label(infile, rate.pop(0), name.pop(0)) )

    for infile in media:
        filters.append( Media(infile) )

    # read all input filter so far
    if verbose or veryverbose:
        sys.stdout.write("reading input files...")
        sys.stdout.flush()

    while any(not(type(f) is Media or type(f) is LabelMedia) for f in filters):
        try: [f() for f in filters]
        except:  break # when EOF is hit in one input

    # now build the ffmpeg command line
    cmd,i = ['ffmpeg'],0
    cmd += ["-nostdin"]
    cmd += [] if not veryverbose else '-loglevel debug'.split()
    cmd += [] if verbose or veryverbose else '-hide_banner -loglevel quiet'.split()
    for f in filters: cmd += f.ffmpeg()
    cmd += "-c:a wavpack -c:v copy -c:s srt".split()
    for f in filters:
        cmd += f.ffmpeg_map(i)
        if not type(f) is Media:
            cmd += ['-metadata:s:%d'%i, 'name=%s'%f.name]
        i += 1
    cmd += ["-f", "matroska"]
    cmd += ["-y", output]

    if verbose or veryverbose:
        sys.stdout.write("done\n")
        sys.stdout.write(" ".join(cmd)+"\n")
        sys.stdout.flush()

    subprocess.check_call(cmd)

    for f in filters:
        f.finalize()

if __name__ == "__main__":
    cmdline = argparse.ArgumentParser("packs csv data and additional media into .mkv file using ffmpeg")
    cmdline.add_argument('-i', '--input', type=str,   default=[], action="append", help="input file, multiple can be given")
    cmdline.add_argument('-l', '--label', type=str,   default=[], action="append", help="label file, multiple can be given")
    cmdline.add_argument('-n', '--name',  type=str,   default=[], action="append", help="a name for each input/label/media file can be given, in the same order as arguments are supplied")
    cmdline.add_argument('-r', '--rate',  type=float, default=[], action="append", help="rate (in Hz) of each input file, mandatory")
    cmdline.add_argument('-m', '--media', type=str,   default=[], action="append", help="additional file to add, multiple can be given, may contain subtitles")
    cmdline.add_argument('-ss', '--seek', type=str,   default=[], action="append", help="provide an offset for media files to align them with subtitle timebase in (hh:mm:ss.ms) format")
    cmdline.add_argument('-v', '--verbose', action="store_true", help="more verbose output")
    cmdline.add_argument('-vv', '--veryverbose', action="store_true", help="even more verbose output")
    cmdline.add_argument('output', nargs='?', type=str, default='-', help="file to pack into")
    args = cmdline.parse_args()

    if len(args.input)==0 and len(args.label)==0:
        args.input = ['-']

    i = [open(f,'rb',buffering=2**21) if f != '-' else sys.stdin.buffer for f in args.input]

    try:
        #
        # make the last rate the default rate
        #
        args.rate += [args.rate[0]] * (len(args.input)+len(args.label)-len(args.rate))


        # need to order rate and names to the specific label or input option
        r1,r2,n1,n2 = [],[],[],[]
        for arg in sys.argv:
            if arg == '-i' or arg == '--input':
                r1.append(args.rate.pop(0))
                n1.append(args.name.pop(0) if len(args.name)>0 else "")
            if arg == '-l' or arg == '--label':
                r2.append(args.rate.pop(0))
                n2.append(args.name.pop(0) if len(args.name)>0 else "")

        if len(r1)>0 or len(r2)>0:
            args.rate = r1+r2
            args.name = n1+n2

        pack(i,args.label,**vars(args))
    except IndexError as e:
        sys.stderr.write("ERR: not enough rates/names supplied for inputs\n")
        sys.exit(-1)
    except Exception as e:
        sys.stderr.write("ERR: %s\n"%str(e))
        sys.exit(-1)