#!/usr/bin/env bash
# ==============================================================================
# repl-test.sh: Integration harness for Records API & Replication (mTLS)
#
# Supports N instances in multiple topologies:
#   star    - one primary, all replicas pull directly from it   (default)
#   chain   - a linear chain: primary -> r1 -> r2 -> ... -> rN
#
# Environment variable overrides:
#   NODES=4            number of total nodes (1 primary + N-1 replicas; default: 2)
#   BASE_PORT=3000     base port; nodes occupy BASE_PORT .. BASE_PORT+NODES-1
#   TOPO=star|chain    replication topology (default: star)
#   TOKEN=...          shared replication bearer token
#   START_DASHBOARD=1  launch the web dashboard (default: 1)
#   SKIP_SWITCHOVER=0  skip the switchover smoke test (default: 0)
#   NO_WAIT=0          exit immediately after smoke tests instead of waiting
#
# Examples:
#   bash repl-test.sh                         # 2 nodes, star, default ports
#   NODES=3 bash repl-test.sh                 # 1 primary + 2 replicas, star
#   NODES=4 TOPO=chain bash repl-test.sh      # 4-node replication chain
#   NODES=5 SKIP_SWITCHOVER=1 bash repl-test.sh
# ==============================================================================
set -euo pipefail

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
BIN="${BIN:-$ROOT_DIR/target/debug/rust_test}"
WORK="${WORK:-$ROOT_DIR/repl-test}"
BASE_PORT="${BASE_PORT:-3000}"
NODES="${NODES:-2}"           # total node count (must be >= 2)
TOPO="${TOPO:-star}"          # star | chain
TOKEN="${TOKEN:-change-me}"
DASHBOARD_PORT="${DASHBOARD_PORT:-8080}"
CORS_ORIGIN="${CORS_ORIGIN:-http://127.0.0.1:$DASHBOARD_PORT}"
START_DASHBOARD="${START_DASHBOARD:-1}"
SKIP_SWITCHOVER="${SKIP_SWITCHOVER:-0}"
NO_WAIT="${NO_WAIT:-0}"

if [[ "$NODES" -lt 2 ]]; then
  echo "Error: NODES must be >= 2." >&2
  exit 1
fi

# ---------------------------------------------------------------------------
# Node naming and port helpers
# ---------------------------------------------------------------------------
# node_name 0  -> "primary"
# node_name 1  -> "replica-1"
# node_name 2  -> "replica-2" ...
node_name() {
  local idx=$1
  if [[ $idx -eq 0 ]]; then echo "primary"; else echo "replica-$idx"; fi
}

# node_port 0  -> BASE_PORT
# node_port 1  -> BASE_PORT+1 ...
node_port() { echo $(( BASE_PORT + $1 )); }

PIDS=()

# ---------------------------------------------------------------------------
# Helper utilities
# ---------------------------------------------------------------------------
check_port() {
  local port=$1
  if command -v lsof >/dev/null 2>&1; then
    lsof -iTCP:"$port" -sTCP:LISTEN -n -P >/dev/null 2>&1
  elif command -v nc >/dev/null 2>&1; then
    nc -z 127.0.0.1 "$port" >/dev/null 2>&1
  else
    (exec 3<>/dev/tcp/127.0.0.1/"$port") 2>/dev/null && { exec 3>&-; exec 3<&-; return 0; } || return 1
  fi
}

cleanup() {
  trap - EXIT INT TERM
  echo -e "\n\033[1;33mStopping all instances and dashboard...\033[0m"
  for pid in "${PIDS[@]}"; do
    if kill -0 "$pid" 2>/dev/null; then
      kill "$pid" 2>/dev/null || true
    fi
  done
  wait 2>/dev/null || true
  echo "Cleaned up."
}
trap cleanup EXIT INT TERM

json_field() {
  # Usage: json_field <json_string> <field_name>
  local json="$1" field="$2"
  if command -v jq >/dev/null 2>&1; then
    echo "$json" | jq -r ".${field} // empty"
  else
    echo "$json" | grep -o "\"${field}\":\"[^\"]*" | head -1 | cut -d'"' -f4
  fi
}

color_echo() {
  local color=$1; shift
  local code
  case "$color" in
    red)    code='\033[1;31m' ;;
    green)  code='\033[1;32m' ;;
    yellow) code='\033[1;33m' ;;
    blue)   code='\033[1;34m' ;;
    *)      code='\033[0m'    ;;
  esac
  echo -e "${code}$*\033[0m"
}

