#!/usr/bin/env bash
# ZonDNS One-Step Installer
# Source: https://github.com/neur3-dev/zondns/blob/main/install.sh
# Review before running: https://github.com/neur3-dev/zondns
#
# Usage:
#   curl -sSL https://install.zondns.cloud | bash
#   bash install.sh [--unattended] [--update] [--uninstall] [--branch <name>]

set -euo pipefail

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
ZONDNS_VERSION="1.0.0"
ZONDNS_REPO="https://github.com/neur3-dev/zondns.git"
ZONDNS_DIR="/opt/zondns"
ZONDNS_BIN="/usr/local/bin/zondns"
BRANCH="${ZONDNS_BRANCH:-main}"

# ---------------------------------------------------------------------------
# Runtime state (set by parse_args)
# ---------------------------------------------------------------------------
UNATTENDED=false
DO_UPDATE=false
DO_UNINSTALL=false

# Config values (set by prompt_config or env)
SERVER_ADDR="${SERVER_ADDR:-}"
ENABLE_AI="${ENABLE_AI:-false}"

# Generated secrets (set by generate_secrets)
DB_PASS=""
DB_APP_PASS=""
JWT_SECRET=""
API_KEY_PEPPER=""
INTELLIGENCE_TOKEN=""

# OS detection
OS_ID=""

# ---------------------------------------------------------------------------
# Error trap
# ---------------------------------------------------------------------------
trap 'echo ""; echo "ERROR: Install failed at line $LINENO. Check logs with: cd $ZONDNS_DIR && docker compose logs"' ERR

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

# ---------------------------------------------------------------------------
# print_banner
# ---------------------------------------------------------------------------
print_banner() {
  echo -e "${CYAN}"
  echo '  ______          _____  _   _  _____ '
  echo ' |___  /         |  __ \| \ | |/ ____|'
  echo '    / / ___  _ __| |  | |  \| | (___  '
  echo '   / / / _ \| '"'"'_ \ |  | | . ` |\___ \ '
  echo '  / /_| (_) | | | | |__| | |\  |____) |'
  echo ' /_____\___/|_| |_|_____/|_| \_|_____/ '
  echo -e "${RESET}"
  echo -e "${BOLD}  v${ZONDNS_VERSION}${RESET}"
  echo ""
}

# ---------------------------------------------------------------------------
# check_root
# ---------------------------------------------------------------------------
check_root() {
  if [[ "$EUID" -ne 0 ]] && ! sudo -n true 2>/dev/null; then
    echo -e "${RED}ERROR: This installer must be run as root.${RESET}"
    echo "Try: sudo bash <(curl -sSL https://install.zondns.cloud)"
    exit 1
  fi
}

# ---------------------------------------------------------------------------
# detect_os
# ---------------------------------------------------------------------------
detect_os() {
  if [[ -f /etc/os-release ]]; then
    # shellcheck disable=SC1091
    OS_ID=$(. /etc/os-release && echo "${ID:-unknown}" | tr '[:upper:]' '[:lower:]')
  else
    echo -e "${RED}ERROR: Cannot detect OS — /etc/os-release not found.${RESET}"
    exit 1
  fi

  case "$OS_ID" in
    ubuntu|debian|centos|rhel|fedora|arch) ;;
    *)
      echo -e "${RED}ERROR: Unsupported OS '${OS_ID}'. Supported: ubuntu, debian, centos, rhel, fedora, arch.${RESET}"
      exit 1
      ;;
  esac

  echo -e "${GREEN}Detected OS: ${OS_ID}${RESET}"
}

