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

# All-in-one helper: prepare circuits (Linux + host), rebuild the image, and run
# the chosen runner binary.
#
# Usage: scripts/run-demo.sh [compose|local|k8s] [run-seconds]
#   compose -> runs examples/src/bin/compose_runner.rs (default)
#   local   -> runs examples/src/bin/local_runner.rs
#   k8s     -> runs examples/src/bin/k8s_runner.rs
#   run-seconds defaults to 60
#
# Env overrides:
#   VERSION                       - circuits version (default v0.3.1)
#   NOMOS_TESTNET_IMAGE           - image tag (default nomos-testnet:local)
#   NOMOS_CIRCUITS_PLATFORM       - override host platform detection
#   NOMOS_CIRCUITS_REBUILD_RAPIDSNARK - set to 1 to force rapidsnark rebuild
#   NOMOS_NODE_REV                - nomos-node git rev for local binaries (default d2dd5a5084e1daef4032562c77d41de5e4d495f8)

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
MODE="${1:-compose}"
RUN_SECS="${2:-60}"
VERSION="${VERSION:-v0.3.1}"
IMAGE="${NOMOS_TESTNET_IMAGE:-nomos-testnet:local}"
NOMOS_NODE_REV="${NOMOS_NODE_REV:-d2dd5a5084e1daef4032562c77d41de5e4d495f8}"
RESTORED_BINARIES=0

case "$MODE" in
  compose) BIN="compose_runner" ;;
  local) BIN="local_runner" ;;
  k8s) BIN="k8s_runner" ;;
  *) echo "Unknown mode '$MODE' (use compose|local)" >&2; exit 1 ;;
esac

restore_binaries_from_tar() {
  local tar_path="${NOMOS_BINARIES_TAR:-${ROOT_DIR}/.tmp/nomos-binaries.tar.gz}"
  local extract_dir="${ROOT_DIR}/.tmp/nomos-binaries"
  if [ ! -f "$tar_path" ]; then
    return 1
  fi
  echo "==> Restoring binaries from ${tar_path}"
  rm -rf "${extract_dir}"
  mkdir -p "${extract_dir}"
  tar -xzf "$tar_path" -C "${extract_dir}"
  local src="${extract_dir}/artifacts"
  local bin_dst="${ROOT_DIR}/testing-framework/assets/stack/bin"
  local circuits_src="${src}/circuits"
  local circuits_dst="${ROOT_DIR}/testing-framework/assets/stack/kzgrs_test_params"
  if [ -f "${src}/nomos-node" ] && [ -f "${src}/nomos-executor" ] && [ -f "${src}/nomos-cli" ]; then
    mkdir -p "${bin_dst}"
    cp "${src}/nomos-node" "${bin_dst}/"
    cp "${src}/nomos-executor" "${bin_dst}/"
    cp "${src}/nomos-cli" "${bin_dst}/"
  else
    echo "Binaries missing in ${tar_path}; fallback to build-from-source path" >&2
    return 1
  fi
  if [ -d "${circuits_src}" ] && [ -f "${circuits_src}/kzgrs_test_params" ]; then
    rm -rf "${circuits_dst}"
    mkdir -p "${circuits_dst}"
    rsync -a --delete "${circuits_src}/" "${circuits_dst}/"
  else
    echo "Circuits missing in ${tar_path}; fallback to download/build path" >&2
    return 1
  fi
  RESTORED_BINARIES=1
  export RESTORED_BINARIES
}

