#!/bin/bash

# ==============================================================================
# Configuration
# ==============================================================================

# Default list of servers to check (SSH host aliases from ~/.ssh/config)
DEFAULT_SERVERS=("thor" "loki" "wanda")

# Remote paths and command
REMOTE_DIR="/root/.local/dpi-ch"
# dpich uses Bubbletea (TUI), which requires a TTY.
# When running in REMOTE_DIR, dpich automatically picks up config.yaml by default.
CMD="cd $REMOTE_DIR && if [ -t 0 ]; then ./dpich --all; else script -q -e -c './dpich --all' /dev/null; fi"

# Minimum free disk space required on remote server (in MB)
MIN_FREE_DISK_MB=100

# ==============================================================================
# Parse options and arguments
# ==============================================================================

PARALLEL=false
CLI_SERVERS=()

for arg in "$@"; do
    case "$arg" in
        -p|--parallel)
            PARALLEL=true
            ;;
        -s|--seq|--sequential)
            PARALLEL=false
            ;;
        *)
            CLI_SERVERS+=("$arg")
            ;;
    esac
done

# Use CLI servers if passed, otherwise fallback to DEFAULT_SERVERS
if [ "${#CLI_SERVERS[@]}" -gt 0 ]; then
    SERVERS=("${CLI_SERVERS[@]}")
else
    SERVERS=("${DEFAULT_SERVERS[@]}")
fi

# Locate template config locally
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
DEST_DIR="${SCRIPT_DIR}"

LOCAL_CONFIG="$SCRIPT_DIR/config.yaml"
if [ ! -f "$LOCAL_CONFIG" ]; then
    LOCAL_CONFIG="$SCRIPT_DIR/config.yml"
fi
if [ ! -f "$LOCAL_CONFIG" ]; then
    LOCAL_CONFIG="$SCRIPT_DIR/config.example.yml"
fi

mkdir -p "$DEST_DIR"

# ------------------------------------------------------------------------------
# Helper functions
# ------------------------------------------------------------------------------

get_config_version() {
    local file="$1"
    grep -m 1 -E '^# *(version|config_version|rev|revision):' "$file" 2>/dev/null \
        | sed -E 's/^# *(version|config_version|rev|revision): *//' \
        | sed -E 's/[[:space:]]+$//' \
        | tr -d '\r'
}

get_remote_config_version() {
    local server="$1"
    ssh "$server" "if [ -f '$REMOTE_DIR/config.yaml' ]; then grep -m 1 -E '^# *(version|config_version|rev|revision):' '$REMOTE_DIR/config.yaml'; elif [ -f '$REMOTE_DIR/config.yml' ]; then grep -m 1 -E '^# *(version|config_version|rev|revision):' '$REMOTE_DIR/config.yml'; fi" 2>/dev/null \
        | sed -E 's/^# *(version|config_version|rev|revision): *//' \
        | sed -E 's/[[:space:]]+$//' \
        | tr -d '\r'
}

validate_report() {
    local file="$1"
    if [ ! -s "$file" ]; then
        echo "[-] Error: Downloaded report is empty ($file)" >&2
        return 1
    fi

    if command -v python3 >/dev/null 2>&1; then
        if ! python3 -c "
import sys
try:
    import yaml
    with open('$file', 'r', encoding='utf-8') as f:
        data = yaml.safe_load(f)
    if not isinstance(data, (dict, list)) or not data:
        sys.exit(1)
except ImportError:
    import json
    with open('$file', 'r', encoding='utf-8') as f:
        data = json.load(f)
    if not isinstance(data, (dict, list)) or not data:
        sys.exit(1)
except Exception:
    sys.exit(1)
" 2>/dev/null; then
            echo "[-] Error: Downloaded report failed validation (corrupted YAML/JSON): $file" >&2
            return 1
        fi
    fi
    return 0
}

