package io.compgen.cgpipe.runner;

import io.compgen.cgpipe.exceptions.RunnerException;
import io.compgen.cgpipe.parser.context.ExecContext;
import io.compgen.cgpipe.parser.variable.VarString;
import io.compgen.common.StringUtils;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

public class SLURMTemplateRunner extends TemplateRunner {
	private String account=null;

	@Override
	public String[] getSubCommand() {
		return new String[] {"sbatch", "--parsable"};
	}

	@Override
	public String[] getReleaseCommand(String jobId) {
		return new String[] {"scontrol", "release", jobId };
	}

	@Override
	public String[] getDelCommand(String jobId) {
		return new String[] {"scancel", jobId};
	}

	@Override
	public boolean isJobIdValid(String jobId) throws RunnerException {
		try {
			Process proc = Runtime.getRuntime().exec(new String[] {"scontrol", "show", "job", jobId});
			int retcode = proc.waitFor();
			if (retcode == 0) {	
				return true;
			}
		} catch (IOException | InterruptedException e) {
		}
		return false;
	}

	protected void updateTemplateContext(ExecContext cxt, JobDef jobdef) {
		if (this.account != null && !cxt.contains("job.account")) {
			cxt.set("job.account",  new VarString(this.account));
		}

		// set the dep list
		if (jobdef.getDependencies().size() > 0) {
		    List<String> depids = new ArrayList<String>();
		    
		    for (JobDependency dep: jobdef.getDependencies()) {
				if (!dep.getJobId().equals("")) {
				    depids.add(dep.getJobId());
				}
		    }
		    
		    cxt.set("job.slurm.depids", new VarString(StringUtils.join(":", depids).replaceAll("::", ":")));
		}

		super.updateTemplateContext(cxt, jobdef);
	}
	
	@Override
	protected void setConfig(String k, String val) {
		switch(k) {
		case "cgpipe.runner.slurm.account":
			this.account = val;
			break;
		default:
			super.setConfig(k, val);
			break;
		}

	}

	@Override
	public String getConfigPrefix() {
		return "cgpipe.runner.slurm";
	}

	protected String buildGlobalHoldScript() {
        return 	"#!" + shell + "\n" +
        		"#SBATCH -H\n" +
        		"#SBATCH -J holding\n" +
        		"#SBATCH -o /dev/null\n" +
        		"#SBATCH -e /dev/null\n" +
        		"#SBATCH -t 00:00:30\n" +
        		"sleep 1\n";
	}

}

@@ -1636,16 +1636,700 @@
 );%0A%09%09%7D%0A%0A
+%09%09// convert 4G to 4000; SLURM uses mem definitions in terms of megabytes.%0A%09%09String mem = jobdef.getSetting(%22job.mem%22);%0A%09%09%0A%09%09String units = %22%22;%0A%09%09float memVal = -1;%0A%09%09%0A%09%09if (mem != null) %7B%0A%09%09%09while (mem.length() %3E 0) %7B%0A%09%09%09    try %7B%0A%09%09%09%09%09memVal = Float.parseFloat(mem);%0A%09%09%09%09%09break;%0A%09%09%09    %7D catch (NumberFormatException e) %7B%0A%09%09%09%09%09units = mem.substring(mem.length()-1) + units;%0A%09%09%09%09%09mem = mem.substring(0, mem.length()-1);%0A%09%09%09    %7D%0A%09%09%09%7D%0A%09%0A%09%09%09if (memVal %3E 0) %7B%0A%09%09%09%09mem = Float.toString(memVal);%0A%09%09%09%09if (units.equalsIgnoreCase(%22G%22) %7C%7C units.equalsIgnoreCase(%22GB%22)) %7B%0A%09%09%09%09%09mem = %22%22+Float.valueOf(memVal * 1000).intValue();%0A%09%09%09%09%7D%0A%09%09%09%7D%0A%09%0A%09%09%09cxt.set(%22job.mem%22, new VarString(mem));%0A%09%09%7D%0A%09%09%0A%09%09%0A
 %09%09super.

Convert 4G to 4000 for SLURM