#!/usr/bin/env python3

from __future__ import unicode_literals
from io import open

class AutoMKVChapters:
    class Template:
        def __init__(self):
            from random import randint
            self.uid = randint(10**4,10**6)
            self.num_editions = 1
            self.lang = ['eng']
            self.country = ['us']
            self.fps = '30'
            self.ofps = '24'
            self.qpf = '0'
            self.idr = False
            self.trims = None
            self.kframes = None

        def toxml(self,chapfile):

            chf = open(chapfile+'.xml','w',encoding='utf-8')
            head = '<?xml version="1.0" encoding="UTF-8"?>\n<!-- <!DOCTYPE Tags SYSTEM "matroskatags.dtd"> -->\n'
            chf.write(head+'<Chapters>\n')

            if self.num_editions > 1:
                tagf = open(chapfile+'tags.xml','w')
                tagf.write(head+'<Tags>\n')
            else:
                tagf = False

            for ed in self.editions:
                chf.write('\t<EditionEntry>\n')
                chf.write('\t\t<EditionFlagHidden>{0:d}</EditionFlagHidden>\n'.format(ed.hidden) if ed.hidden else '')
                chf.write('\t\t<EditionFlagDefault>{0:d}</EditionFlagDefault>\n'.format(ed.default) if ed.default else '')
                chf.write('\t\t<EditionFlagOrdered>{0:d}</EditionFlagOrdered>\n'.format(ed.ordered) if ed.ordered else '')
                chf.write('\t\t<EditionUID>{0:d}</EditionUID>\n'.format(ed.uid))

                if tagf:
                    tagf.write('\t<Tag>\n\t\t<Targets>\n')
                    tagf.write('\t\t\t<EditionUID>{0:d}</EditionUID>\n'.format(ed.uid))
                    tagf.write('\t\t\t<TargetTypeValue>50</TargetTypeValue>\n\t\t</Targets>\n')
                    num_names = len(ed.name) if len(ed.name) < len(self.lang) else len(self.lang)
                    for i in range(num_names):
                        tagf.write('\t\t<Simple>\n\t\t\t<Name>TITLE</Name>\n')
                        tagf.write('\t\t\t<String>{0}</String>\n'.format(ed.name[i] if ed.name[i] != '' else ed.name[i-1]))
                        tagf.write('\t\t\t<TagLanguage>{0}</TagLanguage>\n'.format(self.lang[i]))
                        tagf.write('\t\t\t<DefaultLanguage>{0:d}</DefaultLanguage>\n'.format(1 if i == 0 else 0))
                        tagf.write('\t\t</Simple>\n')
                    tagf.write('\t</Tag>\n')

                for ch in ed.chapters:
                    chf.write('\t\t<ChapterAtom>\n')
                    num_names = len(ch.name) if len(ch.name) < len(self.lang) else len(self.lang)
                    for i in range(num_names):
                        chf.write('\t\t\t<ChapterDisplay>\n')
                        chf.write('\t\t\t\t<ChapterString>{0}</ChapterString>\n'.format(ch.name[i] if ch.name[i] != '' else ch.name[i-1]))
                        chf.write('\t\t\t\t<ChapterLanguage>{0}</ChapterLanguage>\n'.format(self.lang[i]) if self.lang[i] != 'eng' else '')
                        chf.write('\t\t\t\t<ChapterCountry>{0}</ChapterCountry>\n'.format(self.country[i]) if i < len(self.country) else '')
                        chf.write('\t\t\t</ChapterDisplay>\n')
                    chf.write('\t\t\t<ChapterUID>{0:d}</ChapterUID>\n'.format(ch.uid))
                    chf.write('\t\t\t<ChapterTimeStart>{0}</ChapterTimeStart>\n'.format(ch.start))
                    chf.write('\t\t\t<ChapterTimeEnd>{0}</ChapterTimeEnd>\n'.format(ch.end) if ch.end else '')
                    chf.write('\t\t\t<ChapterFlagHidden>{0:d}</ChapterFlagHidden>\n'.format(ch.hidden) if ch.hidden != 0 else '')
                    chf.write('\t\t\t<ChapterFlagEnabled>{0:d}</ChapterFlagEnabled>\n'.format(ch.enabled) if ch.enabled != 1 else '')
                    chf.write('\t\t\t<ChapterSegmentUID format="hex">{0}</ChapterSegmentUID>\n'.format(ch.suid) if ch.suid else '')
                    chf.write('\t\t</ChapterAtom>\n')

                chf.write('\t</EditionEntry>\n')

            chf.write('</Chapters>\n')

            if tagf:
                tagf.write('</Tags>\n')
                tagf.close()

            if self.qpf != '0' and self.kframes:
                from vfr import write_qpfile
                if self.qpf != '1':
                    qpfile = self.qpf
                else:
                    qpfile = chapfile+'.qpfile'
                write_qpfile(qpfile, self.kframes, self.idr)

        def connect_with_vfr(self,avs,label=None,clip=None):
            """
            Connects templates.py with vfr.py, enabling its use outside of vfr.py.
            
            Uses the same quirks as AMkvC but only for 24 and 30 fps.
            Ex: inputfps=30 is understood as being '30*1000/1001'
            
            """

            from vfr import parse_trims, fmt_time

            # compensate for amkvc's fps assumption
            if self.fps in ('24','30'):
                fps = self.fps + '/1.001'
            else:
                fps = str(self.fps)
            if self.ofps and self.ofps in ('24','30'):
                ofps = self.ofps + '/1.001'
            else:
                ofps = str(self.ofps)

            Trims2, Trims2ts = parse_trims(avs, fps, ofps, label=label, clip=clip)[2:4]
            Trims2ts = [(fmt_time(i[0]),fmt_time(i[1]) if i[1] != 0 else None) for i in Trims2ts]

            self.trims = Trims2ts
            self.kframes = Trims2

        def parse_mkv(self, path):
            """Parse a Matroska file for SegmentUID and Duration"""
            import binascii
            import struct
            
            def get_data_len(byte):
                """Get the length (bytes) of the element data"""
                n = ord(byte)
                mask = 0b10000000
                while not n & mask:
                    mask >>= 1
                return n & ~mask
            
            suid = tcscale = duration = 0
            with open(path, 'rb') as file:
                if file.read(4) != b'\x1A\x45\xDF\xA3': # not a Matroska file
                    return suid, duration
                chunk_size = 100000 # 100 kB
                i = 0
                while True:
                    if suid and tcscale and duration:
                        break
                    bin = file.read(chunk_size)
                    if not bin:
                        break
                    suid_pos = bin.find(b'\x73\xA4\x90') # \x90 -> 16 bytes
                    if suid_pos != -1:
                        suid_pos = 4 + i * chunk_size + suid_pos + 3
                        file.seek(suid_pos)
                        suid = binascii.hexlify(file.read(16)).decode()
                    tcscale_pos = bin.find(b'\x2A\xD7\xB1')
                    if tcscale_pos != -1:
                        tcscale_pos = 4 + i * chunk_size + tcscale_pos + 3
                        file.seek(tcscale_pos)
                        tcscale_len = get_data_len(file.read(1))
                        tcscale = int(binascii.hexlify(file.read(tcscale_len)), 16)
                    duration_pos = bin.find(b'\x44\x89\x84') # float (4 bytes)
                    if duration_pos != -1:
                        duration_pos = 4 + i * chunk_size + duration_pos + 3
                        file.seek(duration_pos)
                        duration = struct.unpack('>f', file.read(4))[0]
                    if not duration: # double (8 bytes)
                        duration_pos = bin.find(b'\x44\x89\x88')
                        if duration_pos != -1:
                            duration_pos = 4 + i * chunk_size + duration_pos + 3
                            file.seek(duration_pos)
                            duration = struct.unpack('>d', file.read(8))[0]
                    if bin.find(b'\x1F\x43\xB6\x75') != -1:
                        # segment info should be before the clusters
                        break
                    i += 1
            duration = duration * tcscale / 1000000
            return suid, duration

        class Edition:
            def __init__(self):
                self.default = 0
                self.name = ['Default']
                self.hidden = 0
                self.ordered = 0
                self.num_chapters = 1
                self.uid = 0

        class Chapter:
            def __init__(self):
                self.name = ['Chapter']
                self.chapter = False
                self.start = False
                self.end = False
                self.suid = False
                self.hidden = 0
                self.uid = 0
                self.enabled = 1

    def __init__(self, templatefile, output=None, avs=None, trims=None,
                 kframes=None, uid=None, label=None, ifps=None, clip=None,
                 idr=False):
        try:
            import configparser
        except ImportError:
            import ConfigParser as configparser
        from io import open

        # Init config
        config = configparser.ConfigParser()
        template = open(templatefile, encoding='utf-8')

        # Read template
        config.readfp(template)
        template.close()

        # Template defaults
        self = self.Template()
        self.editions = []
        self.uid = uid if uid else self.uid

        # Set mkvinfo path
        from vfr import mkvmerge, parse_with_mkvmerge, fmt_time
        from os.path import dirname, join, isfile

        # Set placeholder for mkvinfo output
        mkv_globbed = False
        mkvinfo = {}

        for k, v in config.items('info'):
            if k == 'lang':
                self.lang = v.split(',')
            elif k == 'country':
                self.country = v.split(',')
            elif k == 'inputfps':
                self.fps = v
            elif k == 'outputfps':
                self.ofps = v
            elif k == 'createqpfile':
                self.qpf = v
            elif k == 'uid':
                self.uid = int(v)
            elif k == 'editions':
                self.num_editions = int(v)

        if avs and not ifps:
            self.connect_with_vfr(avs, label, clip)
        elif trims:
            self.trims = trims
            self.kframes = kframes
        else:
            self.trims = False
        self.idr = idr

        for i in range(self.num_editions):
            from re import compile
            ed = self.Edition()
            ed.uid = self.uid * 100
            self.uid += 1
            cuid = ed.uid
            ed.num = i+1
            ed.chapters = []
            stuff = {}

            for k, v in config.items('edition{0:d}'.format(ed.num)):
                if k == 'default':
                    ed.default = int(v)
                elif k == 'name':
                    ed.name = v.split(',')
                elif k == 'ordered':
                    ed.ordered = int(v)
                elif k == 'hidden':
                    ed.hidden = int(v)
                elif k == 'chapters':
                    ed.num_chapters = int(v)
                    for i in range(ed.num_chapters):
                        stuff[i+1] = []
                elif k == 'uid':
                    ed.uid = int(v)
                else:
                    opt_re = compile('(\d+)(\w+)')
                    ret = opt_re.search(k)
                    if ret:
                        stuff[int(ret.group(1))].append((ret.group(2),v))

            for j in range(ed.num_chapters):
                ch = self.Chapter()
                cuid += 1
                ch.uid = cuid
                ch.num = j+1

                for k, v in stuff[j+1]:
                    if k == 'name':
                        ch.name = v.split(',')
                    elif k == 'chapter':
                        ch.chapter = int(v)
                    elif k == 'start':
                        ch.start = v
                    elif k == 'end':
                        ch.end = v
                    elif k == 'suid':
                        ch.suid = v.strip() if ret else 0
                    elif k == 'hidden':
                        ch.hidden = int(v)
                    elif k == 'enabled':
                        ch.enabled = int(v)

                if ch.suid and not isfile(ch.suid):
                    ch.suid = ch.suid.replace('0x','').lower().replace(' ','') 

                if ch.chapter and not (ch.start and ch.end):
                    ch.start, ch.end = self.trims[ch.chapter-1] if self.trims else (ch.start, ch.end)
                elif ch.suid:
                    mkvfiles = []
                    if isfile(ch.suid):
                        mkvfiles = [ch.suid]
                    elif not mkv_globbed:
                        from glob import glob
                        mkvfiles = glob('*.mkv') + glob(join(dirname(avs),'*.mkv'))
                        mkv_globbed = True
                    if mkvfiles:
                        if parse_with_mkvmerge:
                            from subprocess import check_output
                            import json
                            for file in mkvfiles:
                                info = check_output([mkvmerge, '-i', '-F', 'json',
                                    '--output-charset', 'utf-8', file]).decode('utf-8')
                                try:
                                    props = json.loads(info).get("container", {}).get("properties", {})
                                    ch.suid = props.get("segment_uid", 0)
                                    duration = props.get("duration", 0)
                                    mkvinfo[ch.suid] = {'file': file,
                                        'duration': fmt_time(duration * 10**6)
                                                    if duration else 0}
                                except Exception:
                                    pass
                        else:
                            for file in mkvfiles:
                                ch.suid, duration = self.parse_mkv(file)
                                mkvinfo[ch.suid] = {'file': file, 
                                    'duration': fmt_time(duration * 10**6) 
                                                if duration else 0}
                    if not (ch.start or ch.end):
                        ch.start = fmt_time(0) if not ch.start else ch.start
                        ch.end = mkvinfo[ch.suid]['duration'] if not ch.end and (ch.suid in mkvinfo) else ch.end

                ed.chapters.append(ch)
            self.editions.append(ed)
        if output:
            self.toxml(output)


def main(args):

    template = args[0]
    output = args[1]
    avs = args[2] if len(args) == 3 else None

    chaps = AutoMKVChapters(template,output,avs)

if __name__ == '__main__':
    from sys import argv, exit
    if len(argv) > 1:
        main(argv[1:])
    else:
        exit("templates.py <template file> <output filenames w/o extension> [<avisynth file>]")