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

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
CONTROL_PLANE_DIR="$ROOT_DIR/app/control_plane"
WORKER_IMAGE=""
POSTGRES_PORT="${POSTGRES_PORT:-5433}"
E2E_CALLER_DATABASE_URL="${DATABASE_URL:-}"

load_dotenv_file() {
  local file="$1"

  if [[ ! -f "$file" ]]; then
    return
  fi

  echo "[e2e] loading dotenv $file"
  set -a
  # These files are local test-only dotenv files. Product provider configuration
  # still comes from Web UI-created provider rows; this only feeds e2e fixtures.
  # shellcheck disable=SC1090
  source "$file"
  set +a
}

load_e2e_dotenv() {
  load_dotenv_file "$CONTROL_PLANE_DIR/.env.test"
  load_dotenv_file "$CONTROL_PLANE_DIR/.env.test.local"
  load_dotenv_file "$CONTROL_PLANE_DIR/.env.local"
}

load_e2e_dotenv
if [[ -n "$E2E_CALLER_DATABASE_URL" ]]; then
  export DATABASE_URL="$E2E_CALLER_DATABASE_URL"
fi

usage() {
  cat <<'EOF'
Usage: tools/e2e/run [--gate|--perf|--chaos|--real-provider|--real-llm|--all]
                     [--providers=name[,name]] [--only=<mix test filter>]
                     [--no-build] [--no-services]

Runs the Ankole end-to-end suites from one entrypoint. Preflight resolves an
immutable published Worker image when it matches the current inputs, or builds
a content-addressed local image. It starts the devkit Postgres and Function
Runtime when either test service is unreachable.
The gate also pulls the pinned pooled release and builds the current control
plane for a real Compose upgrade. Set ANKOLE_E2E_CONTROL_PLANE_IMAGE only to
a locally built image of the current source and resolved Worker image.

Modes:
  --gate           Transport, main-flow, lifecycle, worker/computer, Compose
                   upgrade, and schedule suites (default).
  --perf           Concurrency/throughput baseline suite.
  --chaos          Chaos suite (worker kill, fabric restart, redelivery,
                   provider-send failures).
  --real-provider  AI Gateway real-provider smoke checks (needs API keys).
  --real-llm       Docker worker + real OpenRouter LLM suite (needs API key).
                   The chaos and Compose real-LLM scenarios run with --all, or
                   with --chaos --only=real_llm_chaos and
                   --gate --only=real_llm_compose_chaos.
  --all            Gate + perf + chaos, then real-provider/real-llm for the
                   providers whose API keys are present. With an OpenRouter
                   key, the gate and chaos suites also run their real-LLM
                   scenarios.

Options:
  --providers=...  Provider groups for --real-provider:
                   openrouter,jina,google,available
  --only=VALUE     Forwarded to mix test as --only VALUE (suite modes only).
  --no-build       Fail instead of building a Worker or control-plane image.
  --no-services    Fail instead of starting the devkit support services.

Examples:
  tools/e2e/run
  tools/e2e/run --all
  tools/e2e/run --chaos
  tools/e2e/run --gate --only=compose_upgrade
  tools/e2e/run --real-provider --providers=openrouter,jina
EOF
}

mode="gate"
providers="available"
only_filter=""
allow_build=1
allow_services=1

while [[ $# -gt 0 ]]; do
  case "$1" in
    --gate) mode="gate"; shift ;;
    --perf) mode="perf"; shift ;;
    --chaos) mode="chaos"; shift ;;
    --real-provider) mode="real-provider"; shift ;;
    --real-llm) mode="real-llm"; shift ;;
    --all) mode="all"; shift ;;
    --providers=*) providers="${1#--providers=}"; shift ;;
    --providers) providers="${2:?--providers requires a value}"; shift 2 ;;
    --only=*) only_filter="${1#--only=}"; shift ;;
    --only) only_filter="${2:?--only requires a value}"; shift 2 ;;
    --no-build) allow_build=0; shift ;;
    --no-services) allow_services=0; shift ;;
    -h|--help) usage; exit 0 ;;
    *)
      echo "Unknown argument: $1" >&2
      usage >&2
      exit 64
      ;;
  esac
done

log() { echo "[e2e] $*"; }

postgres_reachable() {
  (exec 3<>"/dev/tcp/127.0.0.1/$POSTGRES_PORT") 2>/dev/null && exec 3>&- && exec 3<&-
}

preflight_docker() {
  local output_file
  local resolver_args
  output_file="$(mktemp "${TMPDIR:-/tmp}/ankole-e2e-worker-image.XXXXXX")"
  resolver_args=(--scope complete --output "$output_file")
  if [[ "$allow_build" -ne 1 ]]; then
    resolver_args+=(--no-build)
  fi

  if ! (cd "$ROOT_DIR" && bun tools/devkit/src/local-worker-image.ts "${resolver_args[@]}"); then
    rm -f "$output_file"
    exit 1
  fi

  WORKER_IMAGE="$(tr -d '\r\n' < "$output_file")"
  rm -f "$output_file"
  if [[ -z "$WORKER_IMAGE" ]]; then
    echo "[e2e] Worker image resolver returned an empty image reference." >&2
    exit 1
  fi

  export ANKOLE_E2E_WORKER_IMAGE="$WORKER_IMAGE"
  log "using Worker image $WORKER_IMAGE"
}