ensure_host_binaries() {
  # Build nomos-node/nomos-executor for the host if not already present.
  HOST_SRC="${ROOT_DIR}/.tmp/nomos-node-host-src"
  HOST_TARGET="${ROOT_DIR}/.tmp/nomos-node-host-target"
  HOST_NODE_BIN_DEFAULT="${HOST_TARGET}/debug/nomos-node"
  HOST_EXEC_BIN_DEFAULT="${HOST_TARGET}/debug/nomos-executor"

  if [ -n "${NOMOS_NODE_BIN:-}" ] && [ -x "${NOMOS_NODE_BIN}" ] && [ -x "${NOMOS_EXECUTOR_BIN:-}" ]; then
    echo "Using provided host binaries:"
    echo "  NOMOS_NODE_BIN=${NOMOS_NODE_BIN}"
    echo "  NOMOS_EXECUTOR_BIN=${NOMOS_EXECUTOR_BIN}"
    return
  fi

  if [ -x "${HOST_NODE_BIN_DEFAULT}" ] && [ -x "${HOST_EXEC_BIN_DEFAULT}" ]; then
    echo "Host binaries already built at ${HOST_TARGET}"
    NOMOS_NODE_BIN="${HOST_NODE_BIN_DEFAULT}"
    NOMOS_EXECUTOR_BIN="${HOST_EXEC_BIN_DEFAULT}"
    export NOMOS_NODE_BIN NOMOS_EXECUTOR_BIN
    return
  fi

  echo "Building host nomos-node/nomos-executor from ${NOMOS_NODE_REV}"
  mkdir -p "${HOST_SRC}"
  if [ ! -d "${HOST_SRC}/.git" ]; then
    git clone https://github.com/logos-co/nomos-node.git "${HOST_SRC}"
  fi
  (
    cd "${HOST_SRC}"
    git fetch --depth 1 origin "${NOMOS_NODE_REV}"
    git checkout "${NOMOS_NODE_REV}"
    git reset --hard
    git clean -fdx
    RUSTFLAGS='--cfg feature="pol-dev-mode"' \
      NOMOS_CIRCUITS="${HOST_BUNDLE_PATH}" \
      cargo build --features "testing" \
        -p nomos-node -p nomos-executor -p nomos-cli \
        --target-dir "${HOST_TARGET}"
  )
  NOMOS_NODE_BIN="${HOST_NODE_BIN_DEFAULT}"
  NOMOS_EXECUTOR_BIN="${HOST_EXEC_BIN_DEFAULT}"
  export NOMOS_NODE_BIN NOMOS_EXECUTOR_BIN
}

restore_binaries_from_tar || true

echo "==> Preparing circuits (version ${VERSION})"
SETUP_OUT="/tmp/nomos-setup-output.$$"
if [ "${RESTORED_BINARIES}" -ne 1 ]; then
  "${ROOT_DIR}/scripts/setup-circuits-stack.sh" "${VERSION}" </dev/null | tee "$SETUP_OUT"
else
  echo "Skipping circuits setup; using restored bundle"
fi

# Prefer the host bundle if it exists; otherwise fall back to Linux bundle.
if [ -d "${ROOT_DIR}/.tmp/nomos-circuits-host" ]; then
  HOST_BUNDLE_PATH="${ROOT_DIR}/.tmp/nomos-circuits-host"
else
  HOST_BUNDLE_PATH="${ROOT_DIR}/testing-framework/assets/stack/kzgrs_test_params"
fi
rm -f "$SETUP_OUT"

# If the host bundle was somehow pruned, repair it once more.
if [ ! -x "${HOST_BUNDLE_PATH}/zksign/witness_generator" ]; then
  echo "Host circuits missing zksign/witness_generator; repairing..."
  "${ROOT_DIR}/scripts/setup-circuits-stack.sh" "${VERSION}"
fi
KZG_HOST_PATH="${HOST_BUNDLE_PATH}/kzgrs_test_params"
if [ ! -f "${KZG_HOST_PATH}" ]; then
  echo "KZG params missing at ${KZG_HOST_PATH}; rebuilding circuits bundle"
  "${ROOT_DIR}/scripts/setup-circuits-stack.sh" "${VERSION}"
fi

if [ "$MODE" != "local" ]; then
  if [ "${NOMOS_SKIP_IMAGE_BUILD:-0}" = "1" ]; then
    echo "==> Skipping testnet image rebuild (NOMOS_SKIP_IMAGE_BUILD=1)"
  else
    echo "==> Rebuilding testnet image (${IMAGE})"
    IMAGE_TAG="${IMAGE}" "${ROOT_DIR}/testing-framework/assets/stack/scripts/build_test_image.sh"
  fi
fi

if [ "$MODE" = "local" ]; then
  ensure_host_binaries
fi

echo "==> Running ${BIN} for ${RUN_SECS}s"
cd "${ROOT_DIR}"
if [ "$MODE" = "compose" ] || [ "$MODE" = "k8s" ]; then
  KZG_PATH="/kzgrs_test_params/kzgrs_test_params"
else
  KZG_PATH="${KZG_HOST_PATH}"
fi
POL_PROOF_DEV_MODE=true \
NOMOS_TESTNET_IMAGE="${IMAGE}" \
NOMOS_CIRCUITS="${HOST_BUNDLE_PATH}" \
NOMOS_KZGRS_PARAMS_PATH="${KZG_PATH}" \
COMPOSE_DEMO_RUN_SECS="${RUN_SECS}" \
LOCAL_DEMO_RUN_SECS="${RUN_SECS}" \
K8S_DEMO_RUN_SECS="${RUN_SECS}" \
NOMOS_NODE_BIN="${NOMOS_NODE_BIN:-}" \
NOMOS_EXECUTOR_BIN="${NOMOS_EXECUTOR_BIN:-}" \
  cargo run -p runner-examples --bin "${BIN}"