#!/usr/bin/env bash
# ============================================================
#  setup-ssh-client.sh
#  Sets up cloudflared SSH client on macOS or Linux (Debian/Ubuntu)
#  Usage: bash setup-ssh-client.sh
# ============================================================

set -euo pipefail

# ── Colours ─────────────────────────────────────────────────
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
CYAN='\033[0;36m'; BOLD='\033[1m'; RESET='\033[0m'

info()    { echo -e "${CYAN}[INFO]${RESET}  $*"; }
success() { echo -e "${GREEN}[OK]${RESET}    $*"; }
warn()    { echo -e "${YELLOW}[WARN]${RESET}  $*"; }
error()   { echo -e "${RED}[ERROR]${RESET} $*"; exit 1; }
prompt()  { echo -e "${BOLD}$*${RESET}"; }

echo ""
echo -e "${BOLD}╔══════════════════════════════════════════════╗${RESET}"
echo -e "${BOLD}║      Cloudflare SSH Client Setup Script      ║${RESET}"
echo -e "${BOLD}╚══════════════════════════════════════════════╝${RESET}"
echo ""

# ── 1. Detect OS ─────────────────────────────────────────────
detect_os() {
  if [[ "$OSTYPE" == "darwin"* ]]; then
    OS="macos"
  elif [[ -f /etc/debian_version ]]; then
    OS="debian"
  elif [[ -f /etc/redhat-release ]]; then
    OS="redhat"
  else
    error "Unsupported OS. This script supports macOS, Debian/Ubuntu, and RHEL/CentOS."
  fi
  info "Detected OS: $OS"
}

# ── 2. Install cloudflared ────────────────────────────────────
install_cloudflared() {
  if command -v cloudflared &>/dev/null; then
    INSTALLED_VER=$(cloudflared --version 2>&1 | head -n1)
    success "cloudflared already installed: $INSTALLED_VER"
    return
  fi

  info "cloudflared not found. Installing..."

  case "$OS" in
    macos)
      if ! command -v brew &>/dev/null; then
        error "Homebrew not found. Install it first: https://brew.sh"
      fi
      brew install cloudflared
      ;;

    debian)
      ARCH=$(dpkg --print-architecture)
      DEB_URL="https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-${ARCH}.deb"
      info "Downloading cloudflared for arch: $ARCH"
      TMP_DEB=$(mktemp /tmp/cloudflared-XXXXXX.deb)
      curl -fsSL "$DEB_URL" -o "$TMP_DEB"
      sudo dpkg -i "$TMP_DEB"
      rm -f "$TMP_DEB"
      ;;

    redhat)
      RPM_URL="https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-x86_64.rpm"
      TMP_RPM=$(mktemp /tmp/cloudflared-XXXXXX.rpm)
      curl -fsSL "$RPM_URL" -o "$TMP_RPM"
      sudo rpm -i "$TMP_RPM"
      rm -f "$TMP_RPM"
      ;;
  esac

  if command -v cloudflared &>/dev/null; then
    success "cloudflared installed successfully: $(cloudflared --version 2>&1 | head -n1)"
  else
    error "cloudflared installation failed."
  fi
}

# ── 3. Collect user input ─────────────────────────────────────
collect_input() {
  echo ""
  prompt "Enter your SSH tunnel hostname (e.g. ssh.yourdomain.com):"
  read -r SSH_HOSTNAME
  [[ -z "$SSH_HOSTNAME" ]] && error "Hostname cannot be empty."

  prompt "Enter your Ubuntu server username (e.g. john):"
  read -r SSH_USER
  [[ -z "$SSH_USER" ]] && error "Username cannot be empty."

  echo ""
  info "Hostname : $SSH_HOSTNAME"
  info "User     : $SSH_USER"
}

# ── 4. Write SSH config ───────────────────────────────────────
write_ssh_config() {
  SSH_CONFIG="$HOME/.ssh/config"
  CLOUDFLARED_BIN=$(command -v cloudflared)

  mkdir -p "$HOME/.ssh"
  chmod 700 "$HOME/.ssh"

  # Check if entry already exists
  if grep -q "Host $SSH_HOSTNAME" "$SSH_CONFIG" 2>/dev/null; then
    warn "An SSH config entry for '$SSH_HOSTNAME' already exists."
    prompt "Overwrite it? [y/N]:"
    read -r OVERWRITE
    if [[ ! "$OVERWRITE" =~ ^[Yy]$ ]]; then
      info "Skipping SSH config update."
      return
    fi
    # Remove old block for this host
    # Use awk to delete the matching Host block
    awk "/^Host $SSH_HOSTNAME$/{found=1} found && /^Host / && !/^Host $SSH_HOSTNAME$/{found=0} !found" \
      "$SSH_CONFIG" > /tmp/ssh_config_tmp && mv /tmp/ssh_config_tmp "$SSH_CONFIG"
    info "Removed old entry for $SSH_HOSTNAME."
  fi

  # Append new config block
  cat >> "$SSH_CONFIG" <<EOF

# Cloudflare Tunnel SSH — added by setup-ssh-client.sh
Host $SSH_HOSTNAME
    ProxyCommand $CLOUDFLARED_BIN access ssh --hostname %h
    User $SSH_USER
    StrictHostKeyChecking accept-new
EOF

  chmod 600 "$SSH_CONFIG"
  success "SSH config written to $SSH_CONFIG"
}

# ── 5. Verify setup ───────────────────────────────────────────
verify_setup() {
  echo ""
  info "Verifying setup..."

  command -v cloudflared &>/dev/null && success "cloudflared: $(cloudflared --version 2>&1 | head -n1)" \
    || warn "cloudflared not found in PATH after install — you may need to restart your shell."

  grep -q "Host $SSH_HOSTNAME" "$HOME/.ssh/config" 2>/dev/null \
    && success "SSH config entry found for $SSH_HOSTNAME" \
    || warn "SSH config entry not found — check $HOME/.ssh/config manually."
}

# ── 6. Print summary ──────────────────────────────────────────
print_summary() {
  echo ""
  echo -e "${GREEN}${BOLD}════════════════════════════════════════${RESET}"
  echo -e "${GREEN}${BOLD}  ✓ Setup complete!${RESET}"
  echo -e "${GREEN}${BOLD}════════════════════════════════════════${RESET}"
  echo ""
  echo -e "  Connect to your server with:"
  echo -e "  ${BOLD}ssh $SSH_HOSTNAME${RESET}"
  echo ""
  echo -e "  Or with a specific user:"
  echo -e "  ${BOLD}ssh ${SSH_USER}@${SSH_HOSTNAME}${RESET}"
  echo ""
  echo -e "  ${YELLOW}Note:${RESET} Make sure cloudflared tunnel is running on your server"
  echo -e "  and the tunnel is routed to: ${BOLD}$SSH_HOSTNAME${RESET}"
  echo ""
}

# ── Main ──────────────────────────────────────────────────────
main() {
  detect_os
  install_cloudflared
  collect_input
  write_ssh_config
  verify_setup
  print_summary
}

main