#!/bin/sh
set -eu
OMNI_DOMAIN="omnigto.dev"
IMAGE_TAG="develop"
INSTALL_DIR="/opt/omni"
err() { echo "Error: $*" >&2; exit 1; }
ask() {
    printf '%s' "$1" > /dev/tty
    IFS= read -r reply < /dev/tty || err "could not read a reply."
    reply=$(printf '%s' "$reply" | tr -d '[:cntrl:][:space:]' | tr '[:upper:]' '[:lower:]')
}
[ "$(id -u)" -eq 0 ] || err "run as root, e.g.  curl -fsSL https://get.omnigto.dev | sudo sh"
[ -r /dev/tty ] || err "no terminal to prompt on; download the script and run it directly instead of piping."
arch=$(uname -m)
case "$arch" in
    x86_64 | amd64) : ;;
    *) err "Omni needs an x86_64 (Intel/AMD) server; yours is $arch (e.g. ARM). It cannot run here." ;;
esac
MODE=install
if [ -e "$INSTALL_DIR/.env" ]; then
    printf 'Found an existing Omni core at %s.\n  [U] Update      — get the latest Omni, keep your current token (default).\n                    Restarts the core. A running solve is gracefully saved and reloaded if possible\n  [R] Reconfigure — enter a new account token or core name\n  [X] Uninstall   — stop and remove this core\n  [Q] Quit\n' "$INSTALL_DIR" > /dev/tty
    while :; do
        ask 'Choose [U/r/x/q]: '
        case "$reply" in
            '' | u | update) MODE=update ; break ;;
            r | reconfigure) MODE=reconfigure ; break ;;
            x | uninstall) MODE=uninstall ; break ;;
            q | quit) echo "Nothing changed. The existing core is untouched."; exit 0 ;;
            *) printf 'Sorry, "%s" is not one of U, R, X or Q.\n' "$reply" > /dev/tty ;;
        esac
    done
fi
if [ "$MODE" = uninstall ]; then
    command -v docker >/dev/null 2>&1 || err "Docker is gone, so nothing can be removed cleanly. Delete $INSTALL_DIR by hand."
    [ -f "$INSTALL_DIR/compose.yaml" ] || \
        err "$INSTALL_DIR/compose.yaml is missing, so the core cannot be stopped from here. Re-run and choose Update to restore it, then uninstall."
    cd "$INSTALL_DIR"
    core_image=$(docker compose images -q core 2>/dev/null || true)
    echo "Stopping the core (waiting for it to checkpoint a running solve, this can take a few minutes)..."
    docker compose down --remove-orphans
    if docker volume inspect omni >/dev/null 2>&1; then
        mount=$(docker volume inspect omni --format '{{.Mountpoint}}' 2>/dev/null || true)
        size=$([ -n "$mount" ] && [ -d "$mount" ] && du -sh "$mount" 2>/dev/null | cut -f1 || true)
        prompt=$(printf 'Also delete the solver data volume "omni"%s? Saved solutions, solve checkpoints and the abstraction cache. This cannot be undone. [y/N] ' \
                        "${size:+ ($size)}")
        while :; do
            ask "$prompt"
            case "$reply" in
                y | yes) if docker volume rm omni >/dev/null 2>&1; then echo "Deleted the data volume."
                         else echo "Could not delete the data volume 'omni' (something else may still use it)."; fi
                         break ;;
                '' | n | no) KEPT_VOLUME=1 ; break ;;
                *) printf 'Sorry, "%s" is not one of Y or N.\n' "$reply" > /dev/tty ;;
            esac
        done
    fi
    [ -n "$core_image" ] && docker image rm "$core_image" >/dev/null 2>&1 || true
    docker logout "registry.$OMNI_DOMAIN" >/dev/null 2>&1 || true
    cd /
    rm -rf "$INSTALL_DIR"
    rm -f /etc/modules-load.d/omni-bbr.conf
    echo
    echo "Done. The Omni core is uninstalled and will stop appearing in your panel."
    if [ "${KEPT_VOLUME:-}" = 1 ]; then
        echo "Kept your solver data in the Docker volume 'omni'. Reinstalling picks it up again;"
        echo "to delete it later:  docker volume rm omni"
    fi
    echo "Reinstall any time with the same command from https://app.$OMNI_DOMAIN"
    exit 0
fi
if ! command -v docker >/dev/null 2>&1; then
    while :; do
        ask 'Docker is not installed. Install it now via get.docker.com? [Y/n] '
        case "$reply" in
            '' | y | yes) echo "Installing Docker via get.docker.com..."
                          curl -fsSL https://get.docker.com | sh || err "Docker install failed; install it manually and re-run."
                          break ;;
            n | no) err "Docker is required. Install it (https://docs.docker.com/engine/install/) and re-run." ;;
            *) printf 'Sorry, "%s" is not one of Y or N.\n' "$reply" > /dev/tty ;;
        esac
    done
