#!/usr/bin/env bash
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
DATA_DIR="$SCRIPT_DIR"
CODE_ARCHIVE="$DATA_DIR/code.zip"
US_MONETARY_LOCALE="en_US.UTF-8"

restore_dump() {
    local label="$1"
    local dump_path="$2"
    local db_name="$3"
    local temp_output
    temp_output="$(mktemp)"

    echo "Recreating database $db_name..."
    if ! psql -h "$PG_HOST" -p "$PG_PORT" -U "$PG_USER" -d postgres \
        -c "DROP DATABASE IF EXISTS \"$db_name\" WITH (FORCE);" \
        -c "CREATE DATABASE \"$db_name\";" \
        -c "ALTER DATABASE \"$db_name\" SET lc_monetary TO '$US_MONETARY_LOCALE';" >"$temp_output" 2>&1; then
        cat "$temp_output"
        rm -f "$temp_output"
        echo "ERROR: Failed to recreate database $db_name"
        exit 1
    fi

    echo "Restoring $label..."
    if ! pg_restore -h "$PG_HOST" -p "$PG_PORT" -U "$PG_USER" \
        --no-owner --no-privileges --clean --if-exists -d "$db_name" \
        "$dump_path" >"$temp_output" 2>&1; then
        grep -v "does not exist, skipping" "$temp_output" || true
        rm -f "$temp_output"
        echo "ERROR: Failed to restore $(basename "$dump_path")"
        exit 1
    fi
    grep -v "does not exist, skipping" "$temp_output" || true
    rm -f "$temp_output"
    echo "Done."
}

echo "=== ReSect Artifact Setup ==="
echo

# Check prerequisites
missing=()
command -v pg_restore >/dev/null 2>&1 || missing+=("pg_restore (PostgreSQL client tools)")
command -v psql >/dev/null 2>&1 || missing+=("psql (PostgreSQL client tools)")
command -v node >/dev/null 2>&1 || missing+=("node (Node.js 22+)")
command -v pnpm >/dev/null 2>&1 || missing+=("pnpm")
command -v unzip >/dev/null 2>&1 || missing+=("unzip")

if [ ${#missing[@]} -ne 0 ]; then
    echo "ERROR: Missing required tools:"
    for tool in "${missing[@]}"; do
        echo "  - $tool"
    done
    echo
    echo "Please install the missing prerequisites (see README.md) and try again."
    exit 1
fi

# Check Node.js version
NODE_MAJOR=$(node -e "console.log(process.versions.node.split('.')[0])")
if [ "$NODE_MAJOR" -lt 22 ]; then
    echo "ERROR: Node.js 22 or later is required (found v$(node --version))"
    exit 1
fi

# Choose which dump(s) to restore
echo "Choose which dataset(s) to restore:"
echo "  1) attacks.dump only (smaller, false negative evaluation only)"
echo "  2) full.dump only (recommended, supports both false negative and false positive evaluation)"
while true; do
    read -rp "Selection [2]: " RESTORE_CHOICE
    RESTORE_CHOICE="${RESTORE_CHOICE:-2}"
    case "$RESTORE_CHOICE" in
        1)
            RESTORE_MODE="attack"
            RESTORE_ATTACK=true
            RESTORE_FULL=false
            DEFAULT_DB="resect-attacks"
            break
            ;;
        2)
            RESTORE_MODE="full"
            RESTORE_ATTACK=false
            RESTORE_FULL=true
            DEFAULT_DB="resect-full"
            break
            ;;
        *)
            echo "Please enter 1 or 2."
            ;;
    esac
done
echo

# Check selected dump files exist
if [ "$RESTORE_ATTACK" = true ] && [ ! -f "$DATA_DIR/attacks.dump" ]; then
    echo "ERROR: attacks.dump not found in $DATA_DIR"
    exit 1
fi
if [ "$RESTORE_FULL" = true ] && [ ! -f "$DATA_DIR/full.dump" ]; then
    echo "ERROR: full.dump not found in $DATA_DIR"
    exit 1
fi
if [ ! -f "$CODE_ARCHIVE" ]; then
    echo "ERROR: code.zip not found in $DATA_DIR"
    exit 1
fi

# Prompt for PostgreSQL connection details
read -rp "PostgreSQL host [localhost]: " PG_HOST
PG_HOST="${PG_HOST:-localhost}"

read -rp "PostgreSQL port [5432]: " PG_PORT
PG_PORT="${PG_PORT:-5432}"

read -rp "PostgreSQL user [postgres]: " PG_USER
PG_USER="${PG_USER:-postgres}"