# ---------------------------------------------------------------------------
# install_docker
# ---------------------------------------------------------------------------
install_docker() {
  if docker version >/dev/null 2>&1; then
    echo -e "${GREEN}Docker already installed.${RESET}"
    return 0
  fi

  echo -e "${YELLOW}Installing Docker...${RESET}"
  case "$OS_ID" in
    ubuntu|debian)
      curl -fsSL https://get.docker.com | sh
      ;;
    centos|rhel|fedora)
      dnf install -y docker-ce docker-ce-cli containerd.io 2>/dev/null \
        || yum install -y docker-ce docker-ce-cli containerd.io
      ;;
    arch)
      pacman -Sy --noconfirm docker
      ;;
  esac

  systemctl enable --now docker
  echo -e "${GREEN}Docker installed and started.${RESET}"
}

# ---------------------------------------------------------------------------
# install_compose
# ---------------------------------------------------------------------------
install_compose() {
  if docker compose version >/dev/null 2>&1; then
    echo -e "${GREEN}Docker Compose already installed.${RESET}"
    return 0
  fi

  echo -e "${YELLOW}Installing Docker Compose plugin v2.27.0...${RESET}"

  local os_lower
  os_lower=$(uname -s | tr '[:upper:]' '[:lower:]')
  local arch
  arch=$(uname -m)

  local compose_url="https://github.com/docker/compose/releases/download/v2.27.0/docker-compose-${os_lower}-${arch}"
  local plugin_dir="/usr/local/lib/docker/cli-plugins"

  mkdir -p "$plugin_dir"
  curl -fsSL "$compose_url" -o "${plugin_dir}/docker-compose"
  chmod +x "${plugin_dir}/docker-compose"

  echo -e "${GREEN}Docker Compose plugin installed.${RESET}"
}

# ---------------------------------------------------------------------------
# clone_or_update
# ---------------------------------------------------------------------------
clone_or_update() {
  if [[ -d "${ZONDNS_DIR}/.git" ]]; then
    echo -e "${YELLOW}Updating existing ZonDNS installation...${RESET}"
    git -C "$ZONDNS_DIR" fetch origin
    git -C "$ZONDNS_DIR" checkout "$BRANCH"
    git -C "$ZONDNS_DIR" pull --ff-only
  else
    echo -e "${YELLOW}Cloning ZonDNS repository...${RESET}"
    git clone --depth 1 --branch "$BRANCH" "$ZONDNS_REPO" "$ZONDNS_DIR"
  fi
  echo -e "${GREEN}Repository ready at ${ZONDNS_DIR}.${RESET}"
}

# ---------------------------------------------------------------------------
# generate_secrets
# ---------------------------------------------------------------------------
generate_secrets() {
  echo -e "${YELLOW}Generating secrets...${RESET}"
  DB_PASS=$(openssl rand -hex 16)
  DB_APP_PASS=$(openssl rand -hex 16)
  JWT_SECRET=$(openssl rand -hex 32)
  API_KEY_PEPPER=$(openssl rand -hex 32)
  INTELLIGENCE_TOKEN=$(openssl rand -hex 32)
  echo -e "${GREEN}Secrets generated.${RESET}"
}

# ---------------------------------------------------------------------------
# prompt_config
# ---------------------------------------------------------------------------
prompt_config() {
  if [[ "$UNATTENDED" == "true" ]]; then
    # Fall back to env vars / defaults
    if [[ -z "$SERVER_ADDR" ]]; then
      SERVER_ADDR=$(curl -s --max-time 5 ifconfig.me 2>/dev/null || echo "127.0.0.1")
    fi
    return 0
  fi

  echo ""
  read -r -p "Enter your server's public IP or domain (used for APP_URL) [default: auto-detect]: " input_addr </dev/tty
  if [[ -z "$input_addr" ]]; then
    SERVER_ADDR=$(curl -s --max-time 5 ifconfig.me 2>/dev/null || echo "127.0.0.1")
    echo "Auto-detected: ${SERVER_ADDR}"
  else
    SERVER_ADDR="$input_addr"
  fi

  read -r -p "Enable AI threat detection (DGA/tunneling analysis)? [y/N]: " input_ai </dev/tty
  case "${input_ai,,}" in
    y|yes) ENABLE_AI=true ;;
    *)     ENABLE_AI=false ;;
  esac
}

