#!/bin/bash

set -ex

WITH_REGISTRY=false
NODES=1
PREFIX="kubevirt-"
SCRIPTS="$PWD/scripts"
MEMORY=""
CPU=""
QEMU_ARGS=""
BACKGROUND=""
NFS_DATA=""
REVERSE=""
RANDOM_PORTS=""
CRIO=""
K8S_VERSION=""

COMMAND=$1
shift

if [[    "${COMMAND}" != "run" \
     && "${COMMAND}" != "provision" \
     && "${COMMAND}" != "ssh" \
     && "${COMMAND}" != "ps" \
     && "${COMMAND}" != "update" \
     && "${COMMAND}" != "port" \
     && "${COMMAND}" != "rm" ]]; then
        echo "No valid command provides. Valid commands  are 'run', 'provision', 'ssh' and 'rm'."
        exit 1
fi

while true; do
  case "$1" in
    -p | --prefix ) PREFIX="${2}-"; shift 2 ;;
    -n | --nodes ) NODES="$2"; shift 2 ;;
    -s | --scripts ) SCRIPTS="$2"; shift 2 ;;
    -t | --tag ) TAG="$2"; shift 2 ;;
    -m | --memory ) MEMORY="$2"; shift 2 ;;
    -c | --cpu ) CPU="$2"; shift 2 ;;
    -q | --qemu-args ) QEMU_ARGS="$2"; shift 2 ;;
    -r | --reverse ) REVERSE="true"; shift ;;
    --background ) BACKGROUND="true"; shift ;;
    -b | --base ) BASE="$2"; shift 2 ;;
    --k8s-port ) K8S_PORT="$2"; shift 2 ;;
    --osp-port ) OSP_PORT="$2"; shift 2 ;;
    --ssh-port ) SSH_PORT="$2"; shift 2 ;;
    --vnc-port ) VNC_PORT="$2"; shift 2 ;;
    --random-ports ) RANDOM_PORTS="true"; shift 1 ;;
    --nfs-data ) NFS_DATA="$2"; shift 2 ;;
    --registry-port ) REGISTRY_PORT="$2"; shift 2 ;;
    --registry-volume ) REGISTRY_VOLUME="$2"; shift 2 ;;
    --k8s-version ) K8S_VERSION="$2"; shift 2 ;;
    --crio ) CRIO="true" ; shift ;;
    -- ) shift; break ;;
    * ) break ;;
  esac
done

if [ "$COMMAND" = "port" ] ; then
  dnsmasq_container="$(docker ps --filter="name=${PREFIX}dnsmasq" -q)"

  if [ "$1" = "k8s" ]; then
    docker port $dnsmasq_container 6443 | cut -d':' -f2
    exit
  fi
  if [ "$1" = "osp" ]; then
    docker port $dnsmasq_container 8443 | cut -d':' -f2
    exit
  fi
  if [ "$1" = "ssh" ]; then
    docker port $dnsmasq_container 2201 | cut -d':' -f2
    exit
  fi
  if [ "$1" = "registry" ]; then
    docker port $dnsmasq_container 5000 | cut -d':' -f2
    exit
  fi
  docker port $dnsmasq_container
  exit
fi

if [ "$COMMAND" = "rm" ] ; then
  running_containers="$(docker ps --filter="name=${PREFIX}" -q)"
  all_containers="$(docker ps -a --filter="name=${PREFIX}" -q)"
  volumes="$(docker volume ls | tr -s ' ' | grep ${PREFIX} | cut -d' ' -f2)"
  set +e
  if [ -n "${running_containers}" ]; then
    docker stop ${running_containers}
  fi
  if [ -n "${all_containers}" ]; then
    docker rm -f ${all_containers}
  fi
  if [ -n "${volumes}" ]; then
    docker volume rm -f ${volumes}
  fi
  exit 0
fi

if [ "$COMMAND" = "ps" ] ; then
  docker ps --filter="name=${PREFIX}"
  exit 0
fi

test -t 1 && USE_TTY="-it"
if [ "$COMMAND" = "ssh" ] ; then
  node=$1
  shift
  docker exec ${USE_TTY} ${PREFIX}${node} ssh.sh "$@"
  exit 0
fi

PORTS=""
if [ "$RANDOM_PORTS" != "true" ]; then
  if [ -n "${K8S_PORT}" ]; then PORTS="${PORTS} -p ${K8S_PORT}:6443"; fi
  if [ -n "${OSP_PORT}" ]; then PORTS="${PORTS} -p ${OSP_PORT}:8443"; fi
  if [ -n "${SSH_PORT}" ]; then PORTS="${PORTS} -p ${SSH_PORT}:2201"; fi
  if [ -n "${VNC_PORT}" ]; then PORTS="${PORTS} -p ${VNC_PORT}:5901"; fi
  if [ -n "${REGISTRY_PORT}" ]; then PORTS="${PORTS} -p ${REGISTRY_PORT}:5000"; fi
else
  PORTS="--expose 6443 --expose 8443 --expose 2201 --expose 5901 --expose 5000 -P"
