#!/usr/bin/env sh
set -eu

REPO_URL="${ARMORER_REPO_URL:-https://github.com/ArmorerLabs/Armorer.git}"
BRANCH="${ARMORER_BRANCH:-main}"
INSTALL_DIR="${ARMORER_INSTALL_DIR:-$HOME/.armorer/source}"
VERSION="${ARMORER_VERSION:-stable}"
VERSION_EXPLICIT="0"
[ "${ARMORER_VERSION+x}" = "x" ] && VERSION_EXPLICIT="1"
IMAGE="${ARMORER_IMAGE:-}"
BIN_DIR="${ARMORER_BIN_DIR:-$HOME/.local/bin}"
YES="0"
OPEN_UI="1"
START_UI="1"
DEV_SOURCE="0"
PORT="${ARMORER_UI_PORT:-3000}"
NO_PULL="${ARMORER_NO_PULL:-0}"

log() {
  printf '%s\n' "$*"
}

fail() {
  printf 'Armorer install failed: %s\n' "$*" >&2
  exit 1
}

usage() {
  cat <<'USAGE'
Install Armorer as a versioned local appliance.

Usage:
  curl -fsSL https://armorerlabs.com/install | sh
  curl -fsSL https://armorerlabs.com/install | sh -s -- --yes

Options:
  --yes          Use safe defaults and avoid prompts.
  --version V    Launcher/image version. Defaults to stable. Use beta for the beta channel or vX.Y.Z for pinned releases.
  --image IMAGE  Container image override. Defaults to ghcr.io/armorerlabs/armorer/appliance:<version>.
  --bin-dir DIR  Install the native launcher here. Defaults to ~/.local/bin.
  --port PORT    Local UI port. Defaults to 3000.
  --no-start     Install launcher only; do not start Armorer.
  --no-pull      Start from an existing local image without pulling.
  --no-open      Do not open the browser after setup.
  --dev          Development install from source with Node/pnpm build.
  --repo URL     Dev source repository.
  --branch NAME  Dev source branch.
  --dir PATH     Dev source checkout directory.

Environment:
  ARMORER_DOCKER_SOCKET=/path/to/docker.sock overrides automatic Docker socket detection.
USAGE
}

while [ "$#" -gt 0 ]; do
  case "$1" in
    --yes|-y)
      YES="1"
      ;;
    --version)
      shift
      [ "$#" -gt 0 ] || fail "--version needs a value"
      VERSION="$1"
      VERSION_EXPLICIT="1"
      ;;
    --image)
      shift
      [ "$#" -gt 0 ] || fail "--image needs a value"
      IMAGE="$1"
      ;;
    --bin-dir)
      shift
      [ "$#" -gt 0 ] || fail "--bin-dir needs a path"
      BIN_DIR="$1"
      ;;
    --port)
      shift
      [ "$#" -gt 0 ] || fail "--port needs a value"
      PORT="$1"
      ;;
    --no-start)
      START_UI="0"
      ;;
    --no-pull)
      NO_PULL="1"
      ;;
    --no-open)
      OPEN_UI="0"
      ;;
    --dev)
      DEV_SOURCE="1"
      ;;
    --repo)
      shift
      [ "$#" -gt 0 ] || fail "--repo needs a URL"
      REPO_URL="$1"
      ;;
    --branch)
      shift
      [ "$#" -gt 0 ] || fail "--branch needs a name"
      BRANCH="$1"
      ;;
    --dir)
      shift
      [ "$#" -gt 0 ] || fail "--dir needs a path"
      INSTALL_DIR="$1"
      ;;
    --help|-h)
      usage
      exit 0
      ;;
    *)
      fail "unknown option: $1"
      ;;
  esac
  shift
done

if [ "$VERSION_EXPLICIT" = "0" ] && [ -n "$IMAGE" ]; then
  case "$IMAGE" in
    *:beta|*:experimental|*:experimental-*)
      VERSION="beta"
      log "Inferred launcher channel beta from ARMORER_IMAGE=$IMAGE"
      ;;
  esac
fi

