#! /usr/bin/python3

import os, sys, subprocess, json

DEFAULT_QUALITY = 23
DEFAULT_AV1_PRESET = 5
DEFAULT_STEREO_BANDWIDTH = "128"
DEFAULT_AC3_BANDWIDTH = "256"
DEFAULT_DTS_BANDWIDTH = "320"

TEMP_FILE_NAME = "ffmpeg2pass-0.log"

def executeProcess(commandSequence):

    process = subprocess.Popen(commandSequence, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    output, error = process.communicate()

    return output


def generateAV1Tokens(q, p):

    return ['-c:v:0',
            'libsvtav1',
            '-svtav1-params',
            f"crf={q}:preset={p}:tune=0:enable-overlays=1:scd=1:scm=0",
            '-pix_fmt', 'yuv420p10le']

def generateVP9Pass1Tokens(q):

    return ['-c:v:0', 'libvpx-vp9', '-row-mt', '1',
                '-crf', str(q), '-pass', '1', '-speed', '4',
                '-frame-parallel', '0', '-g', '9999', '-aq-mode', '0']

def generateVP9Pass2Tokens(q):

    return ['-c:v:0', 'libvpx-vp9', '-row-mt', '1',
                '-crf', str(q), '-pass', '2',
                '-frame-parallel', '0', '-g', '9999', '-aq-mode', '0',
                '-auto-alt-ref', '1', '-lag-in-frames', '25']

def generateCropTokens(start, length):

    return ['-ss', str(start), '-t', str(length)]

def generateOutputTokens(f, q=''):
    if q:
        fTokens = f.split('.')
        paddedFilename = '.'.join(fTokens[:-1]) + f" q{q}" + '.' + fTokens[-1]
        return ['-f', 'webm', paddedFilename]
    else:
        return ['-f', 'webm', f]


inputFilename = sys.argv[1]
outputFilename = sys.argv[2]

targetFormat = 'vp9'
if 'av1' in sys.argv:
    targetFormat = 'av1'
if 'vp9' in sys.argv:
    targetFormat = 'vp9'


qualities = [str(DEFAULT_QUALITY)]
qualitiesTokens = [q for q in sys.argv if q.startswith('q=')]
if qualitiesTokens:
  qualitiesString = qualitiesTokens[0].split('=')[1]
  qualities = qualitiesString.split(',')

preset = DEFAULT_AV1_PRESET
presetTokens = [p for p in sys.argv if p.startswith('p=')]
if presetTokens:
  preset = int(presetTokens[0].split('=')[1])

stereoBandwidth = DEFAULT_STEREO_BANDWIDTH
stereoTokens = [s for s in sys.argv if s.startswith('a=')]
if stereoTokens:
   stereoBandwidth = str(stereoTokens[0].split('=')[1])
if not stereoBandwidth.endswith('k'):
    stereoBandwidth += "k"

ac3Bandwidth = DEFAULT_AC3_BANDWIDTH
ac3Tokens = [a for a in sys.argv if a.startswith('ac3=')]
if ac3Tokens:
   ac3Bandwidth = str(ac3Tokens[0].split('=')[1])
if not ac3Bandwidth.endswith('k'):
    ac3Bandwidth += "k"

dtsBandwidth = DEFAULT_DTS_BANDWIDTH
dtsTokens = [d for d in sys.argv if d.startswith('dts=')]
if dtsTokens:
  dtsBandwidth = str(dtsTokens[0].split('=')[1])
if not dtsBandwidth.endswith('k'):
    dtsBandwidth += "k"

cropStart = ''
cropLength = ''
cropTokens = [c for c in sys.argv if c.startswith('crop')]
if cropTokens:
    if '=' in cropTokens[0]:
        cropString = cropTokens[0].split('=')[1]
        cropStart, cropLength = cropString.split(',')
    else:
        cropStart = 60
        cropLength = 180


output = executeProcess(["ffprobe", "-show_streams", "-of", "json" ,inputFilename])

streamData = json.loads(output)['streams']

videoStreams = [s for s in streamData if s['codec_type'] == 'video']
audioStreams = [s for s in streamData if s['codec_type'] == 'audio']
subtitleStreams = [s for s in streamData if s['codec_type'] == 'subtitle']

for aStream in audioStreams:
    if 'channel_layout' in aStream:
        print(f"audio stream: {aStream['channel_layout']}") #channel_layout
    else:
        print(f"unknown audio stream with {aStream['channels']} channels") #channel_layout

commandTokens = ['ffmpeg', '-y', '-i', inputFilename]


mappingTokens = ['-map', 'v:0']

for a in range(len(audioStreams)):
    mappingTokens += ['-map', f"a:{a}"]

for s in range(len(subtitleStreams)):
    mappingTokens += ['-map', f"s:{s}"]


audioTokens = []
audioStreamIndex = 0
for aStream in audioStreams:

    channels = aStream['channels']

    if 'channel_layout' in aStream.keys():

        channelLayout = aStream['channel_layout']

        if channelLayout == '6.1':
            audioTokens += [f"-c:a:{audioStreamIndex}",
                            'libopus',
                            f"-filter:a:{audioStreamIndex}",
                            'channelmap=channel_layout=6.1',
                            f"-b:a:{audioStreamIndex}",
                            dtsBandwidth]

        if channelLayout == '5.1(side)':
            audioTokens += [f"-c:a:{audioStreamIndex}",
                            'libopus',
                            f"-filter:a:{audioStreamIndex}",
                            "channelmap=FL-FL|FR-FR|FC-FC|LFE-LFE|SL-BL|SR-BR:5.1",
                            f"-b:a:{audioStreamIndex}",
                            ac3Bandwidth]

        if channelLayout == 'stereo':
            audioTokens += [f"-c:a:{audioStreamIndex}",
                            'libopus',
                            f"-b:a:{audioStreamIndex}",
                            stereoBandwidth]
    else:

        if channels == 6:
            audioTokens += [f"-c:a:{audioStreamIndex}",
                            'libopus',
                            f"-b:a:{audioStreamIndex}",
                            ac3Bandwidth]


    audioStreamIndex += 1


nullTokens = ['-f', 'null', '/dev/null']


for quality in qualities:

    if targetFormat == 'av1':

        commandSequence = commandTokens + mappingTokens + generateAV1Tokens(quality, preset) + audioTokens

        if len(qualities) > 1:
            commandSequence += generateOutputTokens(outputFilename, quality)
        else:
            commandSequence += generateOutputTokens(outputFilename)
    
        if cropStart:
            commandSequence += generateCropTokens(cropStart, cropLength)


        print(f"Command Sequence: {commandSequence}")
    
        executeProcess(commandSequence)


    if targetFormat == 'vp9':

        commandSequence1 = commandTokens + mappingTokens + generateVP9Pass1Tokens(quality)


        if cropStart:
            commandSequence1 += generateCropTokens(cropStart, cropLength)

        commandSequence1 += nullTokens


        print(f"Command Sequence 1: {commandSequence1}")

        if os.path.exists(TEMP_FILE_NAME):
            os.remove(TEMP_FILE_NAME)

        executeProcess(commandSequence1)
    
        commandSequence2 = commandTokens + mappingTokens + generateVP9Pass2Tokens(quality) + audioTokens


        if cropStart:
            commandSequence2 += generateCropTokens(cropStart, cropLength)


        if len(qualities) > 1:
            commandSequence2 += generateOutputTokens(outputFilename, quality)
        else:
            commandSequence2 += generateOutputTokens(outputFilename)


        print(f"Command Sequence 2: {commandSequence2}")
    
        executeProcess(commandSequence2)