fi

if [ -z "${BASE}" ]; then echo "Base image is not set. Use  '-b or --base' to set it."; exit 1; fi

if [ "$COMMAND" == "provision" ] ; then
  if [ -z "${TAG}" ]; then echo "Resulting build tag not set. Use  '-t or --tag' to set it."; exit 1; fi
  if [ -z "${SCRIPTS}" ]; then echo "Provision script is not set. Use  '-s or --script' to set it."; exit 1; fi
elif [ "$COMMAND" == "update" ] ; then
  if [ -z "${TAG}" ]; then echo "Resulting build tag not set. Use  '-t or --tag' to set it."; exit 1; fi
else
  if [ "$NODES" -lt "1" ]; then echo "The number of nodes must be greater or equal to 1."; exit 1 ; fi
fi

# Process optional qemu arguments
if [ -n "${MEMORY}" ]; then MEMORY="--memory ${MEMORY}"; fi
if [ -n "${CPU}" ]; then CPU="--cpu ${CPU}"; fi
if [ -n "${QEMU_ARGS}" ]; then QEMU_ARGS="--qemu-args ${QEMU_ARGS}"; fi

function finish() {
    set +e
    for id in ${CONTAINERS}; do
      docker stop ${id}
      docker rm -f ${id}
    done
    for vol in ${VOLUMES}; do
      docker volume rm -f ${vol}
    done
    set -e
}

if [[ "$BACKGROUND" == "true" && "$COMMAND" == "run" ]]; then
  trap finish ERR SIGINT SIGTERM SIGQUIT
else
  trap finish EXIT SIGINT SIGTERM SIGQUIT
fi

