#!/bin/bash
#
# DART DDoS Shield - Linux Installer
#
# Controller install:
#   curl -sSL https://pkg.dev.snaju.com/edge/install.sh | sudo bash
#   or: sudo ./install.sh
#
# Sensor install (HTTP adoption — recommended):
#   curl -sSL https://pkg.dev.snaju.com/edge/install.sh | sudo bash -s -- --controller http://controller:8080 --token <secret>
#
# Sensor install (legacy base64 token):
#   curl -sSL https://pkg.dev.snaju.com/edge/install.sh | sudo bash -s -- --adopt <token>
#

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
BOLD='\033[1m'
NC='\033[0m' # No Color

# Download URL base
DOWNLOAD_BASE="https://pkg.dev.snaju.com/edge"

# Configuration
INSTALL_DIR="/opt/dart-shield"
CONFIG_DIR="/etc/dart-shield"
LOG_DIR="/var/log/dart-shield"
DATA_DIR="/var/lib/dart-shield"
USER="dart-shield"
GROUP="dart-shield"

# Set by install mode (controller or sensor)
INSTALL_MODE=""
BINARY_NAME=""
SERVICE_NAME=""

# Detect architecture
ARCH=$(uname -m)
case $ARCH in
    x86_64)
        BINARY_ARCH="amd64"
        ;;
    aarch64|arm64)
        BINARY_ARCH="arm64"
        ;;
    *)
        echo -e "${RED}Unsupported architecture: $ARCH${NC}"
        exit 1
        ;;
esac

OS=$(uname -s | tr '[:upper:]' '[:lower:]')

set_install_mode() {
    INSTALL_MODE="$1"
    if [ "$INSTALL_MODE" = "sensor" ]; then
        BINARY_NAME="dart-sensor"
        SERVICE_NAME="dart-sensor"
    else
        BINARY_NAME="dart-controller"
        SERVICE_NAME="dart-controller"
    fi
}

print_banner() {
    echo -e "${BLUE}"
    echo "╔═══════════════════════════════════════════════════════════╗"
    echo "║                                                           ║"
    echo "║             DART DDoS Shield Installer                    ║"
    echo "║                                                           ║"
    echo "║       https://pkg.dev.snaju.com/edge/                     ║"
    echo "║                                                           ║"
    echo "╚═══════════════════════════════════════════════════════════╝"
    echo -e "${NC}"
}

log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

section_header() {
    echo ""
    echo -e "${CYAN}${BOLD}── $1 ──${NC}"
    echo ""
}

# Prompt helper: prompt_value "prompt text" "default value"
# Sets REPLY to the entered value (or default if empty)
prompt_value() {
    local prompt="$1"
    local default="$2"
    if [ -n "$default" ]; then
        read -rp "  $prompt [$default]: " REPLY
        REPLY="${REPLY:-$default}"
    else
        read -rp "  $prompt: " REPLY
    fi
}

# Prompt helper for yes/no: prompt_yn "prompt text" "y" or "n"
# Sets REPLY to "true" or "false"
prompt_yn() {
    local prompt="$1"
    local default="$2"
    local hint="y/N"
    [ "$default" = "y" ] && hint="Y/n"
    read -rp "  $prompt [$hint]: " answer
    answer="${answer:-$default}"
    if [[ "$answer" =~ ^[Yy] ]]; then
        REPLY="true"
    else
        REPLY="false"
    fi
}

check_root() {
    if [ "$EUID" -ne 0 ]; then
        log_error "This script must be run as root (use sudo)"
        exit 1
    fi
}

check_systemd() {
    if ! command -v systemctl &> /dev/null; then
        log_error "systemd is required but not found"
        exit 1
    fi
}

create_user() {
    if ! id "$USER" &>/dev/null; then
        log_info "Creating system user: $USER"
        useradd --system --no-create-home --shell /bin/false "$USER"
    else
        log_info "User $USER already exists"
    fi
}

create_directories() {
    log_info "Creating directories..."
    mkdir -p "$INSTALL_DIR"
    mkdir -p "$INSTALL_DIR/staging"
    mkdir -p "$CONFIG_DIR"
    mkdir -p "$LOG_DIR"
    mkdir -p "$DATA_DIR"

    chown -R "$USER:$GROUP" "$LOG_DIR"
    chown -R "$USER:$GROUP" "$DATA_DIR"
    chown -R "$USER:$GROUP" "$INSTALL_DIR/staging"
}

