#!/bin/bash

# Create a standalone HOL Light image using DMTCP.
# The latest version of DMTCP can be downloaded from
# https://github.com/dmtcp/dmtcp .
# To run make-checkpoint.sh, dmtcp_launch and dmtcp_command must be
# evaluable from the PATH environment variable.

if [ $# -ne 1 ] && [ $# -ne 2 ]; then
  echo "make-checkpoint.sh <output script path> [additional OCaml statement to execute (optional, no ';;' added)]"
  exit 1
fi

if ! type "dmtcp_launch" > /dev/null; then
  echo "dmtcp_launch does not exist."
  exit 1
elif ! type "dmtcp_command" > /dev/null; then
  echo "dmtcp_command does not exist."
  exit 1
fi

# Create a directory to store all checkpoint binaries as well as the created
# DMTCP script.
output_script_path="$1"
output_dmtcp_dir_path="${1}.ckpt"

if [ -d ${output_dmtcp_dir_path} ] || [ -f ${output_dmtcp_dir_path} ]; then
  echo "${output_dmtcp_dir_path} already exists."
  exit 1
fi
mkdir -p ${output_dmtcp_dir_path}

# Create an .ml file for initialization
init_ml_path=$(mktemp --suffix=".ml")
# If HOL Light is fully loaded, "loaded" will be written to this file.
hol_is_loaded_log=$(mktemp --suffix=".hol_loaded")

echo '#use "hol.ml";;' >> ${init_ml_path}
if [ $# -eq 2 ]; then
  additional_ocaml_statement=$2
  echo "try begin" >> ${init_ml_path}
  echo "  ${additional_ocaml_statement};" >> ${init_ml_path}
  echo "  Printf.printf \"HOL Light loaded.\n\"" >> ${init_ml_path}
  echo "end with e -> print_endline (Printexc.to_string e);;" >> ${init_ml_path}
else
  echo "Printf.printf \"HOL Light loaded.\n\";;" >> ${init_ml_path}
fi
echo "let __ckpt_f = open_out \"${hol_is_loaded_log}\" in Printf.fprintf __ckpt_f \"loaded\"; close_out __ckpt_f;;" \
  >> ${init_ml_path}

# Use this init_ml_path when hol.sh begins
echo "Setting HOL_ML_PATH to ${init_ml_path}..."
export HOL_ML_PATH=${init_ml_path}

# Locate hol.sh
current_dir=$(dirname -- "$( readlink -f -- "$0"; )")
hol_sh_path="${current_dir}/hol.sh"
HOLLIGHT_DIR="$(${hol_sh_path} -dir)"

# Get any unused port. Use a python script to do this. Doing this directly
# on bash is frustratingly hard. :(
export DMTCP_COORD_PORT=$(python3 -c "import socket; s = socket.socket(); s.bind(('', 0));print(s.getsockname()[1]);s.close()")

function checkpoint_after_load () {
  while true; do
    sleep 1
    cat ${hol_is_loaded_log} | grep "loaded"
    if [ $? -eq 0 ]; then
      # HOL Light is loaded
      dmtcp_command -bc
      while [ ! -f "${output_dmtcp_dir_path}/dmtcp_restart_script.sh" ]; do
        sleep 1
      done
      echo "checkpointed"
      dmtcp_command -k
      echo "done"
      break
    fi
  done

  datenow="$(date)"
  rm -f ${output_script_path}
  echo "#!/bin/bash" >> ${output_script_path}
  echo "# HOL Light checkpointed (generated by make-checkpoint.sh)" \
    >> ${output_script_path}
  echo "# Contents of HOL_ML_PATH:" >> ${output_script_path}
  sed 's/^/#   /' ${init_ml_path} >> ${output_script_path}
  echo "echo \"          HOL Light (checkpointed, ${datenow})\"" \
    >> ${output_script_path}
  echo "args=\"${output_dmtcp_dir_path}/dmtcp_restart_script.sh \$@\"" \
    >> ${output_script_path}
  echo "bash -i -c \"\$args\"" \
    >> ${output_script_path}
  chmod +x ${output_script_path}
}

checkpoint_after_load &

# Launch hol.sh using dmtcp_launch.
bash -i -c "dmtcp_launch --new-coordinator --ckptdir ${output_dmtcp_dir_path} ${hol_sh_path}"