case "$(uname -s)" in
  Darwin)
    OS="darwin"
    ;;
  Linux)
    OS="linux"
    ;;
  *)
    fail "unsupported OS: $(uname -s). Windows support is planned after macOS/Linux."
    ;;
esac

case "$(uname -m)" in
  x86_64|amd64)
    ARCH="amd64"
    ;;
  arm64|aarch64)
    ARCH="arm64"
    ;;
  *)
    fail "unsupported CPU architecture: $(uname -m)"
    ;;
esac

if [ "$OS" = "darwin" ]; then
  PATH="$BIN_DIR:/opt/homebrew/bin:/usr/local/bin:$PATH"
else
  PATH="$BIN_DIR:$HOME/.local/bin:$PATH"
fi
export PATH

confirm() {
  prompt="$1"
  if [ "$YES" = "1" ]; then
    return 0
  fi
  printf '%s [y/N] ' "$prompt"
  read ans || ans=""
  case "$ans" in
    y|Y|yes|YES)
      return 0
      ;;
    *)
      return 1
      ;;
  esac
}

have() {
  command -v "$1" >/dev/null 2>&1
}

detect_user_home() {
  user="$(id -un 2>/dev/null || printf '')"
  detected=""
  if [ -n "$user" ] && [ "$OS" = "darwin" ] && have dscl; then
    detected="$(dscl . -read "/Users/$user" NFSHomeDirectory 2>/dev/null | awk '{print $2}' || true)"
  elif [ -n "$user" ] && have getent; then
    detected="$(getent passwd "$user" 2>/dev/null | cut -d: -f6 || true)"
  fi
  if [ -n "$detected" ]; then
    printf '%s\n' "$detected"
  else
    printf '%s\n' "$HOME"
  fi
}

HOST_HOME="${ARMORER_HOST_HOME:-$(detect_user_home)}"

ensure_homebrew() {
  if have brew; then
    return 0
  fi
  [ "$OS" = "darwin" ] || return 1
  confirm "Homebrew is required to install the open-source container runtime. Install Homebrew now?" || fail "Homebrew is required on macOS when dependencies are missing."
  /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"
  PATH="$BIN_DIR:/opt/homebrew/bin:/usr/local/bin:$PATH"
  export PATH
  have brew || fail "Homebrew install did not put brew on PATH"
}

docker_works() {
  docker version >/dev/null 2>&1 && HOME="$HOST_HOME" docker compose version >/dev/null 2>&1
}

docker_with_socket_works() {
  sock="$1"
  [ -S "$sock" ] || return 1
  DOCKER_HOST="unix://$sock" DOCKER_CONTEXT= docker version >/dev/null 2>&1 &&
    HOME="$HOST_HOME" DOCKER_HOST="unix://$sock" DOCKER_CONTEXT= docker compose version >/dev/null 2>&1
}

try_explicit_runtime() {
  [ -n "${ARMORER_DOCKER_SOCKET:-}" ] || return 1
  if docker_with_socket_works "$ARMORER_DOCKER_SOCKET"; then
    log "Using Armorer Docker socket override: $ARMORER_DOCKER_SOCKET"
    export DOCKER_HOST="unix://$ARMORER_DOCKER_SOCKET"
    export DOCKER_CONTEXT=
    return 0
  fi
  fail "ARMORER_DOCKER_SOCKET is set but Docker is unavailable at $ARMORER_DOCKER_SOCKET"
}

try_existing_macos_runtime() {
  [ "$OS" = "darwin" ] || return 1
  seen_home=""
  for runtime_home in "$HOME" "$HOST_HOME"; do
    [ -n "$runtime_home" ] || continue
    [ "$runtime_home" = "$seen_home" ] && continue
    seen_home="$runtime_home"
    for sock in \
      "$runtime_home/.orbstack/run/docker.sock" \
      "$runtime_home/.docker/run/docker.sock" \
      "$runtime_home/.rd/docker.sock" \
      "$runtime_home/.colima/default/docker.sock"
    do
      if docker_with_socket_works "$sock"; then
        log "Using existing Docker-compatible runtime socket: $sock"
        export DOCKER_HOST="unix://$sock"
        export DOCKER_CONTEXT=
        return 0
      fi
    done
  done
  return 1
}

