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

SCRIPT_NAME="$(basename "$0")"
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
SKILL_ROOT="$(dirname "$SCRIPT_DIR")"

: "${OUTPUT_DIR:=$SKILL_ROOT/outputs}"
: "${BACKUP_DIR:=$OUTPUT_DIR/backups}"

log_info()  { echo "[INFO]  $(date -Iseconds) $*"; }
log_warn()  { echo "[WARN]  $(date -Iseconds) $*" >&2; }
log_error() { echo "[ERROR] $(date -Iseconds) $*" >&2; }

die() { log_error "$@"; exit 1; }

require_confirm() {
  local require="${REQUIRE_CONFIRM:-1}"
  local phrase="${CONFIRM_PHRASE:-I UNDERSTAND THIS CAN LOCK ME OUT}"
  if [[ "$require" != "1" ]]; then
    return 0
  fi
  echo
  echo "CONFIRMATION REQUIRED"
  echo "Type the phrase exactly to continue:"
  echo "  $phrase"
  read -r input
  [[ "$input" == "$phrase" ]] || die "Confirmation phrase mismatch; aborting."
}

detect_ssh_client_ip() {
  if [[ -n "${SSH_CLIENT:-}" ]]; then
    awk '{print $1}' <<<"$SSH_CLIENT"
  elif [[ -n "${SSH_CONNECTION:-}" ]]; then
    awk '{print $1}' <<<"$SSH_CONNECTION"
  else
    echo ""
  fi
}

ensure_ufw() {
  if command -v ufw &>/dev/null; then
    return 0
  fi
  log_warn "ufw not found; attempting install (Debian/Ubuntu)"
  if command -v apt-get &>/dev/null; then
    sudo apt-get update -y
    sudo apt-get install -y ufw
  else
    die "ufw not installed and apt-get not available. Install ufw manually."
  fi
}

main() {
  mkdir -p "$OUTPUT_DIR" "$BACKUP_DIR"

  local dry="${DRY_RUN:-1}"
  [[ "$dry" == "0" ]] || die "Refusing to apply with DRY_RUN=$dry. Export DRY_RUN=0 to proceed."

  require_confirm

  ensure_ufw

  local ssh_port="${SSH_PORT:-22}"
  local allow_http="${ALLOW_HTTP:-true}"
  local allow_https="${ALLOW_HTTPS:-true}"
  local ssh_ip
  ssh_ip="$(detect_ssh_client_ip)"

  log_info "Backing up current UFW status"
  sudo ufw status verbose >"$BACKUP_DIR/ufw_status_before.txt" || true
  sudo iptables -S >"$BACKUP_DIR/iptables_rules_before.txt" || true

  log_info "Applying UFW policy (idempotent)"
  sudo ufw --force reset
  sudo ufw default deny incoming
  sudo ufw default allow outgoing

  # Always allow SSH
  sudo ufw allow "$ssh_port"/tcp
  if [[ -n "$ssh_ip" ]]; then
    sudo ufw allow from "$ssh_ip" to any port "$ssh_port" proto tcp
  fi
  if [[ "$allow_http" == "true" ]]; then
    sudo ufw allow 80/tcp
  fi
  if [[ "$allow_https" == "true" ]]; then
    sudo ufw allow 443/tcp
  fi

  log_info "Enabling UFW"
  sudo ufw --force enable

  log_info "Writing post-state"
  sudo ufw status verbose >"$OUTPUT_DIR/ufw_status_after.txt"

  log_info "UFW apply complete"
}

[[ "${BASH_SOURCE[0]}" == "$0" ]] && main "$@"