download_or_copy() {
    local src_name="$1"
    local dest_path="$2"

    local BINARY_URL="${DOWNLOAD_BASE}/${src_name}-${OS}-${BINARY_ARCH}"

    # Always stage through a temp file to avoid ETXTBSY on the target
    local tmpfile
    tmpfile=$(mktemp /tmp/dart-download.XXXXXX)

    # Check if binary exists in current directory (local install) first
    if [ -f "./${src_name}-${OS}-${BINARY_ARCH}" ]; then
        log_info "Using local binary: ./${src_name}-${OS}-${BINARY_ARCH}"
        cp "./${src_name}-${OS}-${BINARY_ARCH}" "$tmpfile"
    elif [ -f "./${src_name}" ]; then
        log_info "Using local binary: ./${src_name}"
        cp "./${src_name}" "$tmpfile"
    else
        log_info "Downloading from: $BINARY_URL"

        local dl_ok=false
        if command -v wget &> /dev/null; then
            if wget -q "$BINARY_URL" -O "$tmpfile" < /dev/null; then
                dl_ok=true
            else
                log_warn "wget failed, trying curl..."
            fi
        fi

        if [ "$dl_ok" = false ] && command -v curl &> /dev/null; then
            if curl -fsSL "$BINARY_URL" -o "$tmpfile" < /dev/null; then
                dl_ok=true
            fi
        fi

        if [ "$dl_ok" = false ]; then
            rm -f "$tmpfile"
            log_error "Failed to download binary from $BINARY_URL"
            exit 1
        fi

        log_info "Download complete"
    fi

    mv -f "$tmpfile" "$dest_path"
}

# Stop the service and kill any process still holding the binary open
stop_existing_binary() {
    local binary_path="$1"

    # Stop systemd service if running
    if systemctl is-active --quiet "$SERVICE_NAME" 2>/dev/null; then
        log_info "Stopping running $SERVICE_NAME service..."
        systemctl stop "$SERVICE_NAME"
        sleep 1
    fi

    # Kill any remaining processes using the binary
    if [ -f "$binary_path" ] && command -v fuser &> /dev/null; then
        local pids
        pids=$(fuser "$binary_path" 2>/dev/null || true)
        if [ -n "$pids" ]; then
            log_warn "Killing processes still using $binary_path: $pids"
            fuser -k "$binary_path" 2>/dev/null || true
            sleep 1
        fi
    fi
}

install_binary() {
    log_info "Installing $BINARY_NAME..."

    stop_existing_binary "$INSTALL_DIR/$BINARY_NAME"

    # Download to temp file then move (atomic replace, avoids ETXTBSY)
    download_or_copy "$BINARY_NAME" "$INSTALL_DIR/$BINARY_NAME.new"

    chmod +x "$INSTALL_DIR/$BINARY_NAME.new"
    chown root:root "$INSTALL_DIR/$BINARY_NAME.new"
    mv "$INSTALL_DIR/$BINARY_NAME.new" "$INSTALL_DIR/$BINARY_NAME"

    # Verify binary works
    if ! "$INSTALL_DIR/$BINARY_NAME" --version &> /dev/null; then
        log_warn "Binary verification failed - this may be normal if dependencies are missing"
    else
        VERSION=$("$INSTALL_DIR/$BINARY_NAME" --version 2>&1 | head -n1)
        log_info "Installed: $VERSION"
    fi
}

create_symlink() {
    log_info "Creating '$BINARY_NAME' command symlink..."
    rm -f "/usr/local/bin/$BINARY_NAME"
    ln -s "$INSTALL_DIR/$BINARY_NAME" "/usr/local/bin/$BINARY_NAME"
    log_info "You can now use '$BINARY_NAME' command from anywhere"

    # Install shield CLI helper
    if [ -f "./shield.sh" ]; then
        cp ./shield.sh /usr/local/bin/shield
    else
        log_info "Downloading shield CLI helper..."
        wget -q -O /usr/local/bin/shield "${DOWNLOAD_BASE}/shield.sh" 2>/dev/null \
            || curl -sSL -o /usr/local/bin/shield "${DOWNLOAD_BASE}/shield.sh"
    fi
    chmod +x /usr/local/bin/shield
    log_info "Installed 'shield' CLI helper"
}