if [ "$COMMAND" = "update" ] ; then
  script=$1
  shift
  if [ -z "${script}" ]; then echo "No script to update specified"; exit 1; fi
  VM_CID=$(docker run -d --privileged --name ${PREFIX}node01 ${BASE} /bin/bash -c "sleep 100000")
  CONTAINERS="${CONTAINERS} ${VM_CID}"
  if [[ ${script} = *":"* ]]; then
    script=(${script//:/ })
    docker cp ${script[0]} ${VM_CID}:/${script[1]}
  else
    docker cp ${script} ${VM_CID}:/$(basename ${script})
    docker exec ${VM_CID} /bin/bash -c "chmod u+x /$(basename ${script})"
  fi
  docker stop ${VM_CID}
  docker commit --change "ENV SCRIPT_UPDATE $(basename ${script})" ${VM_CID} ${TAG}
  exit 0
fi

DNSMASQ_CID=$(docker run -d ${PORTS} -e NUM_NODES=${NODES} --name ${PREFIX}dnsmasq --privileged ${BASE} /bin/bash -c /dnsmasq.sh)
CONTAINERS=${DNSMASQ_CID}

if [ "$COMMAND" == "provision" ] ; then
  VM_CID=$(docker run -d --privileged --net=container:${DNSMASQ_CID} --name ${PREFIX}node01 ${BASE} /bin/bash -c "/vm.sh ${MEMORY} ${CPU} ${QEMU_ARGS}")
  CONTAINERS="${CONTAINERS} ${VM_CID}"
  docker cp ${SCRIPTS}/. ${VM_CID}:/scripts/
  docker exec ${VM_CID} /bin/bash -c "while [ ! -f /ssh_ready ] ; do sleep 1; done"
  docker exec ${VM_CID} /bin/bash -c "ssh.sh echo VM is up"
  if [ "${SIMPLE_PROVISION}" == "true" ]; then
    docker exec ${VM_CID} /bin/bash -c "scp -r -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -i vagrant.key -P 22 /scripts/manifests/* vagrant@192.168.66.101:/tmp"
  else
    docker cp ${SCRIPTS}/../../common-scripts/. ${VM_CID}:/scripts/
    docker exec ${VM_CID} /bin/bash -c "mkdir /manifests"
    docker cp ${SCRIPTS}/../../manifests/. ${VM_CID}:/manifests/
    docker exec ${VM_CID} /bin/bash -c "scp -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -i vagrant.key -P 22 /manifests/*.yaml vagrant@192.168.66.101:/tmp"
    docker exec ${VM_CID} /bin/bash -c "ssh.sh mkdir -p /tmp/ceph"
    docker exec ${VM_CID} /bin/bash -c "scp -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -i vagrant.key -P 22 /manifests/ceph/*.yaml vagrant@192.168.66.101:/tmp/ceph"
    docker exec ${VM_CID} /bin/bash -c "ssh.sh mkdir -p /tmp/cna"
    docker exec ${VM_CID} /bin/bash -c "scp -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -i vagrant.key -P 22 /manifests/cna/*.yaml vagrant@192.168.66.101:/tmp/cna"
    docker exec ${VM_CID} /bin/bash -c "ssh.sh mkdir -p /tmp/cni"
    docker exec ${VM_CID} /bin/bash -c "scp -r -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -i vagrant.key -P 22 /manifests/cni/ vagrant@192.168.66.101:/tmp/"
    docker exec ${VM_CID} /bin/bash -c "ssh.sh mkdir -p /tmp/scripts"
    docker exec ${VM_CID} /bin/bash -c "scp -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -i vagrant.key -P 22 /scripts/cnis-map.sh vagrant@192.168.66.101:/tmp/scripts/cnis-map.sh"
    docker exec ${VM_CID} /bin/bash -c "scp -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -i vagrant.key -P 22 /scripts/config-cni.sh vagrant@192.168.66.101:/tmp/scripts/config-cni.sh"
  fi

  if [[ ${PREFIX} =~ fedora ]]; then
    # Workaround for for vms that runs fedora31+ and docker,
    # In order to work with docker on Fedora31 it is necessary to revert back to cgroups v1.
    echo "Configuring cgroup version"
    docker exec ${VM_CID} /bin/bash -c "ssh.sh sudo /bin/bash -s < /scripts/cgroups.sh &"
    docker exec ${VM_CID} /bin/bash -c "while [ ! -f /ssh_ready ] ; do sleep 1; done"
    docker exec ${VM_CID} /bin/bash -c "ssh.sh echo VM is up, cgroup configured"
  fi
  docker exec ${VM_CID} /bin/bash -c "ssh.sh sudo version=${K8S_VERSION} /bin/bash -s ${CRIO} < /scripts/provision.sh"
  docker exec ${VM_CID} ssh.sh "sudo shutdown -h"
  docker exec ${VM_CID} /bin/bash -c "rm /usr/local/bin/ssh.sh"
  docker exec ${VM_CID} /bin/bash -c "rm /ssh_ready"
  docker wait ${VM_CID}
  docker commit --change "ENV PROVISIONED TRUE" ${VM_CID} ${TAG}
else
  # Start registry
  if [ -n "$REGISTRY_VOLUME" ]; then
    if [ -z "$(docker volume list | grep ${REGISTRY_VOLUME})" ]; then
      docker volume create --name ${REGISTRY_VOLUME}
    fi
    REGISTRY_VOLUME="-v ${REGISTRY_VOLUME}:/var/lib/registry"
  fi
  REGISTRY_CID=$(docker run -d --net=container:${DNSMASQ_CID} --name ${PREFIX}registry ${REGISTRY_VOLUME} registry:2)
  CONTAINERS="${CONTAINERS} ${REGISTRY_CID}"

  if [ -n "${NFS_DATA}" ]; then
    if [[ "$DIR" != /* ]]; then
      NFS_DATA=${PWD}/${NFS_DATA}
    fi
    NFS_CID=$(docker run -d --privileged --net=container:${DNSMASQ_CID} --name ${PREFIX}nfs-ganesha -v ${NFS_DATA}:/data/nfs  janeczku/nfs-ganesha:latest)
    CONTAINERS="${CONTAINERS} ${NFS_CID}"

  # Let dnsmasq learn the nfs name
  docker exec ${DNSMASQ_CID} /bin/bash -c 'echo 192.168.66.2 nfs >> /etc/hosts && kill -1 1'
  fi
  # Let dnsmasq learn the registry dns name
  docker exec ${DNSMASQ_CID} /bin/bash -c 'echo 192.168.66.2 registry >> /etc/hosts && kill -1 1'

  # Start VMs
  if [ "${REVERSE}" = "true" ]; then
   _seq="seq ${NODES} -1 1"
  else
   _seq="seq 1 ${NODES}"
  fi

  for i in $(${_seq}); do
    NODE_NUM="$(printf "%02d" ${i})"

    vol=${PREFIX}node${NODE_NUM}
    docker volume create --name ${vol}
    VOLUMES="${VOLUMES} ${vol}"

    vol="-v ${vol}:/var/run/disk/"

    VM_CID=$(docker run -d --privileged -e NODE_NUM=${NODE_NUM} ${vol} --net=container:${DNSMASQ_CID} --name ${PREFIX}node${NODE_NUM} ${BASE} /bin/bash -c "/vm.sh -n /var/run/disk/disk.qcow2 ${MEMORY} ${CPU} ${QEMU_ARGS}")
    CONTAINERS="${CONTAINERS} ${VM_CID}"
    docker exec ${VM_CID} /bin/bash -c "while [ ! -f /ssh_ready ] ; do sleep 1; done"
    if docker exec ${VM_CID} /bin/bash -c "test -f /scripts/node${NODE_NUM}.sh" ; then
      docker exec ${VM_CID} /bin/bash -c "ssh.sh sudo /bin/bash < /scripts/node${NODE_NUM}.sh"
    elif docker exec ${VM_CID} /bin/bash -c "test -f /scripts/nodes.sh" ; then
      docker exec ${VM_CID} /bin/bash -c "ssh.sh sudo /bin/bash < /scripts/nodes.sh"
    fi
    if [ "$BACKGROUND" != "true" ]; then
      docker wait ${VM_CID} &
    fi
  done
  if [ "$BACKGROUND" = "true" ]; then
    exit 0
  else
    wait
  fi
fi
