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

# LibreTranslate full setup script (Docker-based)
# Features:
#  - Installs Docker if missing (optional --no-docker-install)
#  - Pulls libretranslate image (tag configurable)
#  - Creates persistent data + cache directories
#  - Optionally pre-downloads language models
#  - Generates or accepts an API key; can disable auth
#  - (Removed) systemd service setup – now always ephemeral
#  - Health check + sample translation
#  - Uninstall mode removes container, image, service, and data (optional keep data)
#  - Idempotent: safe to re-run for upgrades (will pull newer image)

SCRIPT_NAME=$(basename "$0")
VERSION="1.0.0"

# Defaults
IMAGE="libretranslate/libretranslate"
TAG="latest"
SERVICE_NAME="libretranslate"
DOCKER_INSTALL=1
# Systemd removed – always run ephemeral container
API_KEY=""
GENERATE_API_KEY=1
DISABLE_API_KEY=0
PORT=5000
HOST=0.0.0.0
DATA_DIR="/var/lib/libretranslate"
CACHE_DIR="${DATA_DIR}/cache"
CONFIG_DIR="/etc/libretranslate"
ENV_FILE="${CONFIG_DIR}/libretranslate.env"
PULL_ONLY=0
PRELOAD_LANGS=""
UNINSTALL=0
KEEP_DATA=0
HEALTH_TIMEOUT=15
EXTRA_ENV=()
NO_COLOR=0
KEEP_ALIVE=0
RUN_COMMAND=()
DEBUG=0

# Colors
if [[ -t 1 && ${NO_COLOR} -eq 0 ]]; then
  GREEN="\e[32m"
  YELLOW="\e[33m"
  RED="\e[31m"
  BLUE="\e[34m"
  BOLD="\e[1m"
  RESET="\e[0m"
else
  GREEN=""
  YELLOW=""
  RED=""
  BLUE=""
  BOLD=""
  RESET=""
fi

log() { echo -e "${BLUE}[INFO]${RESET} $*"; }
warn() { echo -e "${YELLOW}[WARN]${RESET} $*" >&2; }
err() { echo -e "${RED}[ERR ]${RESET} $*" >&2; }
success() { echo -e "${GREEN}[OK  ]${RESET} $*"; }

usage() {
  cat << EOF
${SCRIPT_NAME} v${VERSION}
Setup or uninstall a self-hosted LibreTranslate instance via Docker.

Usage: ${SCRIPT_NAME} [options]

Primary actions:
	--uninstall              Remove service, container, image (data kept unless --purge or --no-keep-data)
	--pull-only              Only pull/update image & exit

Install behavior options:
	--image NAME             Docker image (default: ${IMAGE})
	--tag TAG                Docker tag (default: ${TAG})
	--port N                 Host port to expose (default: ${PORT})
	--host IP                Bind host (default: ${HOST})
	--data-dir PATH          Persistent data directory (default: ${DATA_DIR})
	--cache-dir PATH         Models cache dir (default: ${CACHE_DIR})
	--no-docker-install      Do not attempt to install Docker
	(systemd support removed; container is ephemeral)
	--keep-alive             Keep running (tail logs) until Ctrl-C
	--                       Treat remaining arguments as a command to run after service is healthy; service stops when command exits
	--api-key KEY            Use specified API key
	--generate-api-key       Force generate new random key (default if none provided)
	--disable-api-key        Disable key requirement (open instance)
	--preload-langs CSV      Pre-download language models (e.g. en,es,fr,de)
	--env K=V                Extra environment variable (repeatable)
	--health-timeout SEC     Wait time for health check (default: ${HEALTH_TIMEOUT})
	--debug                  Verbose output (do not suppress curl errors; follow logs on failure)
	--no-color               Disable colored output

Uninstall options:
	--purge                  Remove data directory (implies --uninstall)
	--keep-data              Keep data on uninstall (default)

Misc:
	-h, --help               Show this help
	-v, --version            Show version

Examples:
	${SCRIPT_NAME} --preload-langs en,es,fr --env LT_LOAD_ONLY=en,es,fr
	${SCRIPT_NAME} --api-key mysecret123 --port 8080
	${SCRIPT_NAME} --uninstall --purge
EOF
}