install_systemd_service() {
    log_info "Installing systemd service: $SERVICE_NAME..."

    cat > "/etc/systemd/system/${SERVICE_NAME}.service" << EOF
[Unit]
Description=DART DDoS Shield - ${INSTALL_MODE^}
Documentation=https://github.com/dantiez/dart-ddos-shield
After=network.target
Wants=network-online.target

[Service]
Type=simple
User=$USER
Group=$GROUP
ExecStart=$INSTALL_DIR/$BINARY_NAME --config $CONFIG_DIR/config.yaml
ExecReload=/bin/kill -HUP \$MAINPID
Restart=on-failure
RestartSec=5
StartLimitInterval=60
StartLimitBurst=3

AmbientCapabilities=CAP_NET_BIND_SERVICE
ProtectSystem=strict
ProtectHome=true
PrivateTmp=true
ReadWritePaths=$CONFIG_DIR $LOG_DIR $DATA_DIR

LimitNOFILE=65535
LimitNPROC=4096

Environment=GOMAXPROCS=4

[Install]
WantedBy=multi-user.target
EOF

    systemctl daemon-reload
}

setup_sudoers() {
    log_info "Configuring sudoers for auto-update..."
    SUDOERS_FILE="/etc/sudoers.d/dart-shield"
    echo "$USER ALL=(root) NOPASSWD: /bin/systemctl restart $SERVICE_NAME" > "$SUDOERS_FILE"
    chmod 440 "$SUDOERS_FILE"

    # Validate sudoers file
    if visudo -c -f "$SUDOERS_FILE" &>/dev/null; then
        log_info "Sudoers rule installed successfully"
    else
        log_warn "Sudoers validation failed, removing rule"
        rm -f "$SUDOERS_FILE"
    fi
}

# ──────────────────────────────────────────────────────────────
# Interactive configuration setup (controller only)
# ──────────────────────────────────────────────────────────────