read -rsp "PostgreSQL password: " PG_PASSWORD
echo

export PGPASSWORD="$PG_PASSWORD"

# Test connection
echo
echo "Testing PostgreSQL connection..."
if ! psql -h "$PG_HOST" -p "$PG_PORT" -U "$PG_USER" -d postgres -c "SELECT 1;" >/dev/null 2>&1; then
    echo "ERROR: Could not connect to PostgreSQL. Please check your credentials and ensure the server is running."
    exit 1
fi
echo "Connection successful."

echo "Checking PostgreSQL locale support..."
if ! psql -h "$PG_HOST" -p "$PG_PORT" -U "$PG_USER" -d postgres \
    -c "SET lc_monetary TO '$US_MONETARY_LOCALE';" >/dev/null 2>&1; then
    echo "ERROR: PostgreSQL on this machine does not accept lc_monetary='$US_MONETARY_LOCALE'."
    echo "This locale must be installed at the operating system level and visible to the PostgreSQL service."
    echo "Please install/generate the locale, restart PostgreSQL if required by your platform, and try again."
    exit 1
fi
echo "Locale check passed."

# Restore databases
echo
if [ "$RESTORE_ATTACK" = true ]; then
    restore_dump "attack dataset (attacks.dump)" "$DATA_DIR/attacks.dump" "resect-attacks"
fi
if [ "$RESTORE_FULL" = true ]; then
    echo "full.dump may take a few minutes to restore."
    restore_dump "full dataset (full.dump)" "$DATA_DIR/full.dump" "resect-full"
fi

# Extract ReSect source archive
echo
read -rp "Path to extract the ReSect source archive [./ReSect]: " REPO_DIR
REPO_DIR="${REPO_DIR:-./ReSect}"

if [ -d "$REPO_DIR" ]; then
    echo "Directory $REPO_DIR already exists, skipping archive extraction."
else
    echo "Extracting ReSect source archive..."
    EXTRACT_DIR="$(mktemp -d)"
    unzip -q "$CODE_ARCHIVE" -d "$EXTRACT_DIR"
    mkdir -p "$(dirname "$REPO_DIR")"
    entries=()
    while IFS= read -r entry; do
        entries+=("$entry")
    done < <(find "$EXTRACT_DIR" -mindepth 1 -maxdepth 1)
    if [ "${#entries[@]}" -eq 1 ] && [ -d "${entries[0]}" ]; then
        mv "${entries[0]}" "$REPO_DIR"
    else
        mkdir -p "$REPO_DIR"
        find "$EXTRACT_DIR" -mindepth 1 -maxdepth 1 -exec mv {} "$REPO_DIR" \;
    fi
    rm -rf "$EXTRACT_DIR"
fi

# Generate .env
ENV_FILE="$REPO_DIR/.env"
cat > "$ENV_FILE" <<EOF
ETHERSCAN_API_KEY=AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
DB_HOST=$PG_HOST
DB_PORT=$PG_PORT
DB_USER=$PG_USER
DB_PASSWORD=$PG_PASSWORD
DB_DATABASE=$DEFAULT_DB
EOF
echo
echo "Generated $ENV_FILE"
echo "Default DB_DATABASE is set to $DEFAULT_DB"

# Install Node.js dependencies
echo
echo "Installing Node.js dependencies..."
cd "$REPO_DIR"
pnpm install
echo

REPO_ABS="$(pwd)"
echo "=== Setup complete ==="
echo
echo "You can now run the evaluations from the ReSect repo ($REPO_ABS):"
echo
echo "  cd $REPO_ABS"
echo
if [ "$RESTORE_MODE" = "attack" ]; then
    echo "  # False negative evaluation (one tx per attack, ~15 seconds)"
    echo "  pnpm run evaluate:first"
    echo
    echo "  # False negative evaluation (all txs, ~1 minute)"
    echo "  pnpm run evaluate:all"
    echo
    echo "  # False positive evaluation is not available yet."
    echo "  # Restore full.dump later if you want to run pnpm run evaluate:fp"
elif [ "$RESTORE_MODE" = "full" ]; then
    echo "  # full.dump already contains the false negative subset."
    echo "  # False negative evaluation (one tx per attack, ~15 seconds)"
    echo "  pnpm run evaluate:first"
    echo
    echo "  # False negative evaluation (all txs, ~1 minute)"
    echo "  pnpm run evaluate:all"
    echo
    echo "  # False positive evaluation (often 5+ hours)"
    echo "  pnpm run evaluate:fp"
fi