gen_api_key() {
  # Avoid SIGPIPE issues under set -o pipefail by capturing output first
  local key
  key=$(head -c 256 /dev/urandom | tr -dc 'A-Za-z0-9' | head -c 40 || true)
  if [[ -z $key || ${#key} -lt 40 ]]; then
    # Fallback using openssl if available
    if command -v openssl > /dev/null 2>&1; then
      key=$(openssl rand -base64 48 | tr -dc 'A-Za-z0-9' | head -c 40 || true)
    fi
  fi
  if [[ -z $key || ${#key} -lt 20 ]]; then
    # Last resort static warning key (should not happen)
    key="LT$(date +%s)$$RANDOM"
  fi
  printf '%s' "$key"
}

need_cmd() {
  command -v "$1" > /dev/null 2>&1 || {
    err "Required command '$1' not found"
    return 1
  }
}

parse_args() {
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --image)
        IMAGE="$2"
        shift 2
        ;;
      --tag)
        TAG="$2"
        shift 2
        ;;
      --port)
        PORT="$2"
        shift 2
        ;;
      --host)
        HOST="$2"
        shift 2
        ;;
      --data-dir)
        DATA_DIR="$2"
        CACHE_DIR="${DATA_DIR}/cache"
        shift 2
        ;;
      --cache-dir)
        CACHE_DIR="$2"
        shift 2
        ;;
      --no-docker-install)
        DOCKER_INSTALL=0
        shift
        ;;
      --keep-alive)
        KEEP_ALIVE=1
        shift
        ;;
      --)
        shift
        RUN_COMMAND=("$@")
        break
        ;;
      --api-key)
        API_KEY="$2"
        GENERATE_API_KEY=0
        shift 2
        ;;
      --generate-api-key)
        GENERATE_API_KEY=1
        shift
        ;;
      --disable-api-key)
        DISABLE_API_KEY=1
        shift
        ;;
      --preload-langs)
        PRELOAD_LANGS="$2"
        shift 2
        ;;
      --env)
        EXTRA_ENV+=("$2")
        shift 2
        ;;
      --pull-only)
        PULL_ONLY=1
        shift
        ;;
      --uninstall)
        UNINSTALL=1
        shift
        ;;
      --purge)
        UNINSTALL=1
        KEEP_DATA=0
        shift
        ;;
      --keep-data)
        KEEP_DATA=1
        shift
        ;;
      --health-timeout)
        HEALTH_TIMEOUT="$2"
        shift 2
        ;;
      --no-color)
        NO_COLOR=1
        shift
        ;;
      --debug)
        DEBUG=1
        shift
        ;;
      -h | --help)
        usage
        exit 0
        ;;
      -v | --version)
        echo "${VERSION}"
        exit 0
        ;;
      *)
        err "Unknown argument: $1"
        usage
        exit 1
        ;;
    esac
  done
}

ensure_root() {
  if [[ $EUID -ne 0 ]]; then
    err "This script must run as root (or via sudo)."
    exit 1
  fi
}

install_docker() {
  if command -v docker > /dev/null 2>&1; then
    log "Docker already installed"
    return 0
  fi
  if [[ ${DOCKER_INSTALL} -eq 0 ]]; then
    err "Docker is not installed and --no-docker-install specified."
    exit 1
  fi
  log "Installing Docker..."
  if command -v apt-get > /dev/null 2>&1; then
    apt-get update -y
    apt-get install -y ca-certificates curl gnupg
    install -d -m 0755 /etc/apt/keyrings
    curl -fsSL "https://download.docker.com/linux/$(
      . /etc/os-release
      echo "$ID"
    )/gpg" | gpg --dearmor -o /etc/apt/keyrings/docker.gpg
    chmod a+r /etc/apt/keyrings/docker.gpg
    echo \
      "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/$(
        . /etc/os-release
        echo "$ID"
      ) $(
        . /etc/os-release
        echo "$VERSION_CODENAME"
      ) stable" \
      > /etc/apt/sources.list.d/docker.list
    apt-get update -y
    apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
  else
    err "Unsupported package manager. Please install Docker manually."
    exit 1
  fi
  # Attempt to start docker daemon if dockerd exists and systemctl available; otherwise rely on user
  if command -v systemctl > /dev/null 2>&1; then
    (systemctl enable --now docker 2> /dev/null && success "Docker installed and started") || warn "Docker installed; ensure dockerd is running"
  else
    warn "Docker installed; please ensure docker daemon is running"
  fi
}

