unpack

#!/usr/bin/env python3
# -*- encoding: utf-8 -*-

import argparse,sys,subprocess,shlex,tempfile,os,shutil,re,struct,time,json,selectors,re
from collections import deque
from threading import Thread as T
from datetime import timedelta as td
from select import select
from string import whitespace as ws
from struct import Struct, error
from queue import Queue

class ToFFmpeg():
    def __init__(self, fr0m, to):
        self.fileno = fr0m.fileno()
        self.to = os.fdopen(to,'wb')
        self.fr0m = fr0m
        self.BUFSIZE = 4096
        self.buf = []

    def __call__(self):
        if len(self.buf) == 0:
            self.buf = self.fr0m.read(self.BUFSIZE)

        n = self.to.write(self.buf)
        self.buf = self.buf[n:]

        if n == 0:
            self.fr0m.close()
            self.to.close()
            return True

class FromFFmpeg():
    def __init__(self, f, channel, inputbuf):
        self.f = os.fdopen(f,'rb')
        self.struct = Struct('%df'%channel)
        self.fileno = f
        self.channel = channel
        self.inputbuf = inputbuf
        self.x = 0

    def __call__(self):
        buf = self.inputbuf[self.fileno]
        self.x += len(buf)
        line = self.f.read(4*self.channel)
        if len(line) == 0:
            return True
        buf.appendleft(self.struct.unpack(line))

class SubtitleReader():
    def __init__(self, subfile,rate,style):
        self.sub = os.fdopen(subfile,'r')
        self.fileno = subfile
        self.labels = []
        self.lastt = None
        self.rate = rate
        self.style = style

    def ass2rate(self,ss):
        h,m,s = [abs(float(x)) for x in ss.split(':')]
        return int( (60*60*h+60*m+s) * self.rate )

    def __call__(self):
        if self.sub is None:
            return

        line = self.sub.readline()
        if line == "":
            self.sub.close()
            return True

        try:
            _,beg,end,style,_,_,_,_,_,label = line.split(',')
            label = re.sub('\s+','_',label.strip())
            beg,end = self.ass2rate(beg),self.ass2rate(end)
            if self.style is None or self.style == style:
                self.labels.append((beg,end,label))
        except ValueError as e:
            pass

    def label(self,t):
        if self.lastt is not None:
            assert (t-self.lastt) == 1, "%f"%(t-self.lastt)

        self.lastt = t
        try:
            b,e,l = self.labels[0]
            if t < b:
                return None
            while t >= e:
                b,e,l = self.labels.pop(0)
                b,e,l = self.labels[0]
            return l
        except IndexError as e:
            return None

def duration2ticks(s, rate):
    """ convert from ffmpeg's duration format to ticks in rate
    """
    m = re.match("(?P<neg>-)?((?P<hours>\d{1,2}):)?(?P<minutes>\d{1,2}):(?P<seconds>\d{1,2})(?P<millis>.\d{1,3})?", s)
    if m is not None:
        g = m.groupdict()
        s =  float(g['hours']   or 0)*60*60
        s += float(g['minutes'] or 0)*60
        s += float(g['seconds'] or 0)
        s += float(g['millis']  or 0)
        return s * rate

    m = re.match("(?P<neg>-)?(?P<seconds>\d+)(?P<millis>.\d{1,3})?", s)
    if m is not None:
        g = m.groupdict()
        s  = float(g['seconds'] or 0)
        s += float(g['millis']  or 0)
        return s * rate

    return None


