#!/usr/bin/env python
import argparse
import os
import shutil
import subprocess
import tempfile

import OpenImageIO

from arch.path import Sequence

IGNORE_EXTNS = [".ifl", ".mov"]

MOV_EXTN = "mov"
TEMP_EXTN = "exr"


def parse_args():
    parser = argparse.ArgumentParser(
        description='Create video files from a directory of images.'
    )
    parser.add_argument('directory',
                        help='The path to a directory containing images.')
    return parser.parse_args()


def create_mov(sequence, directory, framerate):
    """
    Takes a directory containing one or more image sequences. Creates a .mov for
    the passed sequence, assuming that the sequence is in the directory.
    :param Sequence sequence: An Arch-style Sequence.
    :param str directory: A path to the directory containing the image files
    associated with the passed sequence.
    :param float framerate: The framerate which should be used when creating the
    .mov.
    """

    mov_path = None
    temp_dir = None
    # check if the image sequence in question consists of DWAA/DWAB exrs. If
    # so, create lossless exrs and use those. Otherwise, ffmpeg can use the
    # image sequence as it stands.

    bad_exrs = False

    if sequence.extension == 'exr':
        first_frame_path = os.path.join(directory,
                                        '.'.join([sequence.name,
                                                 str(sequence.first_frame),
                                                 sequence.extension]
                                                 )
                                        )
        buf = OpenImageIO.ImageBuf(first_frame_path)
        spec = buf.spec()
        compression = spec.getattribute('compression')
        # TODO: check if these are actually the correct formatting for the
        #  incompatible compression schemes.
        if compression == 'DWAA' or compression == 'DWAB':
            bad_exrs = True

    if bad_exrs:
        temp_dir = tempfile.mkdtemp()
        temp_path = create_temp_images(sequence, directory, temp_dir)
    else:
        temp_path = os.path.join(directory, sequence.pattern)

    if temp_path is not None:
        mov_path = os.path.join(
            directory,
            ".".join([sequence.name, MOV_EXTN])
                                )
        first_frame = sequence.first_frame
        # TODO: figure out how to offset frame number
        # offset = str(1.0 * first_frame / framerate)

        ffmpeg_args = ["ffmpeg", "-y",
                       "-start_number", str(first_frame),
                       "-framerate", str(framerate),
                       "-i", temp_path,
                       # TODO: figure out how to offset frame number
                       # "-output_ts_offset", offset,
                       mov_path,
                       ]

        proc = subprocess.Popen(ffmpeg_args,
                                stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE)
        _stdout, _stderr = proc.communicate()

    if temp_dir is not None:
        shutil.rmtree(temp_dir)

    return mov_path


def create_temp_images(sequence, input_dir, output_dir):
    """
    Converts all images in the passed sequence to .exrs, written to output_dir.
    :param Sequence sequence: An Arch-style Sequence.
    :param str input_dir: A directory which contains the images corresponding
    to sequence.
    :param str output_dir: A valid directory.
    :return str|None: The ffmpeg-ready path to the images or None if no images
    were written.
    """
    # TODO: throw errors
    if not os.path.exists(input_dir):
        print("The input directory, {}, doesn't exist.".format(input_dir))
    elif not os.path.exists(output_dir):
        print("The output directory, {}, doesn't exist.".format(output_dir))
    else:
        extension = sequence.extension
        if extension not in IGNORE_EXTNS:
            first_frame = sequence.first_frame
            last_frame = sequence.last_frame
            number_of_zeroes = sequence.padding[2]
            # create something which can be used to properly-format frame
            # numbers, like '{:0>2d}'.
            padding = "{{:0>{}d}}".format(number_of_zeroes)
            for frame_number in range(first_frame, last_frame + 1):
                frame_pattern = ".".join([sequence.name,
                                          padding.format(frame_number),
                                          extension]
                                         )
                frame_path = os.path.join(input_dir, frame_pattern)
                write_image(frame_path, output_dir)
            # this is the ffmpeg-ready path for the images created,
            # i.e. images/live/here/image.%04d.exr.
            temp_file_name = ".".join([sequence.name,
                                       "%0{}d".format(number_of_zeroes),
                                       TEMP_EXTN])
            return os.path.join(output_dir, temp_file_name)
        else:
            return None


def write_image(input_file_path, output_dir):
    """
    Writes a copy of the passed file as an EXR in the specified directory. The
    name of the file will be identical, except the extension.
    :param str input_file_path: A path to a valid image file.
    :param str output_dir: A path to a valid directory.
    """
    # generate output filename
    input_image_name = os.path.basename(input_file_path)
    output_image_name = ".".join([os.path.splitext(input_image_name)[0],
                                  TEMP_EXTN])
    output_image_path = os.path.join(output_dir, output_image_name)

    # create an image buffer
    input_buf = OpenImageIO.ImageBuf(input_file_path)
    # set the compression to none - this ensures that ffmpeg can read the exrs
    input_buf.specmod().attribute("compression", "none")

    # write the image
    if input_buf.has_error:
        print("Error: {}".format(input_buf.geterror()))
    else:
        input_buf.write(output_image_path)


def main():
    args = parse_args()

    sequences = Sequence.sequences_from_directory(args.directory)

    for sequence in sequences:
        print(create_mov(sequence, args.directory, 24.0))


if __name__ == "__main__":
    main()