colima_docker_works() {
  have colima || return 1
  HOME="$HOST_HOME" colima status >/dev/null 2>&1 || return 1
  sock="$HOST_HOME/.colima/default/docker.sock"
  if docker_with_socket_works "$sock"; then
    log "Using existing Colima runtime socket: $sock"
    export DOCKER_HOST="unix://$sock"
    export DOCKER_CONTEXT=
    return 0
  fi
  HOME="$HOST_HOME" docker context use colima >/dev/null 2>&1 || return 1
  HOME="$HOST_HOME" docker_works
}

wait_for_colima_socket() {
  sock="$HOST_HOME/.colima/default/docker.sock"
  i=0
  while [ "$i" -lt 30 ]; do
    [ -S "$sock" ] && return 0
    i=$((i + 1))
    sleep 1
  done
  fail "Colima started but $sock did not appear"
}

ensure_macos_runtime() {
  if try_explicit_runtime || try_existing_macos_runtime || colima_docker_works || docker_works; then
    return 0
  fi

  ensure_homebrew
  log "Installing Armorer's macOS container runtime: Colima..."
  brew install colima docker docker-compose

  if ! HOME="$HOST_HOME" colima status >/dev/null 2>&1; then
    log "Starting Colima..."
    HOME="$HOST_HOME" colima start --cpu "${ARMORER_COLIMA_CPU:-4}" --memory "${ARMORER_COLIMA_MEMORY:-8}" --disk "${ARMORER_COLIMA_DISK:-40}"
  fi
  wait_for_colima_socket

  HOME="$HOST_HOME" docker context use colima >/dev/null 2>&1 || true
  colima_docker_works || fail "Docker is still unavailable after starting Colima"
}

ensure_linux_runtime() {
  if try_explicit_runtime || docker_works; then
    return 0
  fi
  cat >&2 <<'LINUX_DOCKER'
Docker Engine and the Docker Compose plugin are required.

Install Docker for your Linux distribution, then run:

  docker run --rm hello-world
  docker compose version

After that, rerun the Armorer installer.
LINUX_DOCKER
  exit 1
}

ensure_runtime() {
  if [ "$OS" = "darwin" ]; then
    ensure_macos_runtime
  else
    ensure_linux_runtime
  fi
}

release_base_url() {
  if [ "$VERSION" = "stable" ] || [ "$VERSION" = "latest" ]; then
    printf '%s\n' "https://github.com/ArmorerLabs/Armorer/releases/latest/download"
  else
    printf '%s\n' "https://github.com/ArmorerLabs/Armorer/releases/download/$VERSION"
  fi
}

default_image() {
  if [ -n "$IMAGE" ]; then
    printf '%s\n' "$IMAGE"
  elif [ "$VERSION" = "latest" ]; then
    printf '%s\n' "ghcr.io/armorerlabs/armorer/appliance:stable"
  else
    printf '%s\n' "ghcr.io/armorerlabs/armorer/appliance:$VERSION"
  fi
}

image_exists_locally() {
  docker image inspect "$1" >/dev/null 2>&1
}

start_appliance() {
  image="$1"
  if [ "$NO_PULL" = "1" ]; then
    if ! image_exists_locally "$image"; then
      print_appliance_diagnostics
      fail "local image $image was not found in the selected Docker context. Build/tag it first, for example: docker build -t $image . Or remove --no-pull so Armorer can pull it."
    fi
    if "$BIN_DIR/armorer" start --image "$image" --port "$PORT" --no-pull; then
      return 0
    fi
    print_appliance_diagnostics
    fail "could not start local image $image"
  fi

  if "$BIN_DIR/armorer" start --image "$image" --port "$PORT"; then
    return 0
  fi

  if image_exists_locally "$image"; then
    log "Image pull failed, but $image exists locally. Retrying without pulling..."
    if "$BIN_DIR/armorer" start --image "$image" --port "$PORT" --no-pull; then
      return 0
    fi
  fi

  print_appliance_diagnostics
  fail "could not start $image"
}