# ---------------------------------------------------------------------------
# write_env
# ---------------------------------------------------------------------------
write_env() {
  if [[ -f "${ZONDNS_DIR}/.env" ]]; then
    echo -e "${YELLOW}Existing .env found — keeping secrets. Run with --update to upgrade containers.${RESET}"
    return 0
  fi

  echo -e "${YELLOW}Writing .env file...${RESET}"

  cat > "${ZONDNS_DIR}/.env" <<EOF
DB_PASS=${DB_PASS}
DB_APP_USER=dns_runtime
DB_APP_PASS=${DB_APP_PASS}
JWT_SECRET=${JWT_SECRET}
API_KEY_PEPPER=${API_KEY_PEPPER}
APP_URL=http://${SERVER_ADDR}
DNS_FALLBACK=9.9.9.9:53
POLICY_FAIL_CLOSED=false
# TRUSTED_PROXIES=
# ALLOWED_CLIENT_CIDRS=
EOF

  if [[ "$ENABLE_AI" == "true" ]]; then
    echo "INTELLIGENCE_TOKEN=${INTELLIGENCE_TOKEN}" >> "${ZONDNS_DIR}/.env"
  fi

  chmod 600 "${ZONDNS_DIR}/.env"
  echo -e "${GREEN}.env written to ${ZONDNS_DIR}/.env${RESET}"
}

# ---------------------------------------------------------------------------
# start_services
# ---------------------------------------------------------------------------
start_services() {
  echo -e "${YELLOW}Starting ZonDNS services...${RESET}"

  local compose_cmd="docker compose -f docker-compose.yml -f docker-compose.prod.yml"

  if [[ "$ENABLE_AI" == "true" ]]; then
    $compose_cmd --profile intelligence up -d --remove-orphans
  else
    $compose_cmd up -d --remove-orphans
  fi

  echo -e "${GREEN}Services started.${RESET}"
}

# ---------------------------------------------------------------------------
# wait_healthy
# ---------------------------------------------------------------------------
wait_healthy() {
  echo -e "${YELLOW}Waiting for services to be healthy...${RESET}"

  local spinners=('⣾' '⣽' '⣻' '⢿' '⡿' '⣟' '⣯' '⣷')
  local timeout=60
  local elapsed=0
  local idx=0

  while [[ $elapsed -lt $timeout ]]; do
    # Count containers not yet in "running" state (catches starting AND crashed)
    local not_running
    not_running=$(cd "$ZONDNS_DIR" && docker compose ps --format json 2>/dev/null \
      | grep -cv '"State":"running"' 2>/dev/null || true)

    if [[ "$not_running" -eq 0 ]]; then
      echo -e "\r${GREEN}All services running.       ${RESET}"
      return 0
    fi

    printf "\r  ${spinners[$idx]} Starting... (%ds)" "$elapsed"
    idx=$(( (idx + 1) % ${#spinners[@]} ))
    sleep 1
    elapsed=$((elapsed + 1))
  done

  echo -e "\r${YELLOW}WARNING: Services may still be starting. Check with: zondns status${RESET}"
}

# ---------------------------------------------------------------------------
# install_cli_shim
# ---------------------------------------------------------------------------
install_cli_shim() {
  echo -e "${YELLOW}Installing zondns CLI shim...${RESET}"

  cat > "$ZONDNS_BIN" <<'SHIM'
#!/usr/bin/env bash
# ZonDNS CLI shim
set -euo pipefail
ZONDNS_DIR=/opt/zondns
case "${1:-help}" in
  update)    bash "$ZONDNS_DIR/install.sh" --update ;;
  status)    (cd "$ZONDNS_DIR" && docker compose ps) ;;
  logs)      (cd "$ZONDNS_DIR" && docker compose logs --tail=50 -f "${2:-}") ;;
  uninstall) bash "$ZONDNS_DIR/install.sh" --uninstall ;;
  help|*)
    echo "Usage: zondns {update|status|logs [service]|uninstall}"
    ;;