pull_image() {
  log "Pulling image ${IMAGE}:${TAG}"
  docker pull "${IMAGE}:${TAG}"
  success "Image pulled"
}

detect_container_user() {
  # Determine uid/gid of configured user inside image so host dirs can be chowned
  if ! command -v docker > /dev/null 2>&1; then
    return 0
  fi
  local uid gid
  uid=$(docker run --rm --entrypoint /usr/bin/id "${IMAGE}:${TAG}" -u 2> /dev/null || echo "")
  gid=$(docker run --rm --entrypoint /usr/bin/id "${IMAGE}:${TAG}" -g 2> /dev/null || echo "")
  if [[ -n $uid && -n $gid ]]; then
    CONTAINER_UID=$uid
    CONTAINER_GID=$gid
  fi
}

write_env_file() {
  mkdir -p "${CONFIG_DIR}" "${DATA_DIR}" "${CACHE_DIR}"
  detect_container_user
  if [[ -n ${CONTAINER_UID:-} && -n ${CONTAINER_GID:-} ]]; then
    if command -v stat > /dev/null 2>&1; then
      for d in "${DATA_DIR}" "${CACHE_DIR}"; do
        if [[ -d $d ]]; then
          CUR_UID=$(stat -c %u "$d" 2> /dev/null || echo -1)
          if [[ ${CUR_UID} -ne ${CONTAINER_UID} ]]; then
            chown "${CONTAINER_UID}":"${CONTAINER_GID}" "$d" 2> /dev/null || warn "Unable to chown $d to ${CONTAINER_UID}:${CONTAINER_GID}"
          fi
        fi
      done
    fi
  fi
  if [[ ${DISABLE_API_KEY} -eq 1 ]]; then
    API_KEY_LINE="LT_NO_API_KEY=true"
  else
    if [[ -z ${API_KEY} && ${GENERATE_API_KEY} -eq 1 ]]; then
      API_KEY=$(gen_api_key)
      GENERATED=1
    else
      GENERATED=0
    fi
    API_KEY_LINE="LT_API_KEYS=${API_KEY}"
  fi

  {
    echo "# LibreTranslate environment file"
    echo "# Generated $(date -u +%Y-%m-%dT%H:%M:%SZ)"
    echo "${API_KEY_LINE}"
    [[ -n ${PRELOAD_LANGS} ]] && echo "LT_PRELOAD_LANGS=${PRELOAD_LANGS}"
    for kv in "${EXTRA_ENV[@]:-}"; do echo "$kv"; done
  } > "${ENV_FILE}.tmp"
  mv "${ENV_FILE}.tmp" "${ENV_FILE}"
  chmod 600 "${ENV_FILE}"
  success "Environment file written: ${ENV_FILE}"
}

start_container_ephemeral() {
  docker rm -f "${SERVICE_NAME}" > /dev/null 2>&1 || true
  docker run -d --name "${SERVICE_NAME}" \
    --env-file "${ENV_FILE}" \
    -v "${DATA_DIR}:/home/libretranslate/.local/share/argos-translate" \
    -v "${CACHE_DIR}:/app/cache" \
    -p "${PORT}:${PORT}" \
    "${IMAGE}:${TAG}" \
    --host 0.0.0.0 --port "${PORT}"
  success "Container started (ephemeral)"
  echo
  echo "Endpoint (pending readiness): http://$(hostname -I | awk '{print $1}'):${PORT}"
  echo "Waiting for health..."
}

health_check() {
  local start
  start=$(date +%s)
  local url="http://127.0.0.1:${PORT}/languages"
  local attempt=0
  while true; do
    attempt=$((attempt + 1))
    if curl ${DEBUG:+-v} -fsS "$url" > /dev/null 2>&1; then
      success "Service healthy (attempt $attempt)"
      return 0
    else
      [[ $DEBUG -eq 1 ]] && log "Health attempt $attempt failed"
    fi
    if (($(date +%s) - start > HEALTH_TIMEOUT)); then
      err "Health check failed after ${HEALTH_TIMEOUT}s (attempts: $attempt)"
      docker logs --tail 200 "${SERVICE_NAME}" || true
      return 1
    fi
    sleep 0.5
  done
}