preflight_check() {
    local server="$1"

    # 1. Fail-fast SSH connectivity check
    echo "Testing SSH connection to $server..."
    if ! ssh -o BatchMode=yes -o ConnectTimeout=5 "$server" "true" 2>/dev/null; then
        echo "[-] Error: Cannot connect to $server via SSH (timeout or authentication error)." >&2
        return 1
    fi

    # 2. Check free disk space
    local free_mb
    free_mb=$(ssh "$server" "df -m $REMOTE_DIR 2>/dev/null || df -m /" | awk 'NR>1 {print $4}' | head -n 1)
    if [ -n "$free_mb" ] && [ "$free_mb" -lt "$MIN_FREE_DISK_MB" ] 2>/dev/null; then
        echo "[-] Error: Low disk space on $server (${free_mb}MB available in $REMOTE_DIR). Minimum ${MIN_FREE_DISK_MB}MB required." >&2
        return 1
    fi

    # 3. Check if dpich is installed
    if ! ssh "$server" "test -x $REMOTE_DIR/dpich || command -v dpich >/dev/null 2>&1"; then
        echo "[-] dpich is not installed on $server."
        read -r -p "Install dpich from GitHub releases on $server? [Y/n] " answer
        case "$answer" in
            [nN]|[nN][oO])
                echo "Skipping $server."
                return 1
                ;;
            *)
                echo "Installing dpich on $server..."
                ssh -t "$server" "which curl >/dev/null 2>&1 || (apt-get update && apt-get install -y curl || yum install -y curl); which unzip >/dev/null 2>&1 || (apt-get update && apt-get install -y unzip || yum install -y unzip); bash <(curl -Ls https://hyperion-cs.github.io/dpi-checkers/ru/dpi-ch/install/unix.sh)"
                if ! ssh "$server" "test -x $REMOTE_DIR/dpich || command -v dpich >/dev/null 2>&1"; then
                    echo "[-] Error: dpich installation failed on $server" >&2
                    return 1
                fi
                echo "[+] dpich installed successfully on $server."
                ;;
        esac
    fi

    # 4. Check and sync config
    local local_ver remote_ver
    local_ver=$(get_config_version "$LOCAL_CONFIG")
    remote_ver=$(get_remote_config_version "$server")

    local needs_upload=false

    if ! ssh "$server" "test -f $REMOTE_DIR/config.yaml || test -f $REMOTE_DIR/config.yml" 2>/dev/null; then
        echo "Config not found on $server. Uploading $(basename "$LOCAL_CONFIG")..."
        needs_upload=true
    elif [ -n "$local_ver" ] && [ "$local_ver" != "$remote_ver" ]; then
        echo "Config version changed (local: '$local_ver' vs remote: '${remote_ver:-none}'). Updating on $server..."
        needs_upload=true
    else
        # Fallback to checksum comparison if no version header is defined
        if [ -z "$local_ver" ]; then
            local local_md5 remote_md5
            local_md5=$(md5 -q "$LOCAL_CONFIG" 2>/dev/null || md5sum "$LOCAL_CONFIG" 2>/dev/null | awk '{print $1}')
            remote_md5=$(ssh "$server" "md5sum $REMOTE_DIR/config.yaml 2>/dev/null || md5sum $REMOTE_DIR/config.yml 2>/dev/null" | awk '{print $1}')
            if [ -n "$local_md5" ] && [ "$local_md5" != "$remote_md5" ]; then
                echo "Config content modified (checksum mismatch). Updating on $server..."
                needs_upload=true
            fi
        fi
    fi

    if [ "$needs_upload" = true ]; then
        if [ -f "$LOCAL_CONFIG" ]; then
            ssh "$server" "mkdir -p $REMOTE_DIR && rm -f $REMOTE_DIR/config.yaml $REMOTE_DIR/config.yml"
            if scp "$LOCAL_CONFIG" "$server:$REMOTE_DIR/config.yaml"; then
                echo "[+] Uploaded config.yaml to $server"
            else
                echo "[-] Warning: Failed to upload config to $server" >&2
            fi
        else
            echo "[-] Warning: No local config template found to upload" >&2
        fi
    else
        # Ensure config.yaml symlink exists if only config.yml is present
        ssh "$server" "test -f $REMOTE_DIR/config.yaml || (test -f $REMOTE_DIR/config.yml && ln -sf config.yml $REMOTE_DIR/config.yaml)"
    fi

    # 5. Clean up previous reports on server
    echo "Cleaning up previous reports on $server..."
    ssh "$server" "rm -f $REMOTE_DIR/dpi-ch-result_* $REMOTE_DIR/results_*"

    return 0
}