wait_up() {
  local name=$1 port=$2
  for _ in $(seq 1 80); do
    if curl -k -sf "https://127.0.0.1:$port/api/now" >/dev/null 2>&1; then
      return 0
    fi
    sleep 0.1
  done
  color_echo red "Error: $name failed to start on :$port" >&2
  if [[ -f "$WORK/$name/log.txt" ]]; then
    echo "--- Last 30 lines of $WORK/$name/log.txt ---" >&2
    tail -n 30 "$WORK/$name/log.txt" >&2
  fi
  exit 1
}

mint_join_token() {
  local primary_port=$1
  local resp
  resp=$(curl -k -sf -X POST "https://127.0.0.1:$primary_port/api/cluster/tokens" \
    -H 'Content-Type: application/json' \
    -d '{"ttl_secs": 300}')
  json_field "$resp" "token"
}

wait_replica_caught_up() {
  local name=$1 port=$2 target_seq=$3
  for _ in $(seq 1 80); do
    local out
    out=$(curl -k -s "https://127.0.0.1:$port/api/replication/status" 2>/dev/null || true)
    local applied
    applied=$(json_field "$out" "sync.last_applied_seq" 2>/dev/null || echo "0")
    # Fallback grep for nested .sync.last_applied_seq
    if [[ -z "$applied" || "$applied" == "null" ]]; then
      applied=$(echo "$out" | grep -o '"last_applied_seq":[0-9]*' | head -1 | cut -d: -f2)
    fi
    if [[ "${applied:-0}" -ge "$target_seq" ]]; then
      return 0
    fi
    sleep 0.25
  done
  color_echo yellow "  ⚠ WARNING: $name (:$port) did not reach seq $target_seq in time." >&2
  return 1
}

# ---------------------------------------------------------------------------
# Pre-flight
# ---------------------------------------------------------------------------
for idx in $(seq 0 $(( NODES - 1 ))); do
  port=$(node_port "$idx")
  if check_port "$port"; then
    color_echo red "Error: Port $port is already in use." >&2
    exit 1
  fi
done
if [[ "$START_DASHBOARD" -eq 1 ]]; then
  if check_port "$DASHBOARD_PORT"; then
    color_echo yellow "Dashboard port :$DASHBOARD_PORT is in use; dashboard will be skipped."
    START_DASHBOARD=0
  fi
fi

# Build
if [[ -f "$ROOT_DIR/Cargo.toml" ]]; then
  echo "Ensuring debug binary is up to date..."
  ( cd "$ROOT_DIR" && cargo build )
elif [[ ! -x "$BIN" ]]; then
  color_echo red "Error: Binary not found: $BIN" >&2; exit 1
fi

# ---------------------------------------------------------------------------
# Workspace setup
# ---------------------------------------------------------------------------
rm -rf "$WORK"
for idx in $(seq 0 $(( NODES - 1 ))); do
  name=$(node_name "$idx")
  mkdir -p "$WORK/$name"
  cat > "$WORK/$name/config.json" <<'CFGEOF'
{
  "mode": "loose",
  "fields": {
    "version": { "type": "integer" }
  }
}
CFGEOF
done

# ---------------------------------------------------------------------------
# start_node <name> <port> [extra env vars...]
# ---------------------------------------------------------------------------
start_node() {
  local name=$1 port=$2
  shift 2
  ( cd "$WORK/$name" && exec env \
      APP_CONFIG=config.json \
      APP_BIND="127.0.0.1:$port" \
      APP_REPL_TOKEN="$TOKEN" \
      APP_SYNC_INTERVAL_SECS=1 \
      APP_CORS_ORIGIN="$CORS_ORIGIN" \
      "$@" "$BIN" ) >"$WORK/$name/log.txt" 2>&1 &
  PIDS+=("$!")
}

# ==============================================================================
# 1. Start Primary
# ==============================================================================
PRIMARY_PORT=$(node_port 0)
color_echo green "==> [1/N] Starting Primary on :$PRIMARY_PORT (HTTPS/mTLS)..."
start_node primary "$PRIMARY_PORT" APP_ROLE=primary
wait_up "primary" "$PRIMARY_PORT"
color_echo green "    ✔ Primary is up."

# ==============================================================================
# 2. Retrieve CA cert (sanity check)
# ==============================================================================
color_echo blue "\n==> Verifying Cluster CA is accessible..."
ca_resp=$(curl -k -sf "https://127.0.0.1:$PRIMARY_PORT/api/cluster/ca")
if echo "$ca_resp" | grep -q 'BEGIN CERTIFICATE'; then
  color_echo green "    ✔ Cluster CA certificate downloaded."
else
  color_echo red "    ✖ Failed to retrieve Cluster CA – aborting." >&2; exit 1
fi

# ==============================================================================
# 3. Start replicas according to topology
#
# star  topology: every replica pulls directly from the primary
# chain topology: replica-N pulls from replica-(N-1), which pulls from primary
# ==============================================================================