verify_appliance_health() {
  port="$1"
  container="${ARMORER_CONTAINER:-armorer}"
  tmp="$(mktemp)"
  i=0
  consecutive=0
  last_restart_count=""
  warned_legacy_readyz="0"
  while [ "$i" -lt 60 ]; do
    health_tmp="$tmp.health"
    ready_tmp="$tmp.ready"
    ready_code="000"
    if curl -fsS "http://127.0.0.1:$port/api/healthz" -o "$health_tmp" >/dev/null 2>&1; then
      ready_code="$(curl -sS "http://127.0.0.1:$port/api/readyz" -o "$ready_tmp" -w '%{http_code}' 2>/dev/null || printf '000')"
    fi
    if [ "$ready_code" = "200" ] || [ "$ready_code" = "404" ]; then
      if [ "$ready_code" = "404" ] && [ "$warned_legacy_readyz" = "0" ]; then
        log "Warning: /api/readyz is not available on this appliance version; using /api/healthz compatibility mode."
        warned_legacy_readyz="1"
      fi
      restart_count="$(docker inspect -f '{{.RestartCount}}' "$container" 2>/dev/null || printf 'unknown')"
      if [ "$restart_count" = "$last_restart_count" ]; then
        consecutive=$((consecutive + 1))
      else
        consecutive=1
        last_restart_count="$restart_count"
      fi
      if [ "$consecutive" -ge 2 ]; then
        rm -f "$tmp" "$health_tmp" "$ready_tmp"
        return 0
      fi
    else
      consecutive=0
      cat "$health_tmp" "$ready_tmp" > "$tmp" 2>/dev/null || true
    fi
    rm -f "$health_tmp" "$ready_tmp"
    i=$((i + 1))
    sleep 1
  done
  if [ -s "$tmp" ]; then
    cat "$tmp" >&2
  fi
  rm -f "$tmp"
  return 1
}

print_appliance_diagnostics() {
  log "Armorer appliance diagnostics:"
  "$BIN_DIR/armorer" logs --tail 100 2>/dev/null || true
  "$BIN_DIR/armorer" doctor 2>/dev/null || true
}

install_launcher() {
  target="$OS-$ARCH"
  archive="armorer-$target.tar.gz"
  url="$(release_base_url)/$archive"
  tmp="$(mktemp -d)"
  trap 'rm -rf "$tmp"' EXIT HUP INT TERM

  mkdir -p "$BIN_DIR"
  log "Downloading Armorer launcher: $url"
  curl -fsSL "$url" -o "$tmp/$archive"
  tar -xzf "$tmp/$archive" -C "$tmp"
  [ -f "$tmp/armorer" ] || fail "release archive did not contain armorer launcher"
  chmod +x "$tmp/armorer"
  mv "$tmp/armorer" "$BIN_DIR/armorer"

  case ":$PATH:" in
    *":$BIN_DIR:"*) ;;
    *) log "Add $BIN_DIR to PATH to run 'armorer' from a new shell." ;;
  esac
}

ensure_git() {
  if have git; then
    return 0
  fi
  if [ "$OS" = "darwin" ]; then
    ensure_homebrew
    brew install git
  else
    fail "git is required for --dev. Install git with your package manager, then rerun this installer."
  fi
}

ensure_node() {
  # We need *some* Node to bootstrap the managed-node installer.
  # Any version is fine here; the installer will download pinned Node 22.
  if have node; then
    return 0
  fi
  if [ "$OS" = "darwin" ]; then
    ensure_homebrew
    log "Installing Node.js (bootstrap only; Armorer will manage its own pinned runtime)..."
    brew install node
    have node || fail "Node.js install did not put node on PATH"
    return 0
  fi
  fail "Node.js is required to bootstrap the Armorer installer. Install any Node.js version, then rerun."
}

