#!/bin/bash
#
# DART DDoS Shield - CLI helper
#
# Usage: shield <command> [args...]
#

set -e

INSTALL_DIR="/opt/dart-shield"
CONFIG_DIR="/etc/dart-shield"
CONFIG_FILE="$CONFIG_DIR/config.yaml"
DOWNLOAD_BASE="https://pkg.dev.snaju.com/edge"

# Auto-detect which component is installed
if [ -f "$INSTALL_DIR/dart-controller" ]; then
    BINARY_NAME="dart-controller"
    SERVICE_NAME="dart-controller"
elif [ -f "$INSTALL_DIR/dart-sensor" ]; then
    BINARY_NAME="dart-sensor"
    SERVICE_NAME="dart-sensor"
else
    echo "Error: no DART binary found in $INSTALL_DIR" >&2
    exit 1
fi

# Read a value from a YAML section: get_yaml_value <file> <section> <key>
get_yaml_value() {
    local file="$1" section="$2" key="$3"
    local in_section=false
    while IFS= read -r line; do
        case "$line" in
            "${section}:"*) in_section=true ;;
            [a-z]*:*) [ "$in_section" = true ] && in_section=false ;;
        esac
        if [ "$in_section" = true ]; then
            case "$line" in
                *"${key}:"*)
                    echo "$line" | sed "s/.*${key}:[[:space:]]*//" | tr -d '"' | tr -d "'"
                    return 0
                    ;;
            esac
        fi
    done < "$file"
    return 1
}

# Parse API port from config, default 8080
get_api_port() {
    local port="8080"
    if [ -f "$CONFIG_FILE" ]; then
        # Look for "port:" under "api:" section
        local in_api=false
        while IFS= read -r line; do
            case "$line" in
                api:*) in_api=true ;;
                [a-z]*:*) [ "$in_api" = true ] && in_api=false ;;
            esac
            if [ "$in_api" = true ]; then
                case "$line" in
                    *port:*)
                        port=$(echo "$line" | sed 's/.*port:[[:space:]]*//' | tr -d '"' | tr -d "'")
                        break
                        ;;
                esac
            fi
        done < "$CONFIG_FILE"
    fi
    echo "$port"
}

case "${1:-help}" in
    start)
        sudo systemctl start "$SERVICE_NAME"
        ;;
    stop)
        sudo systemctl stop "$SERVICE_NAME"
        ;;
    restart)
        sudo systemctl restart "$SERVICE_NAME"
        ;;
    status)
        systemctl status "$SERVICE_NAME"
        ;;
    logs)
        shift
        sudo journalctl -u "$SERVICE_NAME" -f "$@"
        ;;
    config)
        sudo "${EDITOR:-nano}" "$CONFIG_FILE"
        ;;
    version)
        "$INSTALL_DIR/$BINARY_NAME" --version
        ;;
    update)
        echo "Downloading installer..."
        curl -sSL "$DOWNLOAD_BASE/install.sh" | sudo bash -s -- --update
        ;;
    health)
        port=$(get_api_port)
        curl -s "http://localhost:${port}/health"
        echo
        ;;
    deploy-sensor)
        # --- Guards ---
        if [ "$(id -u)" -ne 0 ]; then
            echo "Error: deploy-sensor must run as root (use sudo)" >&2
            exit 1
        fi
        if [ ! -f "$INSTALL_DIR/dart-controller" ]; then
            echo "Error: dart-controller not found in $INSTALL_DIR — this command must run on a controller host" >&2
            exit 1
        fi
        if systemctl is-active --quiet dart-sensor 2>/dev/null; then
            echo "Error: dart-sensor service is already running" >&2
            exit 1
        fi

        # --- Generate adoption token via controller API ---
        API_PORT=$(get_api_port)
        API_ADDR="http://127.0.0.1:${API_PORT}"

        echo "Generating adoption token..."
        TOKEN_RESP=$(curl -sS -X POST "${API_ADDR}/api/v1/sensors/adopt-token" \
            -H "Content-Type: application/json" \
            -d "{\"label\":\"local-sensor\",\"ttl_hours\":1,\"api_addr\":\"$API_ADDR\"}" 2>&1)
        CURL_EXIT=$?
        if [ $CURL_EXIT -ne 0 ]; then
            echo "Error: failed to contact controller API at ${API_ADDR} (exit $CURL_EXIT)" >&2
            echo "Response: $TOKEN_RESP" >&2
            echo "Is dart-controller running? Check: systemctl status dart-controller" >&2
            exit 1
        fi

        # Parse the secret from the API response
        SECRET=$(echo "$TOKEN_RESP" | grep -o '"secret":"[^"]*"' | cut -d'"' -f4)
        if [ -z "$SECRET" ]; then
            echo "Error: could not parse adoption token from API response" >&2
            echo "Response: $TOKEN_RESP" >&2
            exit 1
        fi
        echo "Adoption token created."

        # --- Download sensor binary ---
        ARCH=$(uname -m)
        case "$ARCH" in
            x86_64)  ARCH="amd64" ;;
            aarch64|arm64) ARCH="arm64" ;;
            *) echo "Error: unsupported architecture: $ARCH" >&2; exit 1 ;;
        esac
        SENSOR_URL="$DOWNLOAD_BASE/dart-sensor-linux-$ARCH"
        SENSOR_BIN="$INSTALL_DIR/dart-sensor"
        echo "Downloading dart-sensor ($ARCH)..."
        if command -v wget >/dev/null 2>&1; then
            wget -q -O /tmp/dart-sensor "$SENSOR_URL"
        elif command -v curl >/dev/null 2>&1; then
            curl -sSL -o /tmp/dart-sensor "$SENSOR_URL"
        else
            echo "Error: wget or curl required" >&2; exit 1
        fi
        mv /tmp/dart-sensor "$SENSOR_BIN"
        chmod +x "$SENSOR_BIN"

        # --- Write sensor config ---
        # Minimal config: sensor uses HTTP adoption on first start to get mesh config
        SENSOR_CONFIG="$CONFIG_DIR/sensor.yaml"
        SENSOR_NAME="sensor-$(hostname -s)"
        cat > "$SENSOR_CONFIG" <<SEOF
