#!/usr/bin/env python

"""
Wrapper which runs on the render node (or locally) to render a task.

It accepts any args that cnode accepts.
"""

import argparse
import os
import re
import subprocess
import sys

NUMBER_RE = re.compile(r"^(\d+)$")
RANGE_RE = re.compile(r"^(?:(\d+):(\d+)(?:\%(\d+))?)+$")
BIG_NUMBER = 99999999

# Path containing this script and other aux files, like the Clarisse config.
SCRIPT_DIR = os.path.dirname(sys.argv[0])
DEFAULT_CONFIG_FILE = os.path.join(SCRIPT_DIR, "clarisse.cfg")


def get_bounding_frames(*fame_specs):
    """
    Get frames to cover the range being rendered.

    We do this because Clarisse doesn't render images unless they have their
    sequence attributes set, even if the command line flags specify what frames
    to render. Therefore we have to calculate the frames needed and turn them
    on.

    Raises:
        ValueError: Raises if the spec is malformed.

    Returns:
        tuple(int, int): bounding frames
    """
    min_frame = BIG_NUMBER
    max_frame = -BIG_NUMBER

    for spec in fame_specs:
        number_match = NUMBER_RE.match(spec)
        range_match = RANGE_RE.match(spec)

        if not (range_match or number_match):
            raise ValueError("Spec format must be 'start[:end[%step]]")
        if range_match:
            start, end, _ = [int(n or 1) for n in range_match.groups()]
        else:
            start = end = int(number_match.groups()[0])

        if start < min_frame:
            min_frame = start
        if end > max_frame:
            max_frame = end

    return min_frame, max_frame


def generate_prerender_script_arg(args):
    """
    Constructs pre render script with some useful information as args.

    Example:
        /tmp/conductor/cioprep.py -range 1 100 -images project://scene/imgHigh
        project://scene/imgLow

    Args:
        images (list): list of image paths
        range (int, int) The bounding range being rendered.

    Note the choice of images instead of image. This is because you can't
    provide argument names to the pre-render script that are also known by
    cnode. Cnode already has an image arg.

    Returns:
        string: script with args
    """

    prep_script = os.path.join(SCRIPT_DIR, "cioprep.py")
    min_frame, max_frame = get_bounding_frames(*args.image_frames_list)
    return "{} -range {:d} {:d} -images {}".format(
        prep_script, min_frame, max_frame, " ".join(args.image)
    )


def main():
    """
    Main function to run cnode with conductor adjustments.
    """

    parser = argparse.ArgumentParser()

    # images and image frames list are required.
    parser.add_argument(
        "-image_frames_list",
        nargs="*",
        type=str,
        required=True,
        help="List of image frame specs.",
    )
    parser.add_argument(
        "-image", nargs="*", type=str, help="List of images.", required=True
    )
    # User can override the following flags in the task template.

    # Currently the users default clarisse.cfg is copied to /tmp/conductor
    # during submission. User may want a special Config to be in place on the
    # rendernode, in which case they can provide one to override this.
    parser.add_argument(
        "-config_file",
        type=str,
        default=DEFAULT_CONFIG_FILE,
        help="Specify the config file used by CNODE. (default: %(default)s)",
    )

    # Logs in the Conductor UI are more readable when not wrapped at 80 columns.
    parser.add_argument(
        "-log_width",
        type=int,
        default=0,
        help="Set the width of the log. 0 means unlimited. (default: %(default)s)",
    )

    parser.add_argument(
        "-tile_rendering",
        type=int,
        nargs=2,
        default=[1, 1],
        help="Specify : number_of_tiles tile_index (1-based).",
    )

    args, pass_through_args = parser.parse_known_args()

    # We will change to the working directory because cnode cannot deal with
    # spaces in the project filepath. We have eliminated spaces from the project
    # filename. But the directory name is beyond our control, so we cd into it
    # and work from there. https://www.isotropix.com/user/bugtracker/363 .
    proj = pass_through_args[0]
    work_dir, basename = (os.path.dirname(proj), os.path.basename(proj))
    pass_through_args[0] = basename

    cnode_args = ["cnode"] + pass_through_args

    # Add some defaults.
    cnode_args += ["-config_file", args.config_file, "-log_width", str(args.log_width)]

    # Adds a tile rendering flag only if there are actually multiple tiles. We
    # don't want to add the flag in the case of one tile, because it will put
    # extraneous tile number info in the filenames.
    if args.tile_rendering[0] > 1:
        cnode_args += [
            "-tile_rendering",
            str(args.tile_rendering[0]),
            str(args.tile_rendering[1]),
        ]

    # These flags are required so no need to check their existence here.
    cnode_args += ["-image"] + args.image
    cnode_args += ["-image_frames_list"] + args.image_frames_list

    # script arg must go at the end.
    cnode_args += ["-script"] + [generate_prerender_script_arg(args)]

    sys.stdout.write(
        "Running command:..\n{}\nfrom directory:\n{}\n".format(
            " ".join(cnode_args), work_dir
        )
    )
    os.chdir(work_dir)

    p = subprocess.Popen(cnode_args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    for line in iter(p.stdout.readline, b""):
        sys.stdout.write(line)
        sys.stdout.flush()


main()