ensure_managed_node() {
  ensure_node
  log "Ensuring Armorer managed Node runtime..."
  node "$INSTALL_DIR/scripts/install-managed-node.mjs"
  managed_node_path="$(node "$INSTALL_DIR/scripts/install-managed-node.mjs" --print-path 2>/dev/null || true)"
  if [ -n "$managed_node_path" ] && [ -x "$managed_node_path" ]; then
    export ARMORER_MANAGED_NODE="$managed_node_path"
    log "Armorer managed Node ready: $managed_node_path"
  else
    log "Warning: managed Node install succeeded but binary not found; falling back to system Node."
  fi
  # Ensure corepack is available (from managed Node or system Node).
  if ! have corepack; then
    if [ -n "$managed_node_path" ]; then
      "$(dirname "$managed_node_path")/corepack" enable 2>/dev/null || true
    fi
    have corepack || log "Warning: corepack not found; pnpm may not be available."
  fi
}

ensure_rust() {
  if have cargo; then
    return 0
  fi
  if [ "$OS" = "darwin" ]; then
    ensure_homebrew
    log "Installing Rust for the local Armorer development gateway..."
    brew install rust
    have cargo || fail "Rust install did not put cargo on PATH"
    return 0
  fi
  fail "Rust/Cargo is required for --dev because the Armorer gateway is native. Install Rust, then rerun this installer."
}

sync_source() {
  mkdir -p "$(dirname "$INSTALL_DIR")"
  if [ -d "$INSTALL_DIR/.git" ]; then
    log "Updating Armorer source in $INSTALL_DIR..."
    git -C "$INSTALL_DIR" fetch --depth 1 origin "$BRANCH"
    git -C "$INSTALL_DIR" checkout "$BRANCH"
    git -C "$INSTALL_DIR" reset --hard "origin/$BRANCH"
  elif [ -e "$INSTALL_DIR" ]; then
    fail "$INSTALL_DIR exists but is not a git checkout"
  else
    log "Downloading Armorer source..."
    git clone --depth 1 --branch "$BRANCH" "$REPO_URL" "$INSTALL_DIR"
  fi
}

dev_install() {
  log "Installing Armorer from source for development..."
  ensure_git
  ensure_runtime
  ensure_managed_node
  ensure_rust
  sync_source

  cd "$INSTALL_DIR"
  if ! corepack enable >/dev/null 2>&1; then
    log "Skipping global Corepack shim setup; using 'corepack pnpm' directly."
  fi
  corepack pnpm install --frozen-lockfile
  corepack pnpm build
  corepack pnpm --filter armorer-ui-selfhost build

  mkdir -p "$HOME/.armorer"
  if [ -f "$HOME/.armorer/ui.pid" ]; then
    old_pid="$(cat "$HOME/.armorer/ui.pid" 2>/dev/null || true)"
    if [ -n "$old_pid" ] && kill -0 "$old_pid" 2>/dev/null; then
      kill "$old_pid" 2>/dev/null || true
    fi
  fi
  (nohup env PORT="$PORT" corepack pnpm start > "$HOME/.armorer/ui.log" 2>&1 & echo "$!" > "$HOME/.armorer/ui.pid")
  sleep 2
  new_pid="$(cat "$HOME/.armorer/ui.pid" 2>/dev/null || true)"
  if [ -z "$new_pid" ] || ! kill -0 "$new_pid" 2>/dev/null; then
    tail -n 80 "$HOME/.armorer/ui.log" 2>/dev/null || true
    fail "Armorer UI exited during startup"
  fi
  log "Armorer development UI is ready at http://127.0.0.1:$PORT"
}

main() {
  if [ "$DEV_SOURCE" = "1" ]; then
    dev_install
    exit 0
  fi

  log "Installing Armorer appliance launcher..."
  ensure_runtime
  install_launcher

  if [ "$START_UI" = "1" ]; then
    ARMORER_IMAGE_VALUE="$(default_image)"
    log "Starting Armorer appliance with $ARMORER_IMAGE_VALUE..."
    start_appliance "$ARMORER_IMAGE_VALUE"
    if ! verify_appliance_health "$PORT"; then
      print_appliance_diagnostics
      fail "Armorer appliance did not become healthy at http://127.0.0.1:$PORT/api/healthz"
    fi
    if [ "$OPEN_UI" = "1" ]; then
      "$BIN_DIR/armorer" open --port "$PORT" || true
    fi
  fi

  log "Armorer install complete. Run: armorer status"
}

main "$@"