fi
docker compose version >/dev/null 2>&1 || \
    err "Docker Compose v2 plugin is missing; install it (https://docs.docker.com/compose/install/) and re-run."
if [ "$MODE" = update ]; then
    ACCOUNT_TOKEN=$(sed -n 's/^OMNI_ACCOUNT_TOKEN=//p' "$INSTALL_DIR/.env")
    [ -n "$ACCOUNT_TOKEN" ] || err "no OMNI_ACCOUNT_TOKEN in $INSTALL_DIR/.env; re-run and choose Reconfigure."
else
    printf 'Enter your Omni account token (omni_..., shown on your account page): ' > /dev/tty
    IFS= read -r ACCOUNT_TOKEN < /dev/tty || err "could not read the account token."
    ACCOUNT_TOKEN=$(printf '%s' "$ACCOUNT_TOKEN" | tr -d '[:space:]')
    [ -n "$ACCOUNT_TOKEN" ] || err "account token must not be empty."
    case "$ACCOUNT_TOKEN" in
      omni_?*) : ;;
      *) err "that does not look like an Omni account token (expected omni_...): $ACCOUNT_TOKEN" ;;
    esac
    default_core_id=$(hostname 2>/dev/null || uname -n)
    printf 'Name for this core [%s]: ' "$default_core_id" > /dev/tty
    IFS= read -r CORE_ID < /dev/tty || err "could not read a reply."
    CORE_ID=$(printf '%s' "$CORE_ID" | tr -d '[:space:]')
    [ -n "$CORE_ID" ] || CORE_ID=$default_core_id
fi
echo "Logging in to the image registry..."
printf '%s' "$ACCOUNT_TOKEN" | docker login "registry.$OMNI_DOMAIN" -u omni --password-stdin \
    || err "registry login failed — check the account token and that your license is active."
mkdir -p "$INSTALL_DIR"
cat > "$INSTALL_DIR/compose.yaml" <<'COMPOSE'
name: omni-core

services:
  core:
    image: registry.${OMNI_DOMAIN}/omnigto/omni-core:${TAG:-latest}
    environment:
      OMNI_RELAY_URL: "wss://app.${OMNI_DOMAIN}/core"
      OMNI_ACCOUNT_TOKEN: "${OMNI_ACCOUNT_TOKEN:?OMNI_ACCOUNT_TOKEN is missing}"
      OMNI_CORE_ID: "${OMNI_CORE_ID:-}"
    volumes:
      - omni:/omni
    sysctls:
      net.ipv4.tcp_congestion_control: "${OMNI_TCP_CC:-cubic}"
      net.ipv4.tcp_slow_start_after_idle: "0"
      net.ipv4.tcp_wmem: "4096 16384 33554432"
    restart: unless-stopped
    stop_grace_period: 5m

volumes:
  omni:
    name: omni
COMPOSE
if [ "$MODE" != update ]; then
    cat > "$INSTALL_DIR/.env" <<ENV
OMNI_DOMAIN=$OMNI_DOMAIN
TAG=$IMAGE_TAG
OMNI_ACCOUNT_TOKEN=$ACCOUNT_TOKEN
OMNI_CORE_ID=$CORE_ID
ENV
    chmod 600 "$INSTALL_DIR/.env"
fi
TCP_CC=cubic
if modprobe tcp_bbr >/dev/null 2>&1 &&
    grep -qw bbr /proc/sys/net/ipv4/tcp_available_congestion_control 2>/dev/null; then
    TCP_CC=bbr
    echo tcp_bbr > /etc/modules-load.d/omni-bbr.conf 2>/dev/null || true
fi
sed -i '/^OMNI_TCP_CC=/d' "$INSTALL_DIR/.env"
echo "OMNI_TCP_CC=$TCP_CC" >> "$INSTALL_DIR/.env"
cd "$INSTALL_DIR"
echo "Pulling the core image..."
docker compose pull
echo "Starting the core..."
old_core_image=$(docker compose images -q core 2>/dev/null || true)
old_container=$(docker compose ps -q core 2>/dev/null || true)
docker compose up -d
new_core_image=$(docker compose images -q core 2>/dev/null || true)
new_container=$(docker compose ps -q core 2>/dev/null || true)
if [ -n "$old_core_image" ] && [ "$old_core_image" != "$new_core_image" ]; then
    docker image rm "$old_core_image" >/dev/null 2>&1 || true
fi
echo
if [ "$MODE" = update ]; then
    if [ -n "$old_container" ] && [ "$old_container" = "$new_container" ]; then
        echo "Already the latest Omni. Nothing changed and the core was not restarted."
    else
        echo "Done. The core is updated and running."
    fi
else
    echo "Done. The core '$CORE_ID' is running."
fi
echo "Logs:      cd $INSTALL_DIR && docker compose logs -f"
echo "Update:    re-run the same install command and choose Update."
echo "Uninstall: re-run the same install command and choose Uninstall."
echo "Manage this core at https://app.$OMNI_DOMAIN"