setup_complete: true
edge:
    name: "$SENSOR_NAME"
mesh:
    controller_url: "$API_ADDR"
    adoption_secret: "$SECRET"
    config_poll_sec: 60
api:
    port: 0
database:
    path: "/var/lib/dart-shield/sensor.db"
logging:
    level: "info"
    format: "console"
SEOF
        echo "Sensor config written to $SENSOR_CONFIG"

        # --- Create systemd service ---
        cat > /etc/systemd/system/dart-sensor.service <<UEOF
[Unit]
Description=DART DDoS Shield Sensor
After=network-online.target
Wants=network-online.target

[Service]
Type=simple
ExecStart=$SENSOR_BIN --config $SENSOR_CONFIG
Restart=on-failure
RestartSec=5
LimitNOFILE=65535
ReadWritePaths=$CONFIG_DIR /var/lib/dart-shield

[Install]
WantedBy=multi-user.target
UEOF

        # --- Start sensor ---
        systemctl daemon-reload
        systemctl enable dart-sensor
        systemctl start dart-sensor
        echo "dart-sensor service started."

        # --- Verify ---
        sleep 5
        if systemctl is-active --quiet dart-sensor; then
            echo ""
            echo "Local sensor '$SENSOR_NAME' deployed successfully!"
            echo "It should appear on the dashboard within ~60 seconds."
        else
            echo "Warning: dart-sensor started but may not be healthy. Check: journalctl -u dart-sensor" >&2
            exit 1
        fi
        ;;
    uninstall)
        echo "Downloading installer..."
        curl -sSL "$DOWNLOAD_BASE/install.sh" | sudo bash -s -- --uninstall
        ;;
    help|--help|-h)
        echo "DART DDoS Shield CLI — manages the local $BINARY_NAME service"
        echo ""
        echo "Usage: shield <command> [args...]"
        echo ""
        echo "Commands:"
        echo "  start        Start the service"
        echo "  stop         Stop the service"
        echo "  restart      Restart the service"
        echo "  status       Show service status"
        echo "  logs [args]  Follow logs (extra journalctl flags passed through)"
        echo "  config       Open config.yaml in \$EDITOR (default: nano)"
        echo "  version      Show installed binary version"
        echo "  update       Download and run the latest installer with --update"
        echo "  health       Query the local /health endpoint"
        echo "  deploy-sensor Deploy a local sensor on this controller host (requires sudo)"
        echo "  uninstall    Download and run the installer with --uninstall"
        echo "  help         Show this help"
        ;;
    *)
        echo "Unknown command: $1" >&2
        echo "Run 'shield help' for usage." >&2
        exit 1
        ;;
esac