def unpack_to_str(streams,sub,default,style,verbose=False,veryverbose=False,ss=0,sd=None,rate=None,sub_only=False):
    # use -af aresample=resampler=soxr -ar 100 to switch to a common output rate
    # and merge all stream. The amerge complex filter can't be used since it can
    # not be taught to use the soxr resampler (the default one does not work)
    # This is only done when resampling is neccesary.

    files = list(set(f for (d,f) in streams))
    subasfile = sub and not sub[0]['file_index'] in (d['file_index'] for (d,_) in streams)

    # a special case, when the subtitle is in its own file
    if subasfile: files.append(sub[1])

    # setup all input, output and subtitle pipes
    ins   = [py3k(os.pipe()) for x in files]
    outs  = [py3k(os.pipe()) for x in streams]
    subp  =  py3k(os.pipe()) if sub else None

    # figure out the largest sample_rate for upsampling
    rate  = rate or max([int(d['sample_rate']) for (d,f) in streams])
    for r in files:
        os.set_inheritable(r.fileno(),True)

    # stream seek echo output separately
    ssd = ' -ss ' + ss
    ssd += '' if sd is None else (' -t ' + sd)

    # build the command line
    cmd  = 'ffmpeg'
    cmd += '' if not veryverbose else ' -loglevel debug'
    cmd += '' if verbose or veryverbose else ' -hide_banner -loglevel quiet'
    cmd += ''.join(' -i pipe:%d'%r for r,_ in ins)
    cmd += '' if not sub else ' -map %d:%d -f ass %s pipe:%d'%(sub[0]['file_index'],sub[0]['index'],ssd,subp[1])

    for (d,f),(_,w) in zip(streams,outs):
        cmd += ' -map %d:%d'%(d['file_index'],d['index'])
        cmd += ' -f f32le -ar %f'%rate
        cmd += '' #if len(outs)==1 else ' -af aresample=resampler=soxr:linear_interp=1'
        cmd += ssd
        cmd += ' pipe:%d'%w

    if verbose or veryverbose:
        sys.stderr.write(cmd+"\n")

    # cmdline built, now copy all files to their respective w parts of in,
    # and read the results for r parts of outs. Do the same for the sub
    pid = os.fork()
    if not pid:
        for (_,w) in ins: os.close(w)
        for (r,_) in outs+([subp] if sub else []): os.close(r)

        # make sure that we can background the process without blocking
        if sys.stdin.fileno() not in outs:
            os.close(sys.stdin.fileno())

        cmd = cmd.split()
        os.execvp(cmd[0],cmd)
    else:
        for (r,_) in ins: os.close(r)
        for (_,w) in outs+([subp] if sub else []): os.close(w)

        # we need to know the number of channels for each stream to calculate
        # the number of bytes to read for each stream
        chns = [int(d['channels']) for (d,f) in streams]
        fmt  = ("{} "+"{:f} "*sum(chns)).format

        sel = selectors.DefaultSelector()
        ibuf = { k : deque() for (k,_) in outs }

        if sub is not None:
            subr = SubtitleReader(subp[0], rate, style)
            sel.register(subr.fileno, selectors.EVENT_READ, subr)

        for (f,(_,w)) in zip(files,ins):
            sel.register(w, selectors.EVENT_WRITE, ToFFmpeg(f,w))
        for ((r,_),c) in zip(outs,chns):
            sel.register(r, selectors.EVENT_READ, FromFFmpeg(r,c,ibuf))

        #
        # Actually the stream-seek (ss) can be problematic, because ffmpeg
        # removes all subtitle where the beginning is smaller than streamseek.
        # So you can jump in the middle of a subtitle, which results in that
        # label being removed. XXX
        #
        t,block = duration2ticks(ss,rate),0 # keep a tick timer
        while len(sel.get_map().values()) > 0:
            events,block = sel.select(), block+1
            for key, mask in events:
                if key.data(): sel.unregister(key.fileobj)

            if block % 1000000 != 0 and len(sel.get_map().values()) != 0:
                continue

            n = min(len(v) for v in ibuf.values())
            for i in range(n):
                bufs = [ y for x in ibuf.values() for y in x.pop()]
                if sub_only:
                    print(subr.label(t) or default if sub is not None else default)
                else:
                    print(fmt(subr.label(t) or default if sub is not None else default, *bufs))
                t += 1

def py3k(tup):
    """ brain-deadness of py3k """
    for x in tup: os.set_inheritable(x,True)
    return tup

def group(files):
    sub,result,filecount = None,[],0

    for streams in files:
        for desc,fd in streams:
            desc['file_index'] = filecount

            if desc['codec_type'] == 'subtitle':
                sub = (desc,fd)
            elif desc['codec_type'] == 'audio':
                result.append( (desc,fd) )
            else:
                raise Exception("got unhandled codec_type %s"%desc['codec_type'])

        if sub:
            yield result,sub
            sub,result,filecount = None,[],0
        else:
            filecount += 1

    # stuff without subtitle
    if len(result):
        yield result,None


