#!/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 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.
    prerender_script = "{} -images {}".format(os.path.join(SCRIPT_DIR, "cioprep.py"), " ".join(args.image))
    cnode_args += ["-script"] + [prerender_script]

    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()
