#!/bin/bash
# ==============================================================================
# Wallarm Native Node Manager - Install, Configure, Remove, and Control
# ==============================================================================
# Unified single-script manager for the Wallarm Native Node (connector mode,
# NO Docker). Manages multiple isolated nodes under ${BASE_DIR}/nodes with a
# systemd template unit (wallarm-node@<name>.service).
#
# NOTE: This targets the Wallarm Native Node product (go-node, connector-server
# mode, all-in-one installer) - distinct from the NGINX-module based native
# deployment in ./wallarm-ct-deploy.sh.
#
# Commands:
#   --preflight           Run preflight checks only (no installation).
#   --install             Interactive installation of one or more nodes (parallel).
#   --config              Update an existing node's configuration.
#                         Options: --node NAME --address IP:PORT [--token TOKEN] [--labels LABELS]
#   --remove              Remove a node completely.
#                         Options: --node NAME
#   --status [NODE]       Show systemd status for a node, or all nodes.
#   --help|-h             Show help.
# ==============================================================================

set -euo pipefail

# Script location and shared library (logging, detection, connectivity, validation)
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck source=./wallarm-lib.sh
source "$SCRIPT_DIR/wallarm-lib.sh"

# --- Constants ---
BASE_DIR="/opt/wallarm"
NODES_DIR="${BASE_DIR}/nodes"
SYSTEMD_TEMPLATE="/etc/systemd/system/wallarm-node@.service"

# Wallarm Native Node all-in-one installer (latest, x86_64 by default)
# Override WALLARM_INSTALLER_URL to pin a version; WALLARM_INSTALLER_ARCH to
# select the architecture suffix.
INSTALLER_BASE_URL="https://repo.wallarm.com/linux/wallarm-native-node/latest/all-in-one"
INSTALLER_ARCH="${WALLARM_INSTALLER_ARCH:-x86_64}"
INSTALLER_URL="${WALLARM_INSTALLER_URL:-${INSTALLER_BASE_URL}/wallarm-native-node-aio-${INSTALLER_ARCH}-latest.sh}"

# Wallarm cloud endpoints (for connectivity checks)
EU_DATA_NODES=("api.wallarm.com" "node-data0.eu1.wallarm.com" "node-data1.eu1.wallarm.com")
US_DATA_NODES=("us1.api.wallarm.com" "node-data0.us1.wallarm.com" "node-data1.us1.wallarm.com")

# Cloud region selection (populated by preflight, used by select_cloud_region)
US_CLOUD_REACHABLE="false"
EU_CLOUD_REACHABLE="false"
CLOUD_REGION=""
API_HOST=""

# --- Helper functions ---
log() { echo ">>> $*"; }
err() { echo "!!! $*" >&2; }

check_root() {
    if [[ $EUID -ne 0 ]]; then
        err "This script must be run as root (for systemd and /opt write access)."
        exit 1
    fi
}

ensure_dirs() {
    mkdir -p "${NODES_DIR}"
}

generate_systemd_template() {
    if [[ ! -f "${SYSTEMD_TEMPLATE}" ]]; then
        log "Creating systemd template: ${SYSTEMD_TEMPLATE}"
        cat > "${SYSTEMD_TEMPLATE}" <<EOF
[Unit]
Description=Wallarm Native Node - %I
After=network.target

[Service]
Type=simple
WorkingDirectory=${NODES_DIR}/%i
EnvironmentFile=${NODES_DIR}/%i/env
ExecStart=${NODES_DIR}/%i/wallarm-native-node-aio.sh start
ExecStop=${NODES_DIR}/%i/wallarm-native-node-aio.sh stop
Restart=on-failure
RestartSec=5
User=root

[Install]
WantedBy=multi-user.target
EOF
        systemctl daemon-reload
    fi
}

write_env_file() {
    local node_name="$1"
    local api_token="$2"
    local api_host="$3"
    local labels="$4"
    local config_path="${NODES_DIR}/${node_name}/etc/go-node.yaml"
    local env_file="${NODES_DIR}/${node_name}/env"

    cat > "${env_file}" <<EOF
WALLARM_API_TOKEN=${api_token}
WALLARM_API_HOST=${api_host}
WALLARM_LABELS=${labels}
WALLARM_CONFIG_PATH=${config_path}
EOF
    chmod 600 "${env_file}"  # token is sensitive
}

