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

if [[ -z "${IN_NIX_SHELL:-}" ]] && command -v nix >/dev/null 2>&1; then
  exec nix develop -c "$0" "$@"
fi

WORKDIR=$(mktemp -d)
PD_LOG="${WORKDIR}/flaredb-pd.log"
S1_LOG="${WORKDIR}/flaredb-server-1.log"
S2_LOG="${WORKDIR}/flaredb-server-2.log"

run_client() {
  local addr="$1"
  shift
  local output=""
  local status=0
  local attempt=0
  while (( attempt < 20 )); do
    if output=$(cargo run --quiet --bin flaredb-client -- --addr "${addr}" --pd-addr 127.0.0.1:2479 "$@" 2>&1); then
      printf '%s\n' "${output}" | awk 'NF { last = $0 } END { print last }'
      return 0
    fi
    status=$?
    attempt=$((attempt + 1))
    sleep 1
  done
  printf '%s\n' "${output}" >&2
  return "${status}"
}

cleanup() {
  local exit_code=$?
  if [[ -n "${PD_PID:-}" ]]; then
    kill "${PD_PID}" >/dev/null 2>&1 || true
  fi
  if [[ -n "${S1_PID:-}" ]]; then
    kill "${S1_PID}" >/dev/null 2>&1 || true
  fi
  if [[ -n "${S2_PID:-}" ]]; then
    kill "${S2_PID}" >/dev/null 2>&1 || true
  fi
  if (( exit_code != 0 )); then
    echo "verify-sharding failed; logs preserved at ${WORKDIR}" >&2
    [[ -f "${PD_LOG}" ]] && { echo "--- ${PD_LOG} ---" >&2; tail -n 200 "${PD_LOG}" >&2; }
    [[ -f "${S1_LOG}" ]] && { echo "--- ${S1_LOG} ---" >&2; tail -n 200 "${S1_LOG}" >&2; }
    [[ -f "${S2_LOG}" ]] && { echo "--- ${S2_LOG} ---" >&2; tail -n 200 "${S2_LOG}" >&2; }
    return "${exit_code}"
  fi
  rm -rf "${WORKDIR}"
}
trap cleanup EXIT

echo "Building workspace..."
cargo build

echo "Starting PD..."
cargo run --bin flaredb-pd -- --addr 127.0.0.1:2479 >"${PD_LOG}" 2>&1 &
PD_PID=$!
sleep 2

echo "Starting Server 1 (127.0.0.1:50001, data1)..."
cargo run --bin flaredb-server -- \
  --store-id 1 \
  --addr 127.0.0.1:50001 \
  --http-addr 127.0.0.1:8083 \
  --data-dir "${WORKDIR}/data1" \
  --pd-addr 127.0.0.1:2479 \
  --metrics-port 9093 \
  --namespace-mode raw=eventual \
  >"${S1_LOG}" 2>&1 &
S1_PID=$!
sleep 4

echo "Starting Server 2 (127.0.0.1:50002, data2)..."
cargo run --bin flaredb-server -- \
  --store-id 2 \
  --addr 127.0.0.1:50002 \
  --http-addr 127.0.0.1:8084 \
  --data-dir "${WORKDIR}/data2" \
  --pd-addr 127.0.0.1:2479 \
  --metrics-port 9094 \
  --namespace-mode raw=eventual \
  >"${S2_LOG}" 2>&1 &
S2_PID=$!

sleep 5 # Wait for registration and leader routing to settle

echo "Running Client Verification (multi-node routing smoke)..."

echo "Testing Put 'a'..."
run_client 127.0.0.1:50001 --namespace raw raw-put --key a --value val_a >/dev/null

echo "Testing Put 'z'..."
run_client 127.0.0.1:50002 --namespace raw raw-put --key z --value val_z >/dev/null

echo "Testing reads from both nodes..."
VALUE_A=$(run_client 127.0.0.1:50002 --namespace raw raw-get --key a)
VALUE_Z=$(run_client 127.0.0.1:50001 --namespace raw raw-get --key z)
[[ "${VALUE_A}" == "val_a" ]]
[[ "${VALUE_Z}" == "val_z" ]]

echo "Sharding Verification Complete!"