#!/usr/bin/env bash
set -euo pipefail

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
ENV_FILE="${SCRIPT_DIR}/.env"
if [[ ! -f "${ENV_FILE}" ]]; then
  printf 'Missing environment file: %s\n' "${ENV_FILE}" >&2
  exit 1
fi

set -a
# shellcheck disable=SC1090
source "${ENV_FILE}"
set +a

CUDA_BUILD="${CUDA_BUILD:-12}"
DOCKER=(sudo docker)

usage() {
  cat <<EOF
Usage: $(basename "$0") [command]

Commands:
  start           Start the llama-swap gateway in the background (default)
  daemon          Alias for start
  stop            Gracefully stop the llama-swap gateway
  restart         Stop and start the llama-swap gateway
  status          Show the Docker container state and published port
EOF
}

container_state() {
  local state
  if state="$("${DOCKER[@]}" inspect --format='{{.State.Status}}' "${LLAMA_SWAP_CONTAINER_NAME}" 2>/dev/null)"; then
    printf '%s\n' "${state}"
  else
    printf 'missing\n'
  fi
}

remove_stale_container() {
  if [[ "$(container_state)" != "missing" ]]; then
    "${DOCKER[@]}" rm --force "${LLAMA_SWAP_CONTAINER_NAME}" >/dev/null
  fi
}

start_llama_swap() {
  : "${MODELS_ROOT:?MODELS_ROOT must be set in .env}"
  : "${LLAMA_DIR:?LLAMA_DIR must be set in .env}"
  : "${STABLE_DIFFUSION_DIR:?STABLE_DIFFUSION_DIR must be set in .env}"
  : "${PRESETS_DIR:?PRESETS_DIR must be set in .env}"
  : "${MODELS_DIR:?MODELS_DIR must be set in .env}"
  : "${BUILDER_IMAGE:?BUILDER_IMAGE must be set in .env}"
  : "${LLAMA_SWAP_CONTAINER_NAME:?LLAMA_SWAP_CONTAINER_NAME must be set in .env}"
  : "${LLAMA_SWAP_PORT:?LLAMA_SWAP_PORT must be set in .env}"
  : "${LLAMA_SWAP_API_KEY:?LLAMA_SWAP_API_KEY must be set in .env}"
  : "${LLAMA_SWAP_IMAGE:?LLAMA_SWAP_IMAGE must be set in .env}"

  if [[ ! -f "${PRESETS_DIR}/llama-swap.yaml" ]]; then
    printf 'Missing llama-swap configuration: %s\n' "${PRESETS_DIR}/llama-swap.yaml" >&2
    return 1
  fi

  if [[ ! -x "${LLAMA_DIR}/build-${CUDA_BUILD}/bin/llama-server" ]]; then
    printf 'Missing custom llama-server: %s\n' "${LLAMA_DIR}/build-${CUDA_BUILD}/bin/llama-server" >&2
    printf 'Build llama.cpp before starting llama-swap.\n' >&2
    return 1
  fi

  if [[ ! -x "${STABLE_DIFFUSION_DIR}/build-${CUDA_BUILD}/bin/sd-server" ]]; then
    printf 'Missing custom sd-server: %s\n' \
      "${STABLE_DIFFUSION_DIR}/build-${CUDA_BUILD}/bin/sd-server" >&2
    printf 'Build stable-diffusion.cpp before starting llama-swap.\n' >&2
    return 1
  fi

  if [[ "$(container_state)" == "running" ]]; then
    printf '%s is already running\n' "${LLAMA_SWAP_CONTAINER_NAME}" >&2
    return 1
  fi

  mkdir -p "${PRESETS_DIR}/llama-swap-logs" "${PRESETS_DIR}/llama-swap-state"
  remove_stale_container

  "${DOCKER[@]}" run --detach --rm \
    --name "${LLAMA_SWAP_CONTAINER_NAME}" \
    --gpus all \
    --device=/dev/dri \
    --ulimit memlock=-1:-1 \
    --publish "${LLAMA_SWAP_PORT}:8080" \
    --env "CUDA_BUILD=${CUDA_BUILD}" \
    --env "HF_HOME=/models" \
    --env "HF_TOKEN=${HF_TOKEN:-}" \
    --env "LLAMA_SWAP_API_KEY=${LLAMA_SWAP_API_KEY}" \
    --env "LLAMA_SWAP_CONFIG=/presets/llama-swap.yaml" \
    --volume "${LLAMA_DIR}:/src:ro" \
    --volume "${STABLE_DIFFUSION_DIR}:/stable-diffusion:ro" \
    --volume "${PRESETS_DIR}:/presets" \
    --volume "${MODELS_DIR}:/models" \
    --workdir /models \
    "${LLAMA_SWAP_IMAGE}"
  printf 'Started %s in the background\n' "${LLAMA_SWAP_CONTAINER_NAME}"
}

stop_llama_swap() {
  local state timeout
  state="$(container_state)"
  timeout="${STOP_TIMEOUT:-30}"

  if [[ "${state}" == "missing" ]]; then
    printf '%s is not running\n' "${LLAMA_SWAP_CONTAINER_NAME}"
    return 0
  fi

  if [[ "${state}" == "running" || "${state}" == "created" || "${state}" == "restarting" ]]; then
    "${DOCKER[@]}" stop --time="${timeout}" "${LLAMA_SWAP_CONTAINER_NAME}" >/dev/null
  else
    remove_stale_container
  fi
  printf 'Stopped %s\n' "${LLAMA_SWAP_CONTAINER_NAME}"
}

status_llama_swap() {
  local state port
  state="$(container_state)"
  if [[ "${state}" == "missing" ]]; then
    printf '%s: stopped (container missing)\n' "${LLAMA_SWAP_CONTAINER_NAME}"
    return 1
  fi

  port="$("${DOCKER[@]}" port "${LLAMA_SWAP_CONTAINER_NAME}" 8080/tcp 2>/dev/null || true)"
  if [[ -n "${port}" ]]; then
    printf '%s: %s (%s)\n' "${LLAMA_SWAP_CONTAINER_NAME}" "${state}" "${port}"
  else
    printf '%s: %s\n' "${LLAMA_SWAP_CONTAINER_NAME}" "${state}"
  fi
  [[ "${state}" == "running" ]]
}

if [[ $# -gt 0 ]]; then
  COMMAND="$1"
else
  COMMAND="start"
fi
case "${COMMAND}" in
  start|daemon|restart)
    if [[ $# -gt 1 ]]; then
      printf 'Unexpected argument for %s: %s\n' "${COMMAND}" "$2" >&2
      usage >&2
      exit 2
    fi
    ;;
  stop|status)
    if [[ $# -gt 1 ]]; then
      printf 'Unexpected argument for %s: %s\n' "${COMMAND}" "$2" >&2
      usage >&2
      exit 2
    fi
    ;;
  help|-h|--help)
    usage
    exit 0
    ;;
  *)
    printf 'Unknown command: %s\n' "${COMMAND}" >&2
    usage >&2
    exit 2
    ;;
esac

: "${LLAMA_SWAP_CONTAINER_NAME:?LLAMA_SWAP_CONTAINER_NAME must be set in .env}"
case "${COMMAND}" in
  start|daemon)
    start_llama_swap
    ;;
  stop)
    stop_llama_swap
    ;;
  restart)
    stop_llama_swap
    start_llama_swap
    ;;
  status)
    status_llama_swap
    ;;
esac