setup_config() {
    echo ""
    echo -e "${BLUE}${BOLD}╔═══════════════════════════════════════════════════════════╗${NC}"
    echo -e "${BLUE}${BOLD}║           Controller Configuration Wizard                  ║${NC}"
    echo -e "${BLUE}${BOLD}╚═══════════════════════════════════════════════════════════╝${NC}"
    echo ""
    echo "This wizard will configure your DART controller."
    echo "Press Enter to accept the [default] value shown in brackets."
    echo ""

    # Try to auto-detect the server's primary IP
    local SERVER_IP
    SERVER_IP=$(hostname -I 2>/dev/null | awk '{print $1}')

    # ── Edge Identification ──
    section_header "Controller Identification"
    echo "  Give this controller a unique name and location info."
    echo ""

    prompt_value "Controller name (unique identifier)" "$(hostname -s)"
    local EDGE_NAME="$REPLY"

    prompt_value "Region" "US-SOUTH"
    local EDGE_REGION="$REPLY"

    prompt_value "Datacenter / POP name" "dc1"
    local EDGE_DATACENTER="$REPLY"

    prompt_value "Latitude (0 to skip)" "0"
    local EDGE_LAT="$REPLY"

    prompt_value "Longitude (0 to skip)" "0"
    local EDGE_LON="$REPLY"

    # ── Network Prefixes ──
    section_header "Network Prefixes"
    echo "  Enter the IP prefixes (CIDR) your sensors should protect."
    echo "  These are your announced network ranges."
    echo "  Enter one per line. Type 'done' when finished."
    echo ""

    local PREFIXES=()
    while true; do
        read -rp "  Prefix (CIDR, or 'done'): " prefix
        if [ "$prefix" = "done" ] || [ -z "$prefix" ]; then
            break
        fi
        PREFIXES+=("$prefix")
    done

    local AUTO_PROTECT_MODE="null-route"
    if [ ${#PREFIXES[@]} -gt 0 ]; then
        echo ""
        echo "  Auto-protect mode determines how new IPs in your prefixes are protected:"
        echo "    null-route = BGP blackhole announcements (default)"
        echo "    scrub      = BGP FlowSpec vector-specific scrubbing"
        prompt_value "Auto-protect mode" "null-route"
        AUTO_PROTECT_MODE="$REPLY"
    fi

    # ── Mesh Configuration ──
    section_header "Mesh Network"
    echo "  The mesh network allows sensors to connect to this controller."
    echo ""

    prompt_yn "Enable mesh networking?" "y"
    local MESH_ENABLED="$REPLY"

    local MESH_PORT="7946"
    local MESH_KEY=""
    if [ "$MESH_ENABLED" = "true" ]; then
        prompt_value "Mesh listen port" "7946"
        MESH_PORT="$REPLY"

        # Generate a random encryption key
        MESH_KEY=$(openssl rand -hex 16 2>/dev/null || head -c 32 /dev/urandom | xxd -p -c 32 2>/dev/null || echo "")
        if [ -n "$MESH_KEY" ]; then
            log_info "Generated mesh encryption key"
        else
            prompt_value "Mesh encryption key (hex string)" ""
            MESH_KEY="$REPLY"
        fi
    fi

    # ── API Server ──
    section_header "API & Web UI"

    prompt_value "API listen port" "8080"
    local API_PORT="$REPLY"

    # ── DartNode OAuth SSO ──
    section_header "DartNode OAuth SSO"
    echo "  OAuth enables 'Sign in with DartNode' on the web UI."
    echo "  Register an app at https://oauth.dartnode.com to get credentials."
    echo "  Leave blank to skip (can be configured later in config.yaml)."
    echo ""

    prompt_value "OAuth Client ID" ""
    local OAUTH_CLIENT_ID="$REPLY"

    local OAUTH_CLIENT_SECRET=""
    local OAUTH_CALLBACK_URL=""
    if [ -n "$OAUTH_CLIENT_ID" ]; then
        prompt_value "OAuth Client Secret" ""
        OAUTH_CLIENT_SECRET="$REPLY"

        local DEFAULT_CALLBACK="http://${SERVER_IP}:${API_PORT}/oauth/callback"
        prompt_value "OAuth Callback URL" "$DEFAULT_CALLBACK"
        OAUTH_CALLBACK_URL="$REPLY"
    fi

    # ── Logging ──
    section_header "Logging"

    echo "  Log levels: debug, info, warn, error"
    prompt_value "Log level" "info"
    local LOG_LEVEL="$REPLY"

    echo "  Log formats: console (human-readable), json (structured)"
    prompt_value "Log format" "console"
    local LOG_FORMAT="$REPLY"

    # ── Additional Features ──
    section_header "Additional Features"

    prompt_yn "Enable webhooks (attack notifications)?" "y"
    local WH_ENABLED="$REPLY"

    prompt_yn "Enable traffic signature classification?" "y"
    local SIG_ENABLED="$REPLY"

    # ── Build the YAML config ──
    section_header "Writing Configuration"

    # Build prefix YAML
    local PREFIX_YAML=""
    if [ ${#PREFIXES[@]} -gt 0 ]; then
        PREFIX_YAML="    prefixes:"
        for p in "${PREFIXES[@]}"; do
            PREFIX_YAML="${PREFIX_YAML}\n        - \"${p}\""
        done
        PREFIX_YAML="${PREFIX_YAML}\n    auto_protect_mode: \"${AUTO_PROTECT_MODE}\""
    else
        PREFIX_YAML="    prefixes: []\n    auto_protect_mode: \"null-route\""
    fi

    # Build mesh YAML
    local MESH_YAML=""
    if [ "$MESH_ENABLED" = "true" ]; then
        MESH_YAML="mesh:
    enabled: true
    bind_port: ${MESH_PORT}
    encryption_key: \"${MESH_KEY}\""
    else
        MESH_YAML="mesh:
    enabled: false"
    fi

    # Write config file
    log_info "Writing config to $CONFIG_DIR/config.yaml..."

    cat > "$CONFIG_DIR/config.yaml" << CONFIGEOF
# DART DDoS Shield - Controller Configuration
# Generated by installer on $(date -u +"%Y-%m-%dT%H:%M:%SZ")

setup_complete: true

edge:
    name: "${EDGE_NAME}"
    region: "${EDGE_REGION}"
    datacenter: "${EDGE_DATACENTER}"
    latitude: ${EDGE_LAT}
    longitude: ${EDGE_LON}

network:
$(echo -e "$PREFIX_YAML")

${MESH_YAML}

api:
    port: ${API_PORT}
    oauth:
        client_id: "${OAUTH_CLIENT_ID}"
        client_secret: "${OAUTH_CLIENT_SECRET}"
        callback_url: "${OAUTH_CALLBACK_URL}"

database:
    path: "${DATA_DIR}/dart-shield.db"

logging:
    level: "${LOG_LEVEL}"
    format: "${LOG_FORMAT}"
    file: ""

webhooks:
    enabled: ${WH_ENABLED}
    max_retries: 3
    signature_header: "X-Webhook-Signature"

signatures:
    enabled: ${SIG_ENABLED}

update:
    enabled: false
    check_interval: 30
    max_concurrent: 1
    staging_dir: "${INSTALL_DIR}/staging"
    service_name: "${SERVICE_NAME}"
    allow_downgrade: false
CONFIGEOF

    chown root:$GROUP "$CONFIG_DIR/config.yaml"
    chmod 660 "$CONFIG_DIR/config.yaml"

    log_info "Configuration written successfully!"
    echo ""

    # Show summary
    echo -e "${CYAN}${BOLD}── Configuration Summary ──${NC}"
    echo ""
    echo "  Controller:  ${EDGE_NAME} (${EDGE_REGION} / ${EDGE_DATACENTER})"
    if [ ${#PREFIXES[@]} -gt 0 ]; then
        echo "  Prefixes:    ${PREFIXES[*]}"
    else
        echo "  Prefixes:    (none configured)"
    fi
    echo "  Mesh:        ${MESH_ENABLED} (port ${MESH_PORT})"
    echo "  API:         port ${API_PORT}"
    echo "  Logging:     ${LOG_LEVEL} / ${LOG_FORMAT}"
    echo ""
}

enable_and_start_service() {
    echo ""
    log_info "Enabling and starting service..."
    systemctl enable "$SERVICE_NAME"
    systemctl start "$SERVICE_NAME"

    sleep 2

    if systemctl is-active --quiet "$SERVICE_NAME"; then
        log_info "Service started successfully!"
    else
        log_error "Service failed to start. Check logs with: journalctl -u $SERVICE_NAME -f"
    fi
}

update() {
    # Detect which binary is installed
    if [ -f "$INSTALL_DIR/dart-controller" ]; then
        set_install_mode "controller"
    elif [ -f "$INSTALL_DIR/dart-sensor" ]; then
        set_install_mode "sensor"
    else
        log_error "No DART binary found in $INSTALL_DIR"
        exit 1
    fi

    log_info "Updating $BINARY_NAME..."

    OLD_VERSION=""
    if [ -f "$INSTALL_DIR/$BINARY_NAME" ]; then
        OLD_VERSION=$("$INSTALL_DIR/$BINARY_NAME" --version 2>&1 | head -n1 || echo "unknown")
        log_info "Current version: $OLD_VERSION"
    fi

    # Stop the service if running
    if systemctl is-active --quiet "$SERVICE_NAME" 2>/dev/null; then
        log_info "Stopping service..."
        systemctl stop "$SERVICE_NAME"
        SERVICE_WAS_RUNNING=true
    else
        SERVICE_WAS_RUNNING=false
    fi

    # Backup current binary
    if [ -f "$INSTALL_DIR/$BINARY_NAME" ]; then
        log_info "Backing up current binary..."
        cp "$INSTALL_DIR/$BINARY_NAME" "$INSTALL_DIR/$BINARY_NAME.backup"
    fi

    # Download new binary
    download_or_copy "$BINARY_NAME" "$INSTALL_DIR/$BINARY_NAME.new"

    # Replace binary
    chmod +x "$INSTALL_DIR/$BINARY_NAME.new"
    mv "$INSTALL_DIR/$BINARY_NAME.new" "$INSTALL_DIR/$BINARY_NAME"
    chown root:root "$INSTALL_DIR/$BINARY_NAME"

    # Verify new binary
    if ! "$INSTALL_DIR/$BINARY_NAME" --version &> /dev/null; then
        log_error "New binary verification failed, rolling back..."
        if [ -f "$INSTALL_DIR/$BINARY_NAME.backup" ]; then
            mv "$INSTALL_DIR/$BINARY_NAME.backup" "$INSTALL_DIR/$BINARY_NAME"
        fi
        exit 1
    fi

    NEW_VERSION=$("$INSTALL_DIR/$BINARY_NAME" --version 2>&1 | head -n1)
    log_info "New version: $NEW_VERSION"

    rm -f "$INSTALL_DIR/$BINARY_NAME.backup"

    # Install/update shield CLI helper
    if [ -f "./shield.sh" ]; then
        cp ./shield.sh /usr/local/bin/shield
    else
        log_info "Downloading shield CLI helper..."
        wget -q -O /usr/local/bin/shield "${DOWNLOAD_BASE}/shield.sh" 2>/dev/null \
            || curl -sSL -o /usr/local/bin/shield "${DOWNLOAD_BASE}/shield.sh"
    fi
    chmod +x /usr/local/bin/shield
    log_info "Updated 'shield' CLI helper"

    # Restart service if it was running
    if [ "$SERVICE_WAS_RUNNING" = true ]; then
        log_info "Restarting service..."
        systemctl start "$SERVICE_NAME"

        sleep 2
        if systemctl is-active --quiet "$SERVICE_NAME"; then
            log_info "Service restarted successfully!"
        else
            log_error "Service failed to start. Check logs with: journalctl -u $SERVICE_NAME -f"
        fi
    fi

    echo ""
    echo -e "${GREEN}Update complete! $OLD_VERSION -> $NEW_VERSION${NC}"
    echo ""
}

print_completion() {
    local SERVER_IP
    SERVER_IP=$(hostname -I 2>/dev/null | awk '{print $1}')

    echo ""
    echo -e "${GREEN}═══════════════════════════════════════════════════════════${NC}"
    echo -e "${GREEN}         Controller Installation Complete!                 ${NC}"
    echo -e "${GREEN}═══════════════════════════════════════════════════════════${NC}"
    echo ""
    echo -e "${BOLD}Next step: open the setup wizard to configure your controller:${NC}"
    echo ""
    echo -e "  ${CYAN}http://${SERVER_IP}:8080/setup${NC}"
    echo ""
    echo "  Default login:  admin / admin"
    echo ""
    echo "To add sensors, generate an adoption token in the web UI"
    echo "and run this installer with --adopt on each sensor host."
    echo ""
    echo "Useful commands:"
    echo "  Status:      sudo systemctl status $SERVICE_NAME"
    echo "  Logs:        sudo journalctl -u $SERVICE_NAME -f"
    echo "  Restart:     sudo systemctl restart $SERVICE_NAME"
    echo "  Stop:        sudo systemctl stop $SERVICE_NAME"
    echo "  Reconfigure: sudo $0 --reconfigure"
    echo "  Config:      sudo nano $CONFIG_DIR/config.yaml"
    echo ""
    echo "Installation paths:"
    echo "  Binary:   $INSTALL_DIR/$BINARY_NAME"
    echo "  Config:   $CONFIG_DIR/config.yaml"
    echo "  Logs:     $LOG_DIR/"
    echo "  Data:     $DATA_DIR/"
    echo ""
}

print_usage() {
    echo "Usage: $0 [OPTION]"
    echo ""
    echo "Options:"
    echo "  (none)                                Install dart-controller"
    echo "  --controller <url> --token <secret> [--name <name>]  Install dart-sensor via HTTP adoption (recommended)"
    echo "  --adopt <token> [--name <name>] [--api-url <url>]   Install dart-sensor via base64 token (legacy)"
    echo "  --reconfigure                         Re-run the controller configuration wizard"
    echo "  --update                              Update the installed binary to the latest version"
    echo "  --uninstall                           Uninstall DART DDoS Shield"
    echo "  --help                                Show this help message"
    echo ""
}

# Main
main() {
    print_banner

    case "$1" in
        --controller)
            check_root
            check_systemd
            set_install_mode "sensor"

            CONTROLLER_URL="$2"
            if [ -z "$CONTROLLER_URL" ]; then
                log_error "--controller requires a URL argument (e.g., http://controller:8080)"
                print_usage
                exit 1
            fi

            TOKEN_SECRET=""
            CUSTOM_NAME=""
            shift 2
            while [ $# -gt 0 ]; do
                case "$1" in
                    --token)
                        TOKEN_SECRET="$2"
                        shift 2
                        ;;
                    --name)
                        CUSTOM_NAME="$2"
                        shift 2
                        ;;
                    *)
                        log_error "Unknown option: $1"
                        print_usage
                        exit 1
                        ;;
                esac
            done

            if [ -z "$TOKEN_SECRET" ]; then
                log_error "--token is required with --controller"
                print_usage
                exit 1
            fi

            log_info "Installing dart-sensor via HTTP adoption..."
            log_info "Controller: $CONTROLLER_URL"

            # Install system
            create_user
            create_directories
            install_binary
            create_symlink
            install_systemd_service
            setup_sudoers

            # Generate sensor config — minimal, HTTP adoption handles the rest
            SENSOR_NAME="${CUSTOM_NAME:-sensor-$(hostname -s)}"
            log_info "Generating sensor config for: $SENSOR_NAME"

            cat > "$CONFIG_DIR/config.yaml" << HTTPEOF
# DART DDoS Shield - Sensor Configuration
# HTTP adoption from $CONTROLLER_URL on $(date -u +"%Y-%m-%dT%H:%M:%SZ")

setup_complete: true

edge:
    name: "${SENSOR_NAME}"

mesh:
    controller_url: "${CONTROLLER_URL}"
    adoption_secret: "${TOKEN_SECRET}"
    config_poll_sec: 60

database:
    path: "${DATA_DIR}/dart-shield.db"

logging:
    level: "info"
    format: "console"
HTTPEOF

            chown root:$GROUP "$CONFIG_DIR/config.yaml"
            chmod 660 "$CONFIG_DIR/config.yaml"

            log_info "Sensor config written successfully"

            enable_and_start_service

            echo ""
            echo -e "${GREEN}═══════════════════════════════════════════════════════════${NC}"
            echo -e "${GREEN}              Sensor Adoption Started!                     ${NC}"
            echo -e "${GREEN}═══════════════════════════════════════════════════════════${NC}"
            echo ""
            echo "  Sensor name:  ${SENSOR_NAME}"
            echo "  Controller:   ${CONTROLLER_URL}"
            echo ""
            echo "  The sensor will contact the controller via HTTP, receive"
            echo "  mesh configuration, and join the mesh automatically."
            echo ""
            echo "  Check status:  sudo journalctl -u $SERVICE_NAME -f"
            echo ""
            echo "Installation paths:"
            echo "  Binary:   $INSTALL_DIR/$BINARY_NAME"
            echo "  Config:   $CONFIG_DIR/config.yaml"
            echo "  Data:     $DATA_DIR/"
            echo ""
            exit 0
            ;;
        --adopt)
            check_root
            check_systemd
            set_install_mode "sensor"

            if [ -z "$2" ]; then
                log_error "--adopt requires a token argument"
                print_usage
                exit 1
            fi

            ADOPT_TOKEN="$2"
            API_URL=""
            CUSTOM_NAME=""

            # Parse optional flags
            shift 2
            while [ $# -gt 0 ]; do
                case "$1" in
                    --api-url)
                        API_URL="$2"
                        shift 2
                        ;;
                    --name)
                        CUSTOM_NAME="$2"
                        shift 2
                        ;;
                    *)
                        log_error "Unknown option: $1"
                        print_usage
                        exit 1
                        ;;
                esac
            done

            log_info "Installing dart-sensor and adopting from controller..."

            # Decode the base64 token
            DECODED=$(echo "$ADOPT_TOKEN" | base64 -d 2>/dev/null)
            if [ -z "$DECODED" ]; then
                log_error "Failed to decode adoption token (invalid base64)"
                exit 1
            fi

            # Extract fields using python3 (widely available) or jq
            if command -v python3 &>/dev/null; then
                CONTROLLER_ADDR=$(echo "$DECODED" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('c',''))")
                ENCRYPTION_KEY=$(echo "$DECODED" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('k',''))")
                ADOPT_SECRET=$(echo "$DECODED" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('s',''))")
                TOKEN_API_PORT=$(echo "$DECODED" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('p',8080))")
                TOKEN_API_URL=$(echo "$DECODED" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('a',''))")
            elif command -v jq &>/dev/null; then
                CONTROLLER_ADDR=$(echo "$DECODED" | jq -r '.c // ""')
                ENCRYPTION_KEY=$(echo "$DECODED" | jq -r '.k // ""')
                ADOPT_SECRET=$(echo "$DECODED" | jq -r '.s // ""')
                TOKEN_API_PORT=$(echo "$DECODED" | jq -r '.p // 8080')
                TOKEN_API_URL=$(echo "$DECODED" | jq -r '.a // ""')
            else
                log_error "Either python3 or jq is required to decode the adoption token"
                exit 1
            fi

            if [ -z "$CONTROLLER_ADDR" ] || [ -z "$ADOPT_SECRET" ]; then
                log_error "Invalid adoption token (missing controller address or secret)"
                exit 1
            fi

            # Use CLI --api-url override, then token's API URL
            if [ -z "$API_URL" ]; then
                API_URL="$TOKEN_API_URL"
            fi

            # Extract controller host (strip port if present)
            CONTROLLER_HOST=$(echo "$CONTROLLER_ADDR" | cut -d: -f1)

            log_info "Controller: $CONTROLLER_ADDR"
            log_info "API URL: ${API_URL:-not set}"

            # Install system
            create_user
            create_directories
            install_binary
            create_symlink
            install_systemd_service
            setup_sudoers

            # Generate sensor config
            SENSOR_NAME="${CUSTOM_NAME:-sensor-$(hostname -s)}"
            log_info "Generating sensor config for: $SENSOR_NAME"

            cat > "$CONFIG_DIR/config.yaml" << ADOPTEOF
# DART DDoS Shield - Sensor Configuration
# Adopted from controller at $CONTROLLER_ADDR on $(date -u +"%Y-%m-%dT%H:%M:%SZ")

setup_complete: true

edge:
    name: "${SENSOR_NAME}"

mesh:
    enabled: true
    seed_nodes:
        - "${CONTROLLER_ADDR}"
    encryption_key: "${ENCRYPTION_KEY}"
    controller_addr: "${CONTROLLER_HOST}"
    adoption_secret: "${ADOPT_SECRET}"
    config_poll_sec: 60

api:
    port: ${TOKEN_API_PORT:-8080}
    url: "${API_URL}"

database:
    path: "${DATA_DIR}/dart-shield.db"

logging:
    level: "info"
    format: "console"
ADOPTEOF

            chown root:$GROUP "$CONFIG_DIR/config.yaml"
            chmod 660 "$CONFIG_DIR/config.yaml"

            log_info "Sensor config written successfully"

            enable_and_start_service

            echo ""
            echo -e "${GREEN}═══════════════════════════════════════════════════════════${NC}"
            echo -e "${GREEN}              Sensor Adoption Started!                     ${NC}"
            echo -e "${GREEN}═══════════════════════════════════════════════════════════${NC}"
            echo ""
            echo "  Sensor name:  ${SENSOR_NAME}"
            echo "  Controller:   ${CONTROLLER_ADDR}"
            if [ -n "$API_URL" ]; then
                echo "  API URL:      ${API_URL}"
            fi
            echo ""
            echo "  The sensor will connect to the controller mesh and"
            echo "  complete adoption automatically."
            echo ""
            echo "  Check status:  sudo journalctl -u $SERVICE_NAME -f"
            echo ""
            echo "Installation paths:"
            echo "  Binary:   $INSTALL_DIR/$BINARY_NAME"
            echo "  Config:   $CONFIG_DIR/config.yaml"
            echo "  Logs:     $LOG_DIR/"
            echo "  Data:     $DATA_DIR/"
            echo ""
            exit 0
            ;;
        --uninstall|-u)
            check_root
            log_warn "Uninstalling DART DDoS Shield..."

            # Stop and remove both possible services
            for svc in dart-controller dart-sensor; do
                systemctl stop "$svc" 2>/dev/null || true
                systemctl disable "$svc" 2>/dev/null || true
                rm -f "/etc/systemd/system/${svc}.service"
            done
            systemctl daemon-reload

            rm -rf "$INSTALL_DIR"
            rm -f /usr/local/bin/dart-controller
            rm -f /usr/local/bin/dart-sensor
            rm -f /usr/local/bin/shield

            if [ -t 0 ]; then
                read -p "Remove configuration and data? [y/N]: " remove_data
                if [[ "$remove_data" =~ ^[Yy]$ ]]; then
                    rm -rf "$CONFIG_DIR"
                    rm -rf "$DATA_DIR"
                    rm -rf "$LOG_DIR"
                fi

                read -p "Remove system user? [y/N]: " remove_user
                if [[ "$remove_user" =~ ^[Yy]$ ]]; then
                    userdel "$USER" 2>/dev/null || true
                fi
            fi

            rm -f /etc/sudoers.d/dart-shield

            log_info "Uninstallation complete"
            exit 0
            ;;
        --update)
            check_root
            update
            exit 0
            ;;
        --reconfigure)
            check_root
            set_install_mode "controller"
            if ! [ -t 0 ]; then
                log_error "--reconfigure requires an interactive terminal"
                exit 1
            fi

            # Stop service if running
            if systemctl is-active --quiet "$SERVICE_NAME" 2>/dev/null; then
                log_info "Stopping service for reconfiguration..."
                systemctl stop "$SERVICE_NAME"
            fi

            # Backup existing config
            if [ -f "$CONFIG_DIR/config.yaml" ]; then
                local BACKUP="$CONFIG_DIR/config.yaml.bak.$(date +%s)"
                cp "$CONFIG_DIR/config.yaml" "$BACKUP"
                log_info "Backed up existing config to $BACKUP"
            fi

            setup_config
            enable_and_start_service
            print_completion
            exit 0
            ;;
        --help|-h)
            print_usage
            exit 0
            ;;
        "")
            # No arguments, proceed with controller install
            ;;
        *)
            log_error "Unknown option: $1"
            print_usage
            exit 1
            ;;
    esac

    # Default: install controller
    set_install_mode "controller"

    check_root
    check_systemd

    log_info "Installing dart-controller for $OS/$BINARY_ARCH..."
    echo ""

    create_user
    create_directories
    install_binary
    create_symlink
    install_systemd_service
    setup_sudoers

    # Generate default config and direct user to web UI for setup
    if [ ! -f "$CONFIG_DIR/config.yaml" ]; then
        log_info "Generating default config (use web UI to complete setup)..."
        cat > "$CONFIG_DIR/config.yaml" << DEFAULTEOF
# DART DDoS Shield - Controller Configuration
# Complete setup at http://$(hostname -I 2>/dev/null | awk '{print $1}'):8080/setup

setup_complete: false

edge:
    name: "$(hostname -s)"
    region: "default"
    datacenter: "dc1"

api:
    port: 8080

database:
    path: "${DATA_DIR}/dart-shield.db"

logging:
    level: "info"
    format: "console"
DEFAULTEOF
        chown root:$GROUP "$CONFIG_DIR/config.yaml"
        chmod 660 "$CONFIG_DIR/config.yaml"
    fi

    enable_and_start_service
    print_completion
}

main "$@"