download_report() {
    local server="$1"

    echo "Finding result file on $server..."
    local remote_file
    remote_file=$(ssh "$server" "ls -t $REMOTE_DIR/dpi-ch-result_* $REMOTE_DIR/results_* 2>/dev/null | head -n 1")

    if [ -z "$remote_file" ]; then
        echo "[-] Warning: No result file found in $REMOTE_DIR on $server" >&2
        return 1
    fi

    local base_name
    base_name=$(basename "$remote_file")
    local clean_name="${base_name#dpi-ch-result_}"
    clean_name="${clean_name#results_}"
    local local_file="$DEST_DIR/dpi-ch-result_${server}_${clean_name}"

    echo "Downloading $remote_file -> $local_file"
    if scp "$server:$remote_file" "$local_file"; then
        if ! validate_report "$local_file"; then
            rm -f "$local_file"
            return 1
        fi
        echo "[+] Successfully saved and verified: $local_file"

        # Remove previous local reports for this server, keeping only the fresh one
        for old in "$DEST_DIR"/dpi-ch-result_"${server}"_*; do
            if [ -f "$old" ] && [ "$old" != "$local_file" ]; then
                echo "Removing previous local report: $(basename "$old")"
                rm -f "$old"
            fi
        done
        return 0
    else
        echo "[-] Error: Failed to download file from $server" >&2
        return 1
    fi
}

# ------------------------------------------------------------------------------
# Main Execution
# ------------------------------------------------------------------------------

if [ "$PARALLEL" = true ]; then
    echo "Starting DPI check in PARALLEL mode on servers: ${SERVERS[*]}"
    echo ""

    # Phase 1: Pre-flight checks on all servers
    ACTIVE_SERVERS=()
    for server in "${SERVERS[@]}"; do
        echo "--- Pre-flight check: $server ---"
        if preflight_check "$server"; then
            ACTIVE_SERVERS+=("$server")
        else
            echo "[-] Excluding $server from this run."
        fi
        echo ""
    done

    if [ "${#ACTIVE_SERVERS[@]}" -eq 0 ]; then
        echo "No servers available to test."
        exit 1
    fi

    echo "=================================================="
    echo "Launching checks concurrently on: ${ACTIVE_SERVERS[*]}"
    echo "=================================================="

    PIDS=()
    LOGS=()

    # Phase 2: Launch checks simultaneously
    for i in "${!ACTIVE_SERVERS[@]}"; do
        server="${ACTIVE_SERVERS[$i]}"
        log_file="/tmp/checkdpi_${server}.log"
        LOGS[$i]="$log_file"
        rm -f "$log_file"

        echo "[+] [$server] Starting check in background..."
        (
            if ssh "$server" "$CMD" > "$log_file" 2>&1; then
                exit 0
            else
                exit 1
            fi
        ) &
        PIDS[$i]=$!
    done

    echo ""
    echo "Waiting for all servers to complete (this may take 1-3 minutes)..."

    # Phase 3: Wait for completion and download results
    for i in "${!ACTIVE_SERVERS[@]}"; do
        server="${ACTIVE_SERVERS[$i]}"
        pid="${PIDS[$i]}"
        log_file="${LOGS[$i]}"

        echo ""
        echo "--- Waiting for $server (PID $pid) ---"
        if wait "$pid"; then
            echo "[+] [$server] Check completed successfully."
            download_report "$server"
            rm -f "$log_file"
        else
            echo "[-] Error: Check failed on $server" >&2
            if [ -f "$log_file" ]; then
                echo "Last 10 log lines from $server:" >&2
                tail -n 10 "$log_file" >&2
            fi
        fi
    done

else
    # Sequential mode
    echo "Starting DPI check in SEQUENTIAL mode on servers: ${SERVERS[*]}"
    echo "(Tip: use -p or --parallel to check all servers concurrently)"

    for server in "${SERVERS[@]}"; do
        echo ""
        echo "=================================================="
        echo "Server: $server"
        echo "=================================================="

        if ! preflight_check "$server"; then
            continue
        fi

        echo "Running check on $server..."
        if ! ssh -t "$server" "$CMD"; then
            echo "[-] Error: Command failed on $server" >&2
            continue
        fi

        download_report "$server"
    done
fi

echo ""
echo "Done."