# --- Preflight checks ---
run_preflight() {
    local failed=0

    log "Running preflight checks..."

    # 1. Root privileges
    if [[ $EUID -ne 0 ]]; then
        err "Preflight failed: must be run as root."
        failed=1
    fi

    # 2. Init system must be systemd (template unit management)
    local init_system
    init_system=$(detect_init_system)
    if [[ "$init_system" != "systemd" ]]; then
        err "Preflight failed: this manager requires systemd (detected: $init_system)."
        failed=1
    fi

    # 3. Architecture
    local arch
    arch=$(detect_architecture)
    if [[ "$arch" != "x86_64" && "$arch" != "aarch64" ]]; then
        err "Preflight failed: unsupported architecture '$arch' for the native node installer."
        failed=1
    fi
    log "Architecture: ${arch} (installer suffix: ${INSTALLER_ARCH})"

    # 4. Required commands
    local required_cmds=(curl systemctl sed mkdir rm sleep)
    local missing=()
    local cmd
    for cmd in "${required_cmds[@]}"; do
        if ! command_exists "$cmd"; then
            missing+=("$cmd")
        fi
    done
    if [[ ${#missing[@]} -gt 0 ]]; then
        err "Preflight failed: missing required commands: ${missing[*]}"
        failed=1
    fi

    # 5. Installer reachability
    if ! test_connectivity "$INSTALLER_URL" "Wallarm native node installer"; then
        err "Preflight failed: installer not reachable: $INSTALLER_URL"
        failed=1
    fi

    # 6. Wallarm cloud reachability (must have at least one reachable region)
    local us eu
    us=$(test_cloud_endpoints "US" "${US_DATA_NODES[@]}")
    eu=$(test_cloud_endpoints "EU" "${EU_DATA_NODES[@]}")
    US_CLOUD_REACHABLE="$us"
    EU_CLOUD_REACHABLE="$eu"
    if [[ "$us" != "true" && "$eu" != "true" ]]; then
        err "Preflight failed: no Wallarm cloud region reachable (US/EU)."
        failed=1
    fi

    # 7. Disk space (>= 2GB free on the nodes volume)
    local avail_kb
    avail_kb=$(df -k "$BASE_DIR" 2>/dev/null | awk 'NR==2 {print $4}' || true)
    if [[ -n "$avail_kb" ]] && (( avail_kb < 2097152 )); then
        err "Preflight failed: insufficient disk space on $BASE_DIR (need >= 2GB free)."
        failed=1
    fi

    # 8. Memory (>= 2GB recommended; warning only)
    if command_exists free; then
        local mem_mb
        mem_mb=$(free -m 2>/dev/null | awk '/Mem:/ {print $2}')
        if [[ -n "$mem_mb" ]] && (( mem_mb < 2048 )); then
            err "Warning: only ${mem_mb}MB RAM detected (2GB+ recommended)."
        fi
    fi

    if [[ $failed -ne 0 ]]; then
        err "Preflight check FAILED. Resolve the issues above and re-run."
        return 1
    fi

    log "Preflight checks passed."
    return 0
}

# Validate a listen address (IP:PORT) and check its port is free
check_listen_port() {
    local address="$1"
    local port="${address##*:}"
    if [[ ! "$port" =~ ^[0-9]+$ ]] || (( port < 1 || port > 65535 )); then
        err "Invalid listen address (expected IP:PORT): $address"
        return 1
    fi
    if ! check_port_available "$port"; then
        err "Listen port $port (for $address) is already in use."
        return 1
    fi
    return 0
}

# --- Core actions ---
install_single_node() {
    local node_name="$1"
    local listen_address="$2"
    local api_token="$3"
    local api_host="$4"
    local labels="${5:-group=${node_name}}"

    local work_dir="${NODES_DIR}/${node_name}"
    local installer_path="${work_dir}/wallarm-native-node-aio.sh"

    log "[${node_name}] Installing (listening on ${listen_address})..."

    mkdir -p "${work_dir}/etc" "${work_dir}/var/log" "${work_dir}/var/run"

    # 1. Write config
    cat > "${work_dir}/etc/go-node.yaml" <<EOF
mode: connector-server
connector:
  address: "${listen_address}"
EOF

    # 2. Download installer if missing
    if [[ ! -f "${installer_path}" ]]; then
        log "[${node_name}] Downloading installer..."
        curl -fsSL -o "${installer_path}" "${INSTALLER_URL}" || {
            err "[${node_name}] Download failed"
            return 1
        }
        chmod +x "${installer_path}"
    fi

    # 3. Write environment file (used by systemd and manual scripts)
    write_env_file "${node_name}" "${api_token}" "${api_host}" "${labels}"

    # 4. Run installer
    cd "${work_dir}" || return 1
    if WALLARM_API_TOKEN="${api_token}" \
       WALLARM_API_HOST="${api_host}" \
       WALLARM_LABELS="${labels}" \
       WALLARM_CONFIG_PATH="${work_dir}/etc/go-node.yaml" \
       ./wallarm-native-node-aio.sh install \
           -- --config-dir "${work_dir}/etc" \
           -- --log-dir "${work_dir}/var/log" \
           -- --pid-dir "${work_dir}/var/run" \
           > "${work_dir}/install.log" 2>&1; then
        log "[${node_name}] Installation successful."
        # Enable and start the systemd service
        systemctl enable "wallarm-node@${node_name}" 2>/dev/null || true
        systemctl start "wallarm-node@${node_name}"
        log "[${node_name}] Service started (systemctl status wallarm-node@${node_name})"
    else
        err "[${node_name}] Installation failed. Check ${work_dir}/install.log"
        return 1
    fi
}

cmd_preflight() {
    check_root
    echo ""
    if run_preflight; then
        log "Preflight passed - system ready for --install."
        exit 0
    else
        exit 1
    fi
}

cmd_install() {
    check_root
    if ! run_preflight; then
        exit 1
    fi
    ensure_dirs
    generate_systemd_template

    read -p "Enter Wallarm API Token (with Deploy role): " WALLARM_API_TOKEN
    if [[ -z "$WALLARM_API_TOKEN" ]]; then
        err "API Token cannot be empty."
        exit 1
    fi

    # Select Wallarm cloud region (US/EU)
    select_cloud_region

    echo ""
    echo "Enter each node's name and listening address (format: name IP:Port)"
    echo "Example: node1 0.0.0.0:8081"
    echo "Leave name blank to finish."

    declare -a NODE_NAMES=()
    declare -a NODE_ADDRESSES=()

    while true; do
        read -p "Node name (blank to stop): " name
        [[ -z "$name" ]] && break
        read -p "Listening address (e.g., 0.0.0.0:8081): " address
        if [[ -z "$address" ]]; then
            err "Address cannot be empty, skipping."
            continue
        fi
        NODE_NAMES+=("$name")
        NODE_ADDRESSES+=("$address")
    done

    if [[ ${#NODE_NAMES[@]} -eq 0 ]]; then
        err "No nodes provided."
        exit 1
    fi

    # Validate listen ports before installing anything
    local address
    for address in "${NODE_ADDRESSES[@]}"; do
        if ! check_listen_port "$address"; then
            exit 1
        fi
    done

    echo ""
    echo "Will install ${#NODE_NAMES[@]} nodes in parallel:"
    for i in "${!NODE_NAMES[@]}"; do
        echo "  - ${NODE_NAMES[$i]} -> ${NODE_ADDRESSES[$i]}"
    done
    read -p "Proceed? (y/N): " confirm
    [[ ! "$confirm" =~ ^[Yy]$ ]] && { echo "Cancelled."; exit 0; }

    echo ""
    log "Starting parallel installations..."

    declare -a INSTALL_PIDS=()
    for i in "${!NODE_NAMES[@]}"; do
        install_single_node "${NODE_NAMES[$i]}" "${NODE_ADDRESSES[$i]}" "$WALLARM_API_TOKEN" "$API_HOST" &
        INSTALL_PIDS+=($!)
    done

    FAILED=0
    local pid
    for pid in "${INSTALL_PIDS[@]}"; do
        wait "$pid" || ((FAILED++))
    done

    if [[ $FAILED -eq 0 ]]; then
        log "All nodes installed and started via systemd."
    else
        err "$FAILED node(s) failed. Check individual install.log files."
    fi
}

cmd_config() {
    # Usage: --config --node NAME --address IP:PORT [--token TOKEN] [--labels LABELS]
    check_root
    local node_name="" address="" token="" labels=""

    while [[ $# -gt 0 ]]; do
        case "$1" in
            --node)    [[ $# -ge 2 ]] || { err "--node requires a value"; exit 1; }; node_name="$2"; shift 2 ;;
            --address) [[ $# -ge 2 ]] || { err "--address requires a value"; exit 1; }; address="$2"; shift 2 ;;
            --token)   [[ $# -ge 2 ]] || { err "--token requires a value"; exit 1; }; token="$2"; shift 2 ;;
            --labels)  [[ $# -ge 2 ]] || { err "--labels requires a value"; exit 1; }; labels="$2"; shift 2 ;;
            *) err "Unknown config option: $1"; exit 1 ;;
        esac
    done

    if [[ -z "$node_name" ]]; then
        err "Missing --node"
        exit 1
    fi

    local work_dir="${NODES_DIR}/${node_name}"
    if [[ ! -d "$work_dir" ]]; then
        err "Node '$node_name' does not exist in ${NODES_DIR}"
        exit 1
    fi

    # Update config file
    if [[ -n "$address" ]]; then
        if ! check_listen_port "$address"; then
            exit 1
        fi
        log "Updating listening address to $address"
        sed -i "s|^\([[:space:]]*address: \).*|\1\"${address}\"|" "${work_dir}/etc/go-node.yaml"
    fi

    # Update env file if token or labels provided (rewrite to avoid sed escaping issues)
    if [[ -n "$token" || -n "$labels" ]]; then
        local env_file="${work_dir}/env"
        [[ -f "$env_file" ]] || { err "env file not found"; exit 1; }
        local current_token current_labels
        current_token=$(grep '^WALLARM_API_TOKEN=' "$env_file" | cut -d= -f2-)
        current_labels=$(grep '^WALLARM_LABELS=' "$env_file" | cut -d= -f2-)
        write_env_file "$node_name" "${token:-$current_token}" "${labels:-$current_labels}"
        log "Token/labels updated for $node_name."
    fi

    log "Configuration updated for $node_name. Restart with: systemctl restart wallarm-node@${node_name}"
}

cmd_remove() {
    check_root
    local node_name=""
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --node) [[ $# -ge 2 ]] || { err "--node requires a value"; exit 1; }; node_name="$2"; shift 2 ;;
            *) err "Unknown remove option: $1"; exit 1 ;;
        esac
    done

    if [[ -z "$node_name" ]]; then
        err "Missing --node"
        exit 1
    fi

    local work_dir="${NODES_DIR}/${node_name}"
    if [[ ! -d "$work_dir" ]]; then
        err "Node '$node_name' does not exist."
        exit 1
    fi

    log "Stopping and disabling service..."
    systemctl stop "wallarm-node@${node_name}" 2>/dev/null || true
    systemctl disable "wallarm-node@${node_name}" 2>/dev/null || true

    log "Removing directory ${work_dir}..."
    rm -rf "$work_dir"
    log "Node $node_name removed."
}

cmd_status() {
    # Show systemd status for all found nodes or a specific one
    local node_name="${1:-}"
    if [[ -n "$node_name" ]]; then
        systemctl status "wallarm-node@${node_name}" --no-pager
    else
        echo "Wallarm Nodes status:"
        local dir name
        for dir in "${NODES_DIR}"/*/; do
            if [[ -d "$dir" ]]; then
                name=$(basename "$dir")
                echo "--- $name ---"
                systemctl status "wallarm-node@${name}" --no-pager | head -5
                echo ""
            fi
        done
    fi
}

# --- Help ---
show_help() {
    cat <<EOF
Usage: $0 [COMMAND] [OPTIONS]

Commands:
  --preflight           Run preflight checks only (no installation).
  --install             Interactive installation of one or more nodes (parallel).
  --config              Update an existing node's configuration.
                        Options: --node NAME --address IP:PORT [--token TOKEN] [--labels LABELS]
  --remove              Remove a node completely.
                        Options: --node NAME
  --status [NODE]       Show systemd status for a node, or all nodes.

Environment:
  WALLARM_INSTALLER_URL   Override the all-in-one installer URL (default: repo.wallarm.com latest).
  WALLARM_INSTALLER_ARCH  Installer architecture suffix (default: x86_64).

Examples:
  $0 --preflight
  $0 --install
  $0 --config --node node1 --address 0.0.0.0:9090
  $0 --remove --node node2
  $0 --status
EOF
}

# --- Main argument parsing ---
if [[ $# -eq 0 ]]; then
    show_help
    exit 0
fi

case "$1" in
    --preflight) shift; cmd_preflight "$@" ;;
    --install)   shift; cmd_install "$@" ;;
    --config)    shift; cmd_config "$@" ;;
    --remove)    shift; cmd_remove "$@" ;;
    --status)    shift; cmd_status "$@" ;;
    --help|-h)   show_help ;;
    *) err "Unknown command: $1"; show_help; exit 1 ;;
esac