esac
SHIM

  chmod +x "$ZONDNS_BIN"
  echo -e "${GREEN}CLI shim installed at ${ZONDNS_BIN}.${RESET}"
}

# ---------------------------------------------------------------------------
# print_summary
# ---------------------------------------------------------------------------
print_summary() {
  echo ""
  echo -e "${GREEN}┌──────────────────────────────────────────┐${RESET}"
  echo -e "${GREEN}│        ZonDNS installed successfully!    │${RESET}"
  echo -e "${GREEN}│                                          │${RESET}"
  printf  "${GREEN}│${RESET}  Dashboard:  http://%-22s${GREEN}│${RESET}\n" "${SERVER_ADDR}"
  printf  "${GREEN}│${RESET}  DNS server: %-26s${GREEN}│${RESET}\n" "${SERVER_ADDR}:53"
  printf  "${GREEN}│${RESET}  API:        http://%-18s${GREEN}│${RESET}\n" "${SERVER_ADDR}/api"
  echo -e "${GREEN}│                                          │${RESET}"
  echo -e "${GREEN}│  Manage:  zondns {update|status|logs}   │${RESET}"
  echo -e "${GREEN}│  Docs:    https://docs.zondns.cloud      │${RESET}"
  echo -e "${GREEN}└──────────────────────────────────────────┘${RESET}"
  echo ""
}

# ---------------------------------------------------------------------------
# parse_args
# ---------------------------------------------------------------------------
parse_args() {
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --unattended) UNATTENDED=true; shift ;;
      --update)     DO_UPDATE=true;  shift ;;
      --uninstall)  DO_UNINSTALL=true; shift ;;
      --branch)
        if [[ -z "${2:-}" ]]; then
          echo "ERROR: --branch requires a value." >&2
          exit 1
        fi
        BRANCH="$2"; shift 2 ;;
      *)
        echo "Unknown option: $1" >&2
        echo "Usage: install.sh [--unattended] [--update] [--uninstall] [--branch <name>]" >&2
        exit 1 ;;
    esac
  done
}

# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
print_banner
parse_args "$@"
check_root

# ---- Uninstall flow --------------------------------------------------------
if [[ "$DO_UNINSTALL" == "true" ]]; then
  if [[ "$UNATTENDED" != "true" ]]; then
    read -r -p "Are you sure you want to uninstall ZonDNS? This will remove all data. [y/N]: " confirm </dev/tty
    case "${confirm,,}" in
      y|yes) ;;
      *) echo "Uninstall cancelled."; exit 0 ;;
    esac
  fi

  echo -e "${YELLOW}Stopping and removing ZonDNS containers and volumes...${RESET}"
  if [[ -d "$ZONDNS_DIR" ]]; then
    (cd "$ZONDNS_DIR" && docker compose down --volumes 2>/dev/null || true)
    rm -rf "$ZONDNS_DIR"
  fi
  rm -f "$ZONDNS_BIN"
  echo -e "${GREEN}ZonDNS uninstalled.${RESET}"
  exit 0
fi

# ---- Common prerequisites --------------------------------------------------
detect_os
install_docker
install_compose

# ---- Update flow -----------------------------------------------------------
if [[ "$DO_UPDATE" == "true" ]]; then
  clone_or_update
  # Re-read .env so start_services picks up ENABLE_AI and other vars
  if [[ -f "${ZONDNS_DIR}/.env" ]]; then
    # shellcheck disable=SC1091
    set -a; . "${ZONDNS_DIR}/.env"; set +a
  fi
  (cd "$ZONDNS_DIR" && start_services)
  echo -e "${GREEN}ZonDNS updated.${RESET}"
  exit 0
fi

# ---- Fresh install ---------------------------------------------------------
clone_or_update
generate_secrets
prompt_config
write_env
(cd "$ZONDNS_DIR" && start_services)
wait_healthy
install_cli_shim
print_summary