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

usage() {
  cat <<'USAGE'
Run a remote command over SSH with consistent, script-friendly options.

Usage:
  ssh_run.sh [options] HOST -- COMMAND [ARG...]
  ssh_run.sh [options] HOST                # interactive shell

Options:
  -u, --user USER              Override SSH user (or set REMOTE_USER)
  -p, --port PORT              SSH port (default: REMOTE_PORT or 22)
  -i, --key PATH               Identity file (default: REMOTE_KEY)
  -t, --tty                    Force pseudo-tty allocation (useful for sudo prompts)
      --accept-new             Set StrictHostKeyChecking=accept-new
      --sudo                   Prefix command with sudo --
      --sudo-non-interactive   Prefix command with sudo -n -- (fails if password needed)
      --connect-timeout SEC    Connect timeout (default: REMOTE_CONNECT_TIMEOUT or 10)
      --dry-run                Print the ssh command that would run
  -h, --help                   Show help

Environment defaults:
  REMOTE_USER, REMOTE_PORT, REMOTE_KEY, REMOTE_CONNECT_TIMEOUT

Examples:
  ssh_run.sh --user ubuntu 10.0.0.1 -- uname -a
  ssh_run.sh --tty --sudo my-server -- systemctl restart nginx
USAGE
}

fail() {
  echo "Error: $*" >&2
  exit 2
}

require_arg() {
  local value="${1:-}"
  local opt="${2:-option}"
  [[ -n "$value" ]] || fail "$opt requires a value"
  printf '%s' "$value"
}

port="${REMOTE_PORT:-22}"
user="${REMOTE_USER:-}"
key="${REMOTE_KEY:-}"
connect_timeout="${REMOTE_CONNECT_TIMEOUT:-10}"

tty=false
accept_new=false
sudo_mode=""
dry_run=false

while [[ $# -gt 0 ]]; do
  case "$1" in
    -u|--user)
      user="$(require_arg "${2:-}" "$1")"
      shift 2
      ;;
    -p|--port)
      port="$(require_arg "${2:-}" "$1")"
      shift 2
      ;;
    -i|--key)
      key="$(require_arg "${2:-}" "$1")"
      shift 2
      ;;
    -t|--tty)
      tty=true
      shift
      ;;
    --accept-new)
      accept_new=true
      shift
      ;;
    --sudo)
      sudo_mode="sudo"
      shift
      ;;
    --sudo-non-interactive)
      sudo_mode="sudo-n"
      shift
      ;;
    --connect-timeout)
      connect_timeout="$(require_arg "${2:-}" "$1")"
      shift 2
      ;;
    --dry-run)
      dry_run=true
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    --)
      shift
      break
      ;;
    -*)
      echo "Unknown option: $1" >&2
      usage >&2
      exit 2
      ;;
    *)
      break
      ;;
  esac
done

if ! [[ "$port" =~ ^[0-9]+$ ]]; then
  fail "Invalid port: $port"
fi

if ! [[ "$connect_timeout" =~ ^[0-9]+$ ]]; then
  fail "Invalid connect timeout: $connect_timeout"
fi

if [[ $# -lt 1 ]]; then
  usage >&2
  exit 2
fi

host="$1"
shift

dest="$host"
if [[ -n "$user" ]]; then
  host_no_user="${host#*@}"
  dest="${user}@${host_no_user}"
fi

ssh_opts=(
  -p "$port"
  -o "ConnectTimeout=${connect_timeout}"
  -o "ServerAliveInterval=30"
  -o "ServerAliveCountMax=3"
)

if [[ -n "$key" ]]; then
  ssh_opts+=(-i "$key" -o "IdentitiesOnly=yes")
fi

if $accept_new; then
  ssh_opts+=(-o "StrictHostKeyChecking=accept-new")
fi

if $tty; then
  ssh_opts+=(-tt)
fi

cmd=("$@")
if [[ ${#cmd[@]} -gt 0 && "${cmd[0]}" == "--" ]]; then
  cmd=("${cmd[@]:1}")
fi
if [[ -n "$sudo_mode" && ${#cmd[@]} -gt 0 ]]; then
  if [[ "$sudo_mode" == "sudo-n" ]]; then
    cmd=("sudo" "-n" "--" "${cmd[@]}")
  else
    cmd=("sudo" "--" "${cmd[@]}")
  fi
fi

full_cmd=(ssh "${ssh_opts[@]}" "$dest")
if [[ ${#cmd[@]} -gt 0 ]]; then
  full_cmd+=("${cmd[@]}")
fi

if $dry_run; then
  printf '%q ' "${full_cmd[@]}"
  printf '\n'
  exit 0
fi

"${full_cmd[@]}"