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

usage() {
  cat <<'USAGE'
Configure an SSH alias and optionally install the current public key for passwordless login.

Usage:
  ssh_alias_setup.sh [options] ALIAS HOST

Options:
  -u, --user USER     SSH user (default: root)
  -p, --port PORT     SSH port (default: 22)
  -i, --key PATH      Identity file (default: ~/.ssh/id_ed25519)
      --install-key   Run ssh-copy-id after writing the alias
      --no-accept-new Disable StrictHostKeyChecking=accept-new
      --dry-run       Print planned actions without changing files
  -h, --help          Show help

Environment:
  SSH_CONFIG_FILE     Override SSH config path for testing
  REMOTE_KEY          Default key path
  REMOTE_PORT         Default SSH port

Examples:
  ssh_alias_setup.sh bigmengya 192.168.1.233 --install-key
  ssh_alias_setup.sh -u ubuntu -p 2222 devbox 10.0.0.8
USAGE
}

expand_path() {
  local value="$1"
  case "$value" in
    ~) printf '%s\n' "$HOME" ;;
    ~/*) printf '%s/%s\n' "$HOME" "${value#~/}" ;;
    *) printf '%s\n' "$value" ;;
  esac
}

user="root"
port="${REMOTE_PORT:-22}"
key="${REMOTE_KEY:-$HOME/.ssh/id_ed25519}"
accept_new=true
install_key=false
dry_run=false

while [[ $# -gt 0 ]]; do
  case "$1" in
    -u|--user)
      user="${2:-}"
      shift 2
      ;;
    -p|--port)
      port="${2:-}"
      shift 2
      ;;
    -i|--key)
      key="${2:-}"
      shift 2
      ;;
    --install-key)
      install_key=true
      shift
      ;;
    --no-accept-new)
      accept_new=false
      shift
      ;;
    --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 [[ $# -lt 2 ]]; then
  usage >&2
  exit 2
fi

alias_name="$1"
host="$2"
key="$(expand_path "$key")"
config_file="$(expand_path "${SSH_CONFIG_FILE:-$HOME/.ssh/config}")"
ssh_dir="$(dirname "$config_file")"
key_pub="${key}.pub"
key_dir="$(dirname "$key")"

if $dry_run; then
  echo "Would ensure directory: $ssh_dir"
  echo "Would ensure key directory: $key_dir"
else
  mkdir -p "$ssh_dir" "$key_dir"
  chmod 700 "$ssh_dir" "$key_dir"
fi

if [[ ! -f "$key" ]]; then
  if $dry_run; then
    printf 'Would generate key: ssh-keygen -t ed25519 -C %q -f %q -N %q\n' "codex" "$key" ""
  else
    ssh-keygen -t ed25519 -C "codex" -f "$key" -N ""
  fi
fi

if [[ ! -f "$key_pub" ]]; then
  echo "Public key not found: $key_pub" >&2
  exit 1
fi

if $dry_run; then
  cat <<DRYRUN
Would write/update SSH alias:
  Host $alias_name
  HostName $host
  User $user
  Port $port
  ServerAliveInterval 30
  ServerAliveCountMax 3
  ConnectTimeout 10
  StrictHostKeyChecking $( $accept_new && printf 'accept-new' || printf 'ask' )
  IdentityFile $key
  IdentitiesOnly yes
DRYRUN
else
  python3 - "$config_file" "$alias_name" "$host" "$user" "$port" "$key" "$accept_new" <<'PY'
from pathlib import Path
import sys

config_path = Path(sys.argv[1])
alias_name = sys.argv[2]
host = sys.argv[3]
user = sys.argv[4]
port = sys.argv[5]
key = sys.argv[6]
accept_new = sys.argv[7].lower() == 'true'

existing = config_path.read_text(encoding='utf-8') if config_path.exists() else ''
lines = existing.splitlines()
out: list[str] = []
skip = False
for line in lines:
    stripped = line.strip()
    if stripped.lower().startswith('host '):
        host_names = stripped[5:].split()
        if alias_name in host_names:
            skip = True
            continue
        skip = False
        out.append(line)
        continue
    if skip:
        continue
    out.append(line)

block = [
    f'Host {alias_name}',
    f'  HostName {host}',
    f'  User {user}',
    f'  Port {port}',
    '  ServerAliveInterval 30',
    '  ServerAliveCountMax 3',
    '  ConnectTimeout 10',
    f'  StrictHostKeyChecking {"accept-new" if accept_new else "ask"}',
    f'  IdentityFile {key}',
    '  IdentitiesOnly yes',
]
text = '\n'.join(out).strip()
if text:
    text += '\n\n'
text += '\n'.join(block) + '\n'
config_path.write_text(text, encoding='utf-8')
PY
  chmod 600 "$config_file"
fi

printf 'Alias ready: ssh %s\n' "$alias_name"
printf 'Config file: %s\n' "$config_file"

if $install_key; then
  if ! command -v ssh-copy-id >/dev/null 2>&1; then
    echo 'ssh-copy-id is required but not installed.' >&2
    exit 1
  fi
  if ! [ -t 0 ]; then
    echo 'Warning: no TTY detected; password prompt may fail. Prefer running with a terminal/TTY.' >&2
  fi
  if $dry_run; then
    printf 'Would run: ssh-copy-id -i %q %q\n' "$key_pub" "$alias_name"
  else
    ssh-copy-id -i "$key_pub" "$alias_name"
  fi
fi