#!/usr/bin/env python3
import argparse
import subprocess
import os
import os.path
from shlex import quote
from sys import stdin, stdout, stderr
from snakemake.utils import read_job_properties

parser=argparse.ArgumentParser(add_help=False)

parser.add_argument("--depend", help="Space separated list of ids for jobs this job should depend on.", default="", type=str)
parser.add_argument("--extra", help="Any additional qsub args", default="", type=str)
parser.add_argument("jobscript")

args = parser.parse_args()

cmd = ["qsub", "-l wd -j oe"]  # always start in wd, always join out/err logs

job_properties = read_job_properties(args.jobscript)
resources = job_properties.get("resources", {})
cluster = job_properties.get("cluster", {})

# dependencies
if args.depend:
    depstr = ":".join(args.depend.split(" "))
    cmd.append("-W")
    cmd.append(quote(f"depend=afterok:{depstr}"))

# threads
threads = job_properties.get("threads", 1)
cmd.append(f"-l 'ncpus={threads}'")

# memory
mem = int(resources.get("mem_gb", 2))
cmd.append(f"-l 'mem={mem}G'")

# walltime
time_min = resources.get("time_min", 60)
cmd.append(f"-l 'walltime={time_min}:00'")

# jobfs
localdisk_gb = resources.get("localdisk_gb", 10)
cmd.append(f"-l 'jobfs={localdisk_gb}G'")

# storage
# a special NCI pbs resource that mounts FSes

if "storage" in cluster:
    cmd.append(f"-l storage={quote(cluster['storage'])}")

project = cluster.get("project", "xe2")
cmd.append(f"-P {quote(project)}")

# queue
# defaults: normal, unless internet=true
queue = "normal"
if resources.get("internet", False):
    queue="copyq"
# override with queue from resources
queue = resources.get("queue", queue)
cmd.append(f"-q '{queue}'")

# logs
out = None
if "output" in cluster:
    out = cluster["output"]
if "output" in resources:
    out = resources["output"]
if out:
    os.makedirs(out, exist_ok=True, mode=0o775)
    cmd.append(f"-o {out}")

cmd.append(args.extra)
cmd.append(args.jobscript)
cmd = " ".join(cmd)

if cluster.get("DEBUG", False):
    print(cmd, file=stderr)
res = subprocess.check_output(cmd, shell=True).decode()
print(res.strip())
