#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# run as: blender -b BLENDFILE -P distributed_render.py
# outputs to 'dist_output.mkv' in lossless 'FFV1 + PCM' in Matroska container.
#
# (c) 2019 Bernd Busse
#

"""Blender script to export animation in lossess format using all cores.

Export blender animation in lossless FFV1 video- and PCM audio-codec.
Result is stored as 'dist-output_ISODATE.mkv' in Matroska container.

Run as: blender -p BLENDFILE -P distributed_render.py [--buildroot PATH]

The output will be placed in 'output/' next to BLENDFILE or in 'PATH/output/' if specified.
"""

import sys
import os
import datetime as dt
import subprocess

from concurrent import futures
from concurrent.futures import ThreadPoolExecutor

import bpy
from bpy import ops as op
from bpy import context as ctx

DIST_OUTPUT = "dist_output"
FINAL_OUTPUT = "output"
FRAME_COUNT = 250
WORKER_COUNT = 4


def eprint(*args, **kwargs):
    """Print to STDERR."""
    return print("Error:", *args, file=sys.stderr, **kwargs)


def print_usage():
    """Print help message."""
    print(__doc__, file=sys.stderr, end="")


def export_audio(dest):
    """Export complete sound."""
    op.sound.mixdown('EXEC_DEFAULT',
                     filepath=os.path.join(dest, "audio.wav"),
                     check_existing=False,
                     container='WAV',
                     codec='PCM',
                     format='S16')


def render_frames(dest, start, end):
    """Export animation range (start, end]."""
    # Get current scene and renderer
    scene = ctx.scene
    render = scene.render

    # Set start and end frames
    scene.frame_start = start
    scene.frame_end = end

    # Set video codec
    render.image_settings.file_format = 'FFMPEG'
    render.ffmpeg.format = 'MKV'
    render.ffmpeg.codec = 'FFV1'
    render.ffmpeg.audio_codec = 'NONE'

    # Set output path
    render.filepath = os.path.join(dest, "{:05d}-{:05d}.mkv".format(start, end))

    # Render given frames
    op.render.render(animation=True)


