#!/bin/bash

set -euxo pipefail

############################################################
# Default Values                                           #
############################################################

nodes=1
cores=16
mem_per_cpu="8G"
time="24:00:00"
partition="amd-hdr100,medium"
outdir="/data/rc/gpfs-policy/data"
policy="./policy-def/list-path-external"
outfile=""
dry_run=""

############################################################
# Help                                                     #
############################################################
usage()
{
>&2 cat << EOF
Usage: $0 [ -h ] [ -o | --outdir ] [ -f | --outfile ] [ --with-dirs ] 
          [ -N | --nodes ] [ -c | --cores ] [ -p | --partition] 
          [ -t | --time ] [ -m | --mem-per-cpu ] [ --dry_run ]
          device
EOF
exit 1
}

help()
{
>&2 cat << EOF
Wraps the  run-mmpol.sh script for applying a policy file. Must be run directly 
as root or via the run-submit-pol-job.py script. The default policy file is
./policy/list-path-external

Usage: $0 [ -h ] [ -o | --outdir ] [ -f | --outfile ] [ -P | --policy ] 
          [ -N | --nodes ] [ -c | --cores ] [ -p | --partition ] 
          [ -t | --time ] [ -m | --mem ] [ --dry-run ]
          device

options:
    -h|--help           Print this Help.
    --dry-run           Do not submit a Slurm job running the policy. Instead, pass --dry-run to run-mmpol.sh and call 
                            it normally to just print the output to STDOUT

Required:
    device              GPFS fileset/directory apply the policy to. Can be 
                            specified as either the name of the fileset or the 
                            full path to the directory 
                            (Examples: scratch, /data/user/[username])

Path:   
    -o|--outdir         Parent directory to save policy output to 
                            (default: /data/rc/gpfs-policy/data)
    -f|--outfile        Name of the policy output file

Policy Options:
    -P|--policy         Path to policy file to apply to the given GPFS device

sbatch options:
    -N|--nodes          Number of nodes to run the job on (default: 1)
    -c|--cores         Number of cores (default: 16)
    -p|--partition      Partition to submit tasks to 
                            (default: amd-hdr100,medium)
    -t|--time           Max walltime (default: 24:00:00)
    -m|--mem-per-cpu    RAM per task (default: 8G)
EOF
exit 0
}

args=$(getopt -a -o ho:f:P:N:c:p:t:m: \
                 --long help,outdir:,outfile:,policy:,nodes:,cores:,partition:,time:,mem:,dry-run -- "$@")

if [[ $? -gt 0 ]]; then
  usage
fi

eval set -- ${args}

while :
do
  case $1 in
    -h | --help)        help            ;;
    -o | --outdir)      outdir=$2       ; shift 2 ;;
    -f | --outfile)     outfile=$2      ; shift 2 ;;
    -P | --policy)      policy=$2       ; shift 2 ;;
    -N | --nodes)       nodes=$2        ; shift 2 ;;
    -c | --cores)       cores=$2        ; shift 2 ;;
    -p | --partition)   partition=$2    ; shift 2 ;;
    -t | --time)        time=$2         ; shift 2 ;;
    -m | --mem-per-cpu) mem_per_cpu=$2  ; shift 2 ;;
    --dry-run)          dry_run=true    ; shift 1 ;;
    --) shift; break ;;
    *) >&2 echo Unsupported option: $1
       usage ;;
  esac
done

if [[ $# -eq 0 ]]; then
  usage
fi

device="$1"

# Ensure gpfs_logdir is set
if [[ -z "$device" ]]; then
    echo "Error: Specify either the name of a fileset or a directory path"
    usage
fi

slurm_out="out/pol-%A-$(basename ${policy})-$(basename ${device}).out"

run_mmpol_cmd_base="./run-mmpol.sh -o ${outdir} -f ${outfile} -P ${policy}"

if [[ -z "${dry_run}" ]]; then
    mkdir -p out

    run_mmpol_cmd="${run_mmpol_cmd_base} ${device}"

    sbatch \
       -N $nodes \
       -c $cores \
       -t $time \
       --mem-per-cpu=$mem_per_cpu \
       -p $partition \
       -o ${slurm_out} \
       --wrap "${run_mmpol_cmd}"
else
    run_mmpol_cmd="${run_mmpol_cmd_base} --dry-run ${device}"
    ${run_mmpol_cmd}
fi