# Arrays for later use
PORTS=()        # PORTS[i] = port for node i
NAMES=()        # NAMES[i] = name for node i
SOURCES=()      # SOURCES[i] = upstream URL for node i (empty for primary)

PORTS[0]=$PRIMARY_PORT
NAMES[0]="primary"
SOURCES[0]=""

for idx in $(seq 1 $(( NODES - 1 ))); do
  name=$(node_name "$idx")
  port=$(node_port "$idx")
  PORTS[$idx]=$port
  NAMES[$idx]="$name"

  case "$TOPO" in
    chain)
      # Replica pulls from the previous node in the chain
      upstream_port=$(node_port $(( idx - 1 )))
      upstream_url="https://127.0.0.1:$upstream_port"
      upstream_name=$(node_name $(( idx - 1 )))
      ;;
    star|*)
      # Replica pulls directly from the primary
      upstream_port=$PRIMARY_PORT
      upstream_url="https://127.0.0.1:$PRIMARY_PORT"
      upstream_name="primary"
      ;;
  esac

  SOURCES[$idx]="$upstream_url"

  # The join token must come from the Cluster CA holder (always the primary)
  color_echo blue "\n==> Minting join token for $name from primary..."
  JOIN_TOKEN=$(mint_join_token "$PRIMARY_PORT")
  color_echo green "    ✔ Token: $JOIN_TOKEN"

  color_echo green "\n==> [$idx/$(( NODES - 1 ))] Starting $name on :$port"
  echo "    Topology: $TOPO  |  upstream: $upstream_name ($upstream_url)"

  # In chain topology the replication upstream differs from the CA issuer (primary).
  # APP_JOIN_URL overrides enrollment so the new replica can get a cert from the
  # actual primary even though it pulls data from an intermediate relay node.
  start_node "$name" "$port" \
    APP_ROLE=replica \
    APP_PRIMARY_URL="$upstream_url" \
    APP_JOIN_URL="https://127.0.0.1:$PRIMARY_PORT" \
    APP_JOIN_TOKEN="$JOIN_TOKEN"
  wait_up "$name" "$port"
  color_echo green "    ✔ $name is up and enrolled."

  if [[ -f "$WORK/$name/messages.db" ]]; then
    color_echo green "    ✔ mTLS credentials persisted in SQLite."
  fi
done

# ---------------------------------------------------------------------------
# Dashboard
# ---------------------------------------------------------------------------
if [[ "$START_DASHBOARD" -eq 1 && -f "$ROOT_DIR/dashboard.py" ]]; then
  color_echo green "\n==> Starting Dashboard on :$DASHBOARD_PORT..."
  ( cd "$ROOT_DIR" && exec python3 dashboard.py ) >"$WORK/dashboard.log" 2>&1 &
  PIDS+=("$!")
fi

# ==============================================================================
# 4. Write smoke-test records to primary
# ==============================================================================
RECORD_COUNT=3
color_echo blue "\n==> Writing $RECORD_COUNT test records to Primary..."
for i in $(seq 1 "$RECORD_COUNT"); do
  curl -k -s -X POST "https://127.0.0.1:$PRIMARY_PORT/api/records" \
    -H 'Content-Type: application/json' \
    -d "{\"data\":{\"version\":1,\"msg\":\"hello $i\"}}" >/dev/null
  echo "    Record #$i written"
done

# ==============================================================================
# 5. Verify every replica has converged
# ==============================================================================
color_echo blue "\n==> Waiting for all replicas to converge to seq $RECORD_COUNT..."
all_ok=1
for idx in $(seq 1 $(( NODES - 1 ))); do
  name="${NAMES[$idx]}"
  port="${PORTS[$idx]}"
  echo -n "    Checking $name (:$port)... "
  if wait_replica_caught_up "$name" "$port" "$RECORD_COUNT"; then
    color_echo green "✔ seq $RECORD_COUNT reached"
  else
    echo "✖ timeout"
    all_ok=0
  fi
done
if [[ $all_ok -eq 1 ]]; then
  color_echo green "✔ All replicas are in sync."
else
  color_echo yellow "⚠ Some replicas did not converge in time."
fi

# ==============================================================================
# 6. Read-only guard check on every replica
# ==============================================================================
color_echo blue "\n==> Verifying read-only guard on all replicas..."
for idx in $(seq 1 $(( NODES - 1 ))); do
  name="${NAMES[$idx]}"
  port="${PORTS[$idx]}"
  http_code=$(curl -k -s -o /dev/null -w "%{http_code}" \
    -X POST "https://127.0.0.1:$port/api/records" \
    -H 'Content-Type: application/json' \
    -d '{"data":{"version":1,"msg":"should be rejected"}}')
  if [[ "$http_code" == "403" ]]; then
    color_echo green "    ✔ $name (:$port) → HTTP 403 Forbidden (correct)"
  else
    color_echo red "    ✖ $name (:$port) → HTTP $http_code (expected 403)"
  fi