def main(blend_argv, argv):
    """Process `argv` and start export worker."""
    # Parse '--buildroot' argument
    buildroot = "//"  # default to BLENDFILE path
    if "--buildroot" in argv:
        opt_idx = argv.index("--buildroot") + 1
        if opt_idx >= len(argv):
            raise RuntimeError("Missing argument for '--buildroot'")

        buildroot = argv[opt_idx]

        del argv[opt_idx]
        del argv[opt_idx - 1]

    # Parse '--out' argument    
    if "--out" in argv:
        opt_idx = argv.index("--out") + 1
        if opt_idx >= len(argv):
            raise RuntimeError("Missing argument for '--out'")

        outname = argv[opt_idx]

        del argv[opt_idx]
        del argv[opt_idx - 1]

    buildroot = os.path.abspath(bpy.path.abspath(buildroot))
    if not os.path.exists(buildroot):
        raise RuntimeError("buildroot path does not exist: '{}'".format(buildroot))
    elif not os.path.isdir(buildroot):
        raise RuntimeError("buildroot is not a directory: '{}'".format(buildroot))

    # Get blendfile path
    output = os.path.join(buildroot, FINAL_OUTPUT)
    dist_output = os.path.join(buildroot, DIST_OUTPUT)
    if not os.path.exists(output):
        os.mkdir(output)
    if not os.path.exists(dist_output):
        os.mkdir(dist_output)

    script_args = ['--buildroot', buildroot]

    # Handle different actions
    if len(argv) == 0:
        print(" == [ distributed_render.py v0.2 for blender ] == ")
        print("Export to {}".format(buildroot))

        # Get current scene
        scene = ctx.scene

        # Get start and end frame
        scene.frame_start
        scene.frame_end

        dt_start = dt.datetime.now()
        print("Start rendering of {} frames at {}"
              .format(scene.frame_end - scene.frame_start + 1,
                      dt_start.time().isoformat()))

        # Export sound
        def export_audio_async(blender_cmd, args):
            print("Export audio")
            task_start = dt.datetime.now()
            subprocess.run(blender_cmd + ['--', 'audio'] + args,
                           check=True, stdout=subprocess.DEVNULL)
            task_end = dt.datetime.now()
            print("Finshed exporting audio. Time elapsed: {}"
                  .format(str(task_end - task_start)))

        # Render frames
        def render_frames_async(blender_cmd, start, end, args):
            print("Export frames " + start + " to " + end)
            task_start = dt.datetime.now()
            subprocess.run(blender_cmd + ['--', 'frames', start, end] + args,
                           check=True, stdout=subprocess.DEVNULL).check_returncode()
            task_end = dt.datetime.now()
            print("Finished exporting frames {} to {}. Time elapsed: {}"
                  .format(start, end, str(task_end - task_start)))

        # Start pool execution
        tasks = []
        frames = []
        with ThreadPoolExecutor(max_workers=WORKER_COUNT) as tpe:
            task = tpe.submit(export_audio_async, blend_argv, script_args)
            tasks.append(task)

            for i in range(scene.frame_start,
                           scene.frame_end + 1,
                           FRAME_COUNT):
                start_frame = i
                end_frame = min(i + FRAME_COUNT - 1, scene.frame_end)
                task = tpe.submit(render_frames_async, blend_argv,
                                  str(start_frame), str(end_frame),
                                  script_args)
                tasks.append(task)
                frames.append((start_frame, end_frame))

            # Check return codes
            for fut in futures.as_completed(tasks):
                if fut.exception():
                    for task in tasks:
                        task.cancel()
                    raise fut.exception()

        if outname:
            filename = outname + ".mkv"
        else:
            filename = "dist_output_{:05d}-{:05d}.mkv".format(scene.frame_start,
                                                              scene.frame_end)

        ffmpeg_args = ["ffmpeg", "-i", os.path.join(dist_output, "audio.wav"),
                       "-safe", "0", "-protocol_whitelist", "file,pipe",
                       "-f", "concat", "-i", "pipe:0",
                       "-c:v", "copy", "-c:a", "copy",
                       "-y", os.path.join(output, filename)]

        ffmpeg = subprocess.Popen(ffmpeg_args,
                                  stdin=subprocess.PIPE,
                                  universal_newlines=False)
        files = ["file '{}/{:05d}-{:05d}.mkv'\n"
                 .format(dist_output, start, end).encode('utf8')
                 for start, end in frames]
        ffmpeg.stdin.writelines(files)
        ffmpeg.stdin.flush()
        ffmpeg.stdin.close()
        retval = ffmpeg.wait()
        if retval is not None and retval != 0:
            eprint("merging output files with ffmpeg failed")

        print("Cleanup intermediate files")
        os.remove(os.path.join(dist_output, "audio.wav"))
        for start, end in frames:
            os.remove("{}/{:05d}-{:05d}.mkv".format(dist_output, start, end))
        try:
            os.rmdir(dist_output)
        except OSError:
            pass

        dt_end = dt.datetime.now()
        print("Finished rendering of {} frames at {}. Time elapsed: {}"
              .format(scene.frame_end - scene.frame_start + 1,
                      dt_end.time().isoformat(),
                      str(dt_end - dt_start)))

    elif argv[0] == 'audio':
        # Export sound
        print("Export audio")
        export_audio(dist_output)
    elif argv[0] == 'frames':
        # Render given frames
        print("Export frames " + argv[1] + " to " + argv[2])
        render_frames(dist_output, int(argv[1]), int(argv[2]))
    elif argv[0] in ('help', 'h'):
        # Display help message
        print_usage()
        return
    else:
        raise RuntimeError("Unsupported action: " + str(argv))


if __name__ == '__main__':
    try:
        index = sys.argv.index('--')
    except ValueError:
        index = len(sys.argv)

    try:
        main(sys.argv[:index], sys.argv[index + 1:])
    except RuntimeError as err:
        eprint(str(err))
        sys.exit(3)
    except Exception as ex:
        eprint("An unhandled Exception occured: " + str(ex))
        sys.exit(1)