"""
A helper script to create SLURM submission scripts for all the assets
defined in asset_build_packages for a given genome
"""
from .asset_build_packages import asset_build_packages
from ubiquerg import expandpath
import os
import argparse
import divvy

parser = argparse.ArgumentParser(
    description="Builds submission scripts for all assets for a genome"
)
parser.add_argument(
    "-g",
    "--genome",
    dest="genome",
    type=str,
    help="genome to build the submission scripts for",
)
parser.add_argument(
    "-p",
    "--path",
    dest="path",
    type=str,
    help="path to the desired submission directory location",
)
parser.add_argument(
    "-pt",
    "--partition",
    dest="PARTITION",
    type=str,
    help="partition in SLURM submission script",
    default="standard",
)
parser.add_argument(
    "-m",
    "--mem",
    dest="MEM",
    type=str,
    help="mem in SLURM submission script",
    default="200000",
)
parser.add_argument(
    "-t",
    "--time",
    dest="TIME",
    type=str,
    help="time in SLURM submission script",
    default="10:00:00",
)
parser.add_argument(
    "-c",
    "--cores",
    dest="CORES",
    type=str,
    help="cpus-per-task in SLURM submission script",
    default="4",
)
parser.add_argument(
    "-o",
    "--output",
    dest="LOGFILE",
    type=str,
    help="output in SLURM submission script",
    default=None,
)
parser.add_argument(
    "-j",
    "--job-name",
    dest="JOBNAME",
    type=str,
    help="job-name in SLURM submission script",
    default=None,
)

args = parser.parse_args()


def _make_sub_dir(path, genome):
    """
    create submission scripts directory

    :param str path: path where the submission scripts directory should be created
    :param str genome: name of the genome
    :return str: created path
    """
    path = os.path.join(expandpath(path), "submission_scripts", genome)
    if not os.path.exists(path):
        os.makedirs(path)
    return path


def _req_input_to_args(req_input):
    """
    Given a list of the required inputs for the build command, create an args
    string

    :param list[str] req_input: input names
    :return str: args string
    """
    return ["--" + x + " <arg_here>" for x in req_input]


subdir_path = _make_sub_dir(args.path, args.genome)
dcc = divvy.ComputingConfiguration()
dcc.activate_package("slurm")
cmd_template = "refgenie build -g {g} -a {a} {req_input_str}"
genome = args.genome
to_remove = ["genome", "path"]

data = vars(args)
for i in to_remove:
    data.pop(i)

for asset in asset_build_packages:
    sub_script = os.path.join(subdir_path, asset + ".sub")
    req_input = asset_build_packages[asset]["required_inputs"]
    if req_input:
        print(
            "{} asset requires additional input in the command ({}), so '{}'"
            " requires manual edit".format(asset, req_input, sub_script)
        )
        req_str = " ".join(_req_input_to_args(req_input))
    else:
        req_str = ""
    data["CODE"] = cmd_template.format(g=genome, a=asset, req_input_str=req_str)
    data["LOGFILE"] = asset + ".log"
    data["JOBNAME"] = asset + "Build"
    dcc.write_script(sub_script, data)