done

# ==============================================================================
# 7. Graceful switchover (primary -> first replica)
#    Skipped if SKIP_SWITCHOVER=1 or only 1 replica exists (would leave no primary)
# ==============================================================================
if [[ "$SKIP_SWITCHOVER" -eq 0 && "$NODES" -ge 2 ]]; then
  FIRST_REPLICA_PORT="${PORTS[1]}"
  FIRST_REPLICA_NAME="${NAMES[1]}"

  color_echo blue "\n==> Initiating Graceful Switchover: primary (:$PRIMARY_PORT) ─► $FIRST_REPLICA_NAME (:$FIRST_REPLICA_PORT)..."
  switch_resp=$(curl -k -sf -X POST "https://127.0.0.1:$PRIMARY_PORT/api/admin/switchover" \
    -H 'Content-Type: application/json' \
    -d "{\"replica_url\":\"https://127.0.0.1:$FIRST_REPLICA_PORT\",\"timeout_secs\":15}" || true)
  color_echo green "    ✔ Switchover response: $switch_resp"

  # Old primary is now replica – verify 403
  http_code=$(curl -k -s -o /dev/null -w "%{http_code}" \
    -X POST "https://127.0.0.1:$PRIMARY_PORT/api/records" \
    -H 'Content-Type: application/json' \
    -d '{"data":{"version":1,"msg":"should fail on demoted primary"}}')
  if [[ "$http_code" == "403" ]]; then
    color_echo green "    ✔ Demoted primary returned HTTP 403 (read-only)"
  else
    color_echo red "    ✖ Demoted primary returned HTTP $http_code (expected 403)"
  fi

  # Write to new primary
  NEXT_SEQ=$(( RECORD_COUNT + 1 ))
  color_echo blue "\n==> Writing record #$NEXT_SEQ to new primary ($FIRST_REPLICA_NAME :$FIRST_REPLICA_PORT)..."
  write_resp=$(curl -k -sf -X POST "https://127.0.0.1:$FIRST_REPLICA_PORT/api/records" \
    -H 'Content-Type: application/json' \
    -d "{\"data\":{\"version\":1,\"msg\":\"post-switchover write\"}}" || true)
  color_echo green "    ✔ Created: $write_resp"

  # After switchover in star topology, the other replicas still point to the old primary
  # which now points to the new primary – they will converge via chaining.
  # In chain topology they chain through naturally.
  # Verify the demoted primary (now replica) catches up:
  color_echo blue "\n==> Verifying all nodes converge to seq $NEXT_SEQ post-switchover..."
  all_ok=1
  for idx in $(seq 0 $(( NODES - 1 ))); do
    port="${PORTS[$idx]}"
    name="${NAMES[$idx]}"
    if [[ $idx -eq 1 ]]; then continue; fi   # new primary already has it
    echo -n "    Checking $name (:$port)... "
    if wait_replica_caught_up "$name" "$port" "$NEXT_SEQ"; then
      color_echo green "✔ seq $NEXT_SEQ reached"
    else
      echo "✖ timeout"
      all_ok=0
    fi
  done
  if [[ $all_ok -eq 1 ]]; then
    color_echo green "✔ All nodes converged after switchover."
  else
    color_echo yellow "⚠ Some nodes did not converge after switchover."
  fi
fi

# ==============================================================================
# Ready Banner
# ==============================================================================
echo ""
echo "======================================================================"
echo "  Cluster is healthy  •  Topology: $TOPO  •  Nodes: $NODES"
echo ""
echo "  Nodes:"
for idx in $(seq 0 $(( NODES - 1 ))); do
  port="${PORTS[$idx]}"
  name="${NAMES[$idx]}"
  if [[ $idx -eq 0 ]]; then
    echo "    [$(( idx + 1 ))] Primary  (:$port)  https://127.0.0.1:$port"
  else
    src="${SOURCES[$idx]}"
    echo "    [$(( idx + 1 ))] Replica  (:$port)  https://127.0.0.1:$port  <─ $src"
  fi
done
echo ""
if [[ "$START_DASHBOARD" -eq 1 ]]; then
  echo "  Web Dashboard:  http://127.0.0.1:$DASHBOARD_PORT"
fi
echo "  Logs Dir:       $WORK/"
echo ""
echo "  Press Ctrl+C to stop all instances."
echo "======================================================================"

if [[ "$NO_WAIT" -ne 1 ]]; then
  wait
fi