def probe(f):
    buf = b""
    while len(buf) < 2e6:
        newdata = f.peek(int(2e6))
        if len(newdata) == 0:
            sys.exit(-1)
        buf += f.peek(int(2e6))

    cmd = "ffprobe -hide_banner -loglevel quiet -of json -show_streams -".split()
    cmd = subprocess.Popen(cmd,stdin=subprocess.PIPE,stdout=subprocess.PIPE)
    o,e = cmd.communicate(buf)
    return json.loads(o.decode("utf-8"))['streams']

def sid(s):
    if not 'tags' in s.keys():
        return None
    elif 'LANGUAGE' in s['tags'].keys():
        return s['tags']['LANGUAGE']
    elif 'NAME' in s['tags'].keys():
        return s['tags']['NAME']
    else:
        return None

if __name__ == '__main__':
    cmdline = argparse.ArgumentParser('unpacks .mkv file (using ffmpeg)')
    cmdline.add_argument('files', type=str, default=["-"], nargs='*',       help="input files")
    cmdline.add_argument('-r', '--rate', type=float, default=None,          help="resample stream to given rate in Hz")
    cmdline.add_argument('-d', '--default', type=str, default="NULL",       help="default class label for unlabelled data")
    cmdline.add_argument('-s', '--style', type=str, default=None,           help="the class label style to select")
    cmdline.add_argument('-t', '--stream', type=str, nargs='*', default=[], help="select streams to consider by their meta-tag *language* or *name*, multiple can be given, can be comma-seperated ")
    cmdline.add_argument('-ss', type=str, default="0",                      help="seeks to a start position, format is documented in ffmpeg-utils(1)")
    cmdline.add_argument('-sd', type=str, default=None,                     help="duration of the playback, format is documented in ffmpeg-utils(1)")
    cmdline.add_argument('-so', '--sub-only', action="store_true"         , help="output only subtitle labels, no numerical data")
    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")
    args = cmdline.parse_args()

    #
    # check for correct version of ffmpeg for stream seek
    #
    version = subprocess.check_output('ffmpeg -version'.split())
    if args.ss != "0" and not b'2016' in version:
        sys.stderr.write("WARN: streamseeking with this version of ffmpeg is buggy")

    # we open each file here already and create an input buffer
    # for probing each input and then writing the read data to
    # ffmpeg in order to not have to read something twice
    iopen = lambda f: open(f,'rb',buffering=2**21)
    sopen = lambda f: os.fdopen(sys.stdin.fileno(),'rb',buffering=2**21)
    args.files = [ item for sublist in [ x.split(',') for x in args.files ] for item in sublist ]
    files = (iopen(f) if f != '-' else sopen(f) for f in args.files )
    files = [ (probe(f),f) for f in files ]
    sids  = [ sid(s) for (ss,f) in files for s in ss if s['codec_type']=='audio' ]

    args.stream = [ item for sublist in [ x.split(',') for x in args.stream ] for item in sublist ]
    #for i in args.stream: sys.stderr.write(i + " ")
    #sys.stderr.write('\n')

    # selector audio and subtitle stream
    okay = lambda d: d['codec_type']=='subtitle' or\
            args.stream==[] and d['codec_type']=='audio' or\
            'tags' in d.keys() and d['codec_type']=='audio' and (\
                   'LANGUAGE' in d['tags'].keys() and d['tags']['LANGUAGE'] in args.stream or\
                       'NAME' in d['tags'].keys() and d['tags']['NAME'] in args.stream)

    # filter out all streams which are not neccesary or selected
    files = [[(s,f) for s in ss if okay(s)] for (ss,f) in files ]

    for audio,sub in group(files):
        if not audio:
            err = "no data stream selected, possible are: %s\n"%str(sids)
            sys.stderr.write(err)
            sys.exit(-1)

        try:
            unpack_to_str(audio,sub,args.default,args.style,args.verbose,args.veryverbose,args.ss,args.sd,args.rate,args.sub_only)
        except struct.error:
            raise
            sys.exit(0)