sample_request() {
  if [[ ${DISABLE_API_KEY} -eq 0 ]]; then
    local key="${API_KEY}"
  else
    local key=""
  fi
  log "Performing sample translation (en->es)..."
  local DATA='{"q":"Hello world","source":"en","target":"es","format":"text"}'
  if [[ -n $key ]]; then
    curl -fsS -H "Content-Type: application/json" -H "Authorization: ${key}" -d "$DATA" "http://127.0.0.1:${PORT}/translate" || warn "Sample request failed"
  else
    curl -fsS -H "Content-Type: application/json" -d "$DATA" "http://127.0.0.1:${PORT}/translate" || warn "Sample request failed"
  fi
  echo
}

uninstall_all() {
  log "Uninstalling LibreTranslate (ephemeral mode)..."
  docker rm -f "${SERVICE_NAME}" 2> /dev/null || true
  docker rmi "${IMAGE}:${TAG}" 2> /dev/null || true
  if [[ ${KEEP_DATA} -eq 0 ]]; then
    rm -rf "${DATA_DIR}" "${CONFIG_DIR}" || true
    success "Data directories removed"
  else
    log "Data kept in ${DATA_DIR} and ${CONFIG_DIR}"
  fi
  success "Uninstall complete"
  exit 0
}

main() {
  parse_args "$@"
  ensure_root

  if [[ ${UNINSTALL} -eq 1 ]]; then
    uninstall_all
  fi

  install_docker
  pull_image
  if [[ ${PULL_ONLY} -eq 1 ]]; then
    log "Pull-only requested, exiting."
    exit 0
  fi

  write_env_file

  # Always ephemeral now
  start_container_ephemeral

  health_check
  sample_request || true

  # If a command is provided, run it and then shutdown container
  if [[ ${#RUN_COMMAND[@]} -gt 0 ]]; then
    log "Running user command: ${RUN_COMMAND[*]}"
    set +e
    "${RUN_COMMAND[@]}"
    CMD_STATUS=$?
    set -e
    log "Command exited with status ${CMD_STATUS}; stopping container"
    docker stop "${SERVICE_NAME}" > /dev/null 2>&1 || true
    exit ${CMD_STATUS}
  fi

  if [[ ${KEEP_ALIVE} -eq 1 ]]; then
    log "Tailing logs (Ctrl-C to stop and remove container)"
    trap 'log "Stopping container"; docker stop "${SERVICE_NAME}" >/dev/null 2>&1 || true; exit 0' INT TERM
    docker logs -f "${SERVICE_NAME}"
    log "Logs ended; stopping container"
    docker stop "${SERVICE_NAME}" > /dev/null 2>&1 || true
  else
    log "Ephemeral container left running in background (id: $(docker inspect --format '{{.Id}}' ${SERVICE_NAME} 2> /dev/null || echo unknown))"
    log "Stop manually with: docker stop ${SERVICE_NAME}"
  fi

  echo
  echo "${BOLD}LibreTranslate is ready.${RESET}"
  echo "Endpoint: http://$(hostname -I | awk '{print $1}'):${PORT}"
  if [[ ${DISABLE_API_KEY} -eq 0 ]]; then
    if [[ ${GENERATED:-0} -eq 1 ]]; then
      echo "Generated API key: ${API_KEY}"
    else
      echo "API key: ${API_KEY}"
    fi
    echo "Use header: Authorization: <API_KEY>"
  else
    echo "API key authentication DISABLED (public instance)."
  fi
  if [[ -n ${PRELOAD_LANGS} ]]; then
    echo "Preloaded languages requested: ${PRELOAD_LANGS}"
  fi
  echo "Environment file: ${ENV_FILE}"
  echo "Manage: docker logs -f ${SERVICE_NAME} | docker stop ${SERVICE_NAME}"
  echo "Uninstall: sudo ${SCRIPT_NAME} --uninstall"
  echo
}

main "$@"