preflight_postgres() {
  export ANKOLE_FUNCTION_RUNTIME_URL="${ANKOLE_FUNCTION_RUNTIME_URL:-http://127.0.0.1:5435}"
  export ANKOLE_FUNCTION_RUNTIME_TOKEN="${ANKOLE_FUNCTION_RUNTIME_TOKEN:-dev-function-runtime-token}"
  if postgres_reachable && curl -fsS "${ANKOLE_FUNCTION_RUNTIME_URL}/_internal/health" >/dev/null; then
    return
  fi

  if [[ "$allow_services" -ne 1 ]]; then
    echo "[e2e] Postgres or Function Runtime is unreachable and --no-services was given." >&2
    exit 1
  fi

  log "Starting devkit Postgres and Function Runtime"
  (cd "$ROOT_DIR" && bun run services:start)

  for _attempt in $(seq 1 30); do
    if postgres_reachable && curl -fsS "${ANKOLE_FUNCTION_RUNTIME_URL}/_internal/health" >/dev/null; then
      return
    fi
    sleep 1
  done

  echo "[e2e] Postgres or Function Runtime did not become healthy." >&2
  exit 1
}

run_mix() {
  local task="$1"
  shift
  (cd "$CONTROL_PLANE_DIR" && mix "$task" "$@")
}

suite_args=()
if [[ -n "$only_filter" ]]; then
  suite_args+=(--only "$only_filter")
fi

preflight_compose_upgrade() {
  if [[ -n "$only_filter" && "$only_filter" != "compose_upgrade" &&
        "$only_filter" != "authn_compose" &&
        "$only_filter" != "real_llm_compose_chaos" ]]; then
    return
  fi

  local deployment_images
  deployment_images="$(ANKOLE_HOST=127.0.0.1 ACME_EMAIL=e2e@example.test \
    ANKOLE_SECRET_BASE=compose-test-only RELEASE_COOKIE=compose-test-only \
    POSTGRES_PASSWORD=compose-test-only ANKOLE_DOCKER_GID=0 \
    ANKOLE_CONTROL_PLANE_IMAGE= ANKOLE_POSTGRESQL_IMAGE= \
    docker compose --env-file /dev/null -f "$ROOT_DIR/tools/deploy/docker-compose/compose.yaml" config --images)"
  while IFS= read -r image; do
    if ! docker image inspect "$image" >/dev/null 2>&1; then
      docker pull "$image"
    fi
  done <<< "$deployment_images"

  local legacy_revision
  legacy_revision="$(tr -d '\r\n' < "$ROOT_DIR/tools/e2e/pooled-worker-revision")"
  for repository in ankole-agent-control-plane ankole-agent-computer-worker; do
    local image="ghcr.io/agentbull/$repository:$legacy_revision"
    if ! docker image inspect "$image" >/dev/null 2>&1; then
      docker pull "$image"
    fi
  done

  local legacy_postgres_image
  legacy_postgres_image="$(tr -d '\r\n' < "$ROOT_DIR/tools/e2e/pooled-postgres-image")"
  if ! docker image inspect "$legacy_postgres_image" >/dev/null 2>&1; then
    docker pull "$legacy_postgres_image"
  fi

  if [[ -z "${ANKOLE_E2E_CONTROL_PLANE_IMAGE:-}" ]]; then
    if [[ "$allow_build" -ne 1 ]]; then
      echo '[e2e] Compose upgrade needs ANKOLE_E2E_CONTROL_PLANE_IMAGE when --no-build is used.' >&2
      exit 1
    fi
    export ANKOLE_E2E_CONTROL_PLANE_IMAGE=ankole-control-plane:e2e
    docker build --file "$CONTROL_PLANE_DIR/Dockerfile" \
      --build-arg "ANKOLE_AGENT_COMPUTER_IMAGE=$WORKER_IMAGE" \
      --tag "$ANKOLE_E2E_CONTROL_PLANE_IMAGE" "$ROOT_DIR"
  fi
  docker image inspect "$ANKOLE_E2E_CONTROL_PLANE_IMAGE" >/dev/null
}

run_gate() {
  preflight_compose_upgrade
  run_mix e2e.gate "${suite_args[@]:+${suite_args[@]}}"
}
run_perf() { run_mix e2e.perf "${suite_args[@]:+${suite_args[@]}}"; }
run_chaos() { run_mix e2e.chaos "${suite_args[@]:+${suite_args[@]}}"; }

run_real_llm() {
  (cd "$CONTROL_PLANE_DIR" && ANKOLE_REAL_LLM_E2E=1 mix e2e.real_llm "${suite_args[@]:+${suite_args[@]}}")
}

run_real_provider() {
  run_mix e2e.ai_gateway_real_provider -- --providers="$providers"
}

has_openrouter_key() {
  [[ -n "${OPENROUTER_API_KEY:-}" ]] || [[ -n "${OPEN_ROUTER_API_KEY:-}" ]]
}

has_real_provider_env() {
  [[ -n "${OPENAI_API_KEY:-}" ]] ||
    has_openrouter_key ||
    [[ -n "${JINA_API_KEY:-}" ]] ||
    [[ -n "${GOOGLE_AI_STUDIO_API_KEY:-}" ]]
}

case "$mode" in
  gate)
    preflight_docker
    preflight_postgres
    run_gate
    ;;
  perf)
    preflight_docker
    preflight_postgres
    run_perf
    ;;
  chaos)
    preflight_docker
    preflight_postgres
    run_chaos
    ;;
  real-provider)
    preflight_postgres
    run_real_provider
    ;;
  real-llm)
    preflight_docker
    preflight_postgres
    run_real_llm
    ;;
  all)
    preflight_docker
    preflight_postgres
    if has_openrouter_key; then export ANKOLE_REAL_LLM_E2E=1; fi
    run_gate
    run_perf
    run_chaos
    if has_real_provider_env; then
      providers="available"
      run_real_provider
    else
      log "skipping AI Gateway real-provider e2e: no provider API keys are present"
    fi
    if has_openrouter_key; then
      run_real_llm
    else
      log "skipping real-LLM worker e2e: no OpenRouter API key is present"
    fi
    ;;
esac
