#!/usr/bin/env bash
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

# Shared helpers for local gateway-backed e2e wrappers. Driver-specific setup,
# cleanup, and runtime behavior stay in the Docker/Podman wrapper scripts.

# E2E traffic is synthetic and must not contribute to product usage metrics.
# Keep an explicit override so telemetry-specific tests can opt back in.
export OPENSHELL_TELEMETRY_ENABLED="${OPENSHELL_TELEMETRY_ENABLED:-false}"

# Resolve a test image override. Repository-only values inherit the caller's
# tag, while tagged and digest-pinned references are already complete.
e2e_image_reference_is_complete() {
  local image=$1
  local last_component="${image##*/}"

  [[ "${image}" == *@* || "${last_component}" == *:* ]]
}

e2e_image_reference_has_digest() {
  [[ "$1" == *@* ]]
}

e2e_resolve_image_reference() {
  local image=$1
  local tag=$2

  if e2e_image_reference_is_complete "${image}"; then
    printf '%s\n' "${image}"
  else
    printf '%s:%s\n' "${image%/}" "${tag}"
  fi
}

e2e_image_reference_repository() {
  local image=$1
  local repository="${image%%@*}"
  local last_component="${repository##*/}"

  if [[ "${last_component}" == *:* ]]; then
    repository="${repository%:*}"
  fi
  printf '%s\n' "${repository}"
}

# Return the registry portion of an image repository. Docker treats the first
# path component as a registry when it contains a dot or colon, or is
# `localhost`; otherwise the image uses the configured/default registry.
e2e_image_reference_registry() {
  local repository
  local first_component

  repository="$(e2e_image_reference_repository "$1")"
  first_component="${repository%%/*}"
  if [[ "${repository}" == */* ]] \
    && { [[ "${first_component}" == *.* ]] || [[ "${first_component}" == *:* ]] || [[ "${first_component}" == "localhost" ]]; }; then
    printf '%s\n' "${first_component}"
  fi
}

# Return the repository path without its registry, suitable for Helm's
# <component>.image.repository values.
e2e_image_reference_repository_path() {
  local repository
  local registry

  repository="$(e2e_image_reference_repository "$1")"
  registry="$(e2e_image_reference_registry "$1")"
  if [ -n "${registry}" ]; then
    printf '%s\n' "${repository#"${registry}"/}"
  else
    printf '%s\n' "${repository}"
  fi
}

e2e_image_reference_tag() {
  local image=$1
  local repository="${image%%@*}"
  local last_component="${repository##*/}"

  if [[ "${image}" == *@* || "${last_component}" != *:* ]]; then
    return 0
  fi
  printf '%s\n' "${last_component##*:}"
}

e2e_image_reference_digest() {
  if [[ "$1" == *@* ]]; then
    printf '%s\n' "${1#*@}"
  fi
}

e2e_cargo_target_dir() {
  local root=$1
  shift
  local cargo_command=(cargo)

  if [ "$#" -gt 0 ]; then
    cargo_command=("$@")
  fi

  if [ -n "${CARGO_TARGET_DIR:-}" ]; then
    case "${CARGO_TARGET_DIR}" in
      /*) printf '%s\n' "${CARGO_TARGET_DIR}" ;;
      *) printf '%s\n' "${root}/${CARGO_TARGET_DIR}" ;;
    esac
    return 0
  fi

  "${cargo_command[@]}" metadata --format-version=1 --no-deps \
    | python3 -c 'import json, sys; print(json.load(sys.stdin)["target_directory"])'
}

e2e_endpoint_port() {
  python3 - "$1" <<'PY'
import sys
from urllib.parse import urlparse

parsed = urlparse(sys.argv[1])
print(parsed.port or (443 if parsed.scheme == "https" else 80))
PY
}

e2e_pick_port() {
  python3 -c 'import socket; s=socket.socket(); s.bind(("",0)); print(s.getsockname()[1]); s.close()'
}

e2e_generate_pki() {
  local gateway_bin=$1
  local pki_dir=$2
  shift 2
  # Remaining args are extra --server-san values (e.g. host.containers.internal).
  # host.docker.internal and localhost are already in the default SAN list.

  local san_args=()
  san_args+=(--server-san host.openshell.internal)
  for san in "$@"; do
    san_args+=(--server-san "${san}")
  done

  "${gateway_bin}" generate-certs --output-dir "${pki_dir}" "${san_args[@]}"
}

e2e_preserve_mise_dirs() {
  if ! command -v mise >/dev/null 2>&1; then
    return 0
  fi

  if [ -z "${MISE_DATA_DIR:-}" ]; then
    export MISE_DATA_DIR="${XDG_DATA_HOME:-${HOME}/.local/share}/mise"
  fi

  if [ -z "${MISE_CACHE_DIR:-}" ]; then
    case "$(uname -s)" in
      Darwin) export MISE_CACHE_DIR="${HOME}/Library/Caches/mise" ;;
      *) export MISE_CACHE_DIR="${XDG_CACHE_HOME:-${HOME}/.cache}/mise" ;;
    esac
  fi

  if [ -z "${MISE_STATE_DIR:-}" ]; then
    export MISE_STATE_DIR="${XDG_STATE_HOME:-${HOME}/.local/state}/mise"
  fi
}

e2e_align_docker_host_with_cli_context() {
  if [ -n "${DOCKER_HOST:-}" ] || ! command -v docker >/dev/null 2>&1; then
    return 0
  fi

  local endpoint
  endpoint="$(docker context inspect --format '{{ .Endpoints.docker.Host }}' 2>/dev/null || true)"
  if [ -z "${endpoint}" ] || [ "${endpoint}" = "<no value>" ]; then
    return 0
  fi

  export DOCKER_HOST="${endpoint}"
  echo "Using Docker endpoint from active context: ${DOCKER_HOST}"
}

e2e_register_plaintext_gateway() {
  local config_home=$1
  local name=$2
  local endpoint=$3
  local port=$4
  local gateway_config_dir="${config_home}/openshell/gateways/${name}"

  mkdir -p "${gateway_config_dir}"
  cat >"${gateway_config_dir}/metadata.json" <<EOF
{
  "name": "${name}",
  "gateway_endpoint": "${endpoint}",
  "is_remote": false,
  "gateway_port": ${port},
  "auth_mode": "plaintext"
}
EOF
  printf '%s' "${name}" >"${config_home}/openshell/active_gateway"
}

e2e_register_mtls_gateway() {
  local config_home=$1
  local name=$2
  local endpoint=$3
  local port=$4
  local pki_dir=$5
  local oidc_issuer="${6:-}"
  local gateway_config_dir="${config_home}/openshell/gateways/${name}"

  mkdir -p "${gateway_config_dir}/mtls"
  cp "${pki_dir}/ca.crt"         "${gateway_config_dir}/mtls/ca.crt"
  cp "${pki_dir}/client/tls.crt" "${gateway_config_dir}/mtls/tls.crt"
  cp "${pki_dir}/client/tls.key" "${gateway_config_dir}/mtls/tls.key"

  local oidc_line=""
  if [ -n "${oidc_issuer}" ]; then
    oidc_line="$(printf ',\n  "oidc_issuer": "%s"' "${oidc_issuer}")"
  fi

  cat >"${gateway_config_dir}/metadata.json" <<EOF
{
  "name": "${name}",
  "gateway_endpoint": "${endpoint}",
  "is_remote": false,
  "gateway_port": ${port}${oidc_line}
}
EOF
  printf '%s' "${name}" >"${config_home}/openshell/active_gateway"
}

# Import the example provider profiles at platform scope.
#
# OpenShell compiles no provider profile into a binary, so a freshly started
# gateway serves an empty catalog. The e2e suites exercise providers built from
# github, openai, nvidia and the rest, which means the lane has to import them
# first — the same step the upgrade notes give operators.
e2e_import_example_provider_profiles() {
  local cli_bin=$1
  local root=$2

  echo "Importing example provider profiles from ${root}/providers..."
  if ! "${cli_bin}" provider profile import --from "${root}/providers" --global; then
    echo "ERROR: failed to import example provider profiles" >&2
    return 1
  fi
}

# Register an administrator OIDC session for a gateway, non-interactively.
#
# The browser PKCE flow the CLI normally uses cannot run unattended, so mint an
# admin access token with Keycloak's password grant and write the same token
# bundle `openshell gateway login` would have stored. Used only to establish a
# setup identity; the tests themselves still authenticate however they choose.
e2e_register_oidc_admin_session() {
  local config_home=$1
  local name=$2
  local endpoint=$3
  local port=$4
  local issuer=$5
  local username=$6
  local password=$7
  local pki_dir=$8
  local cli_bin=$9
  local client_id="${10:-openshell-cli}"
  local gateway_config_dir="${config_home}/openshell/gateways/${name}"

  # The OpenShell scopes are optional client scopes on openshell-cli, so
  # Keycloak mints them only when they are asked for. Without an explicit
  # scope the token carries the realm defaults alone and every authorized RPC
  # fails with "scope '<name>' required".
  local token
  token=$(curl -sf -X POST "${issuer}/protocol/openid-connect/token" \
    -d "grant_type=password" \
    -d "client_id=${client_id}" \
    -d "username=${username}" \
    -d "password=${password}" \
    --data-urlencode "scope=openid openshell:all" \
    | python3 -c 'import json,sys; print(json.load(sys.stdin)["access_token"])' 2>/dev/null) || true

  if [ -z "${token}" ]; then
    echo "ERROR: could not obtain an admin OIDC token from ${issuer}" >&2
    return 1
  fi

  mkdir -p "${gateway_config_dir}"

  # Trust the gateway's self-signed serving certificate. Only the CA is
  # installed: these lanes start the gateway without --tls-client-ca, and with
  # no client cert/key on disk the CLI falls back to CA-only server
  # verification and authenticates with the bearer token instead of an mTLS
  # identity.
  mkdir -p "${gateway_config_dir}/mtls"
  cp "${pki_dir}/ca.crt" "${gateway_config_dir}/mtls/ca.crt"

  # auth_mode must be "oidc": the CLI dispatches on this field alone when
  # deciding to load a stored bearer token, so omitting it leaves
  # oidc_token.json on disk and unread, and the request goes unauthenticated.
  cat >"${gateway_config_dir}/metadata.json" <<EOF
{
  "name": "${name}",
  "gateway_endpoint": "${endpoint}",
  "is_remote": false,
  "gateway_port": ${port},
  "auth_mode": "oidc",
  "oidc_issuer": "${issuer}",
  "oidc_client_id": "${client_id}",
  "oidc_scopes": "openid openshell:all"
}
EOF
  cat >"${gateway_config_dir}/oidc_token.json" <<EOF
{
  "access_token": "${token}",
  "issuer": "${issuer}",
  "client_id": "${client_id}"
}
EOF
  chmod 600 "${gateway_config_dir}/oidc_token.json"
  printf '%s' "${name}" >"${config_home}/openshell/active_gateway"

  # Assert the CLI reaches the gateway as an authenticated administrator before
  # anything depends on it. ListProviderProfiles is annotated
  # auth_mode: "bearer", so it cannot succeed unless the stored token was
  # loaded and accepted -- this fails here, with the cause named, rather than
  # surfacing later as an opaque profile import error.
  if ! "${cli_bin}" provider list-profiles --global --output json >/dev/null 2>&1; then
    echo "ERROR: the CLI could not make an authenticated call as the OIDC administrator" >&2
    echo "       gateway config: ${gateway_config_dir}" >&2
    echo "       CLI output follows:" >&2
    "${cli_bin}" provider list-profiles --global --output json >&2 || true
    return 1
  fi

  echo "Established an authenticated OIDC administrator session for '${name}'."
}

e2e_toml_string() {
  local value="$1"
  value="${value//\\/\\\\}"
  value="${value//\"/\\\"}"
  printf '"%s"' "${value}"
}

e2e_generate_gateway_jwt() {
  local jwt_dir=$1

  mkdir -p "${jwt_dir}"
  (
    umask 077
    openssl genpkey -algorithm Ed25519 -out "${jwt_dir}/signing.pem" >/dev/null 2>&1
  )
  openssl pkey -in "${jwt_dir}/signing.pem" -pubout -out "${jwt_dir}/public.pem" >/dev/null 2>&1
  openssl rand -hex 16 >"${jwt_dir}/kid"
}

e2e_write_gateway_jwt_config() {
  local jwt_dir=$1
  local gateway_id=$2

  printf '[openshell.gateway.gateway_jwt]\n'
  printf 'signing_key_path = %s\n' "$(e2e_toml_string "${jwt_dir}/signing.pem")"
  printf 'public_key_path = %s\n'  "$(e2e_toml_string "${jwt_dir}/public.pem")"
  printf 'kid_path = %s\n'         "$(e2e_toml_string "${jwt_dir}/kid")"
  printf 'gateway_id = %s\n'       "$(e2e_toml_string "${gateway_id}")"
  # Local Docker/Podman e2e gateways exercise the single-player default:
  # sandbox JWTs identify the supervisor and do not expire.
  printf '\n'
}

e2e_write_gateway_mtls_auth_config() {
  printf '[openshell.gateway.mtls_auth]\n'
  printf 'enabled = true\n\n'
}

e2e_write_gateway_oidc_config() {
  local issuer=$1
  local scopes_claim="${2:-scope}"

  printf '[openshell.gateway.oidc]\n'
  printf 'issuer = %s\n'         "$(e2e_toml_string "${issuer}")"
  case "${issuer}" in
    http://127.*|http://\[::1\]*)
      printf 'dangerously_allow_insecure_http = true\n'
      ;;
  esac
  printf 'audience = "openshell-cli"\n'
  printf 'jwks_ttl_secs = 60\n'
  printf 'roles_claim = "realm_access.roles"\n'
  printf 'admin_role = "openshell-admin"\n'
  printf 'user_role = "openshell-user"\n'
  printf 'scopes_claim = %s\n\n' "$(e2e_toml_string "${scopes_claim}")"
}

e2e_build_gateway_binaries() {
  local root=$1
  local target_var=$2
  local gateway_var=$3
  local cli_var=$4
  local target_dir
  local jobs=()

  if [ -n "${CARGO_BUILD_JOBS:-}" ]; then
    jobs=(-j "${CARGO_BUILD_JOBS}")
  fi

  target_dir="$(e2e_cargo_target_dir "${root}")"
  printf -v "${target_var}" '%s' "${target_dir}"
  printf -v "${gateway_var}" '%s' "${OPENSHELL_GATEWAY_BIN:-${target_dir}/debug/openshell-gateway}"
  printf -v "${cli_var}" '%s' "${OPENSHELL_BIN:-${target_dir}/debug/openshell}"

  if [ -z "${OPENSHELL_GATEWAY_BIN:-}" ]; then
    echo "Building openshell-gateway..."
    if [ "${OPENSHELL_E2E_EXTERNAL_COMPUTE_DRIVER:-0}" = "1" ]; then
      cargo build ${jobs[@]+"${jobs[@]}"} \
        -p openshell-gateway --bin openshell-gateway \
        --no-default-features --features telemetry
    else
      cargo build ${jobs[@]+"${jobs[@]}"} \
        -p openshell-gateway --bin openshell-gateway
    fi
  else
    echo "Using prebuilt openshell gateway at ${OPENSHELL_GATEWAY_BIN}"
  fi

  if [ -z "${OPENSHELL_BIN:-}" ]; then
    echo "Building openshell-cli..."
    cargo build ${jobs[@]+"${jobs[@]}"} \
      -p openshell-cli
  else
    echo "Using prebuilt openshell CLI at ${OPENSHELL_BIN}"
  fi

  if [ ! -x "${!gateway_var}" ]; then
    echo "ERROR: expected openshell-gateway binary at ${!gateway_var}" >&2
    exit 1
  fi
  if [ ! -x "${!cli_var}" ]; then
    echo "ERROR: expected openshell CLI binary at ${!cli_var}" >&2
    exit 1
  fi
}

e2e_build_external_driver() {
  local root=$1
  local package=$2
  local binary=$3
  local output_var=$4
  local target_dir
  local jobs=()

  if [ -n "${CARGO_BUILD_JOBS:-}" ]; then
    jobs=(-j "${CARGO_BUILD_JOBS}")
  fi
  target_dir="$(e2e_cargo_target_dir "${root}")"
  if [ -n "${OPENSHELL_EXTERNAL_DRIVER_BIN:-}" ]; then
    printf -v "${output_var}" '%s' "${OPENSHELL_EXTERNAL_DRIVER_BIN}"
    echo "Using prebuilt external driver at ${OPENSHELL_EXTERNAL_DRIVER_BIN}"
  else
    printf -v "${output_var}" '%s' "${target_dir}/debug/${binary}"
    echo "Building external ${binary}..."
    cargo build ${jobs[@]+"${jobs[@]}"} -p "${package}" --bin "${binary}"
  fi
  if [ ! -x "${!output_var}" ]; then
    echo "ERROR: expected external driver binary at ${!output_var}" >&2
    exit 1
  fi
}

e2e_wait_for_socket() {
  local socket_path=$1
  local process_pid=$2
  local process_label=$3
  local timeout="${4:-30}"
  local elapsed=0

  while [ "${elapsed}" -lt "${timeout}" ]; do
    if [ -S "${socket_path}" ]; then
      return 0
    fi
    if ! kill -0 "${process_pid}" 2>/dev/null; then
      echo "ERROR: ${process_label} exited before creating ${socket_path}" >&2
      return 1
    fi
    sleep 1
    elapsed=$((elapsed + 1))
  done
  echo "ERROR: ${process_label} did not create ${socket_path} within ${timeout}s" >&2
  return 1
}

e2e_stop_process() {
  local process_pid=$1
  local process_label=$2
  if [ -n "${process_pid}" ] && kill -0 "${process_pid}" 2>/dev/null; then
    echo "Stopping ${process_label} (pid ${process_pid})..."
    kill "${process_pid}" 2>/dev/null || true
    wait "${process_pid}" 2>/dev/null || true
  fi
}

e2e_write_gateway_args_file() {
  local args_file=$1
  shift

  : >"${args_file}"
  for arg in "$@"; do
    printf '%s\0' "${arg}" >>"${args_file}"
  done
}

e2e_export_gateway_restart_metadata() {
  local gateway_bin=$1
  local args_file=$2
  local log_file=$3
  local pid_file=$4

  export OPENSHELL_E2E_GATEWAY_BIN="${gateway_bin}"
  export OPENSHELL_E2E_GATEWAY_ARGS_FILE="${args_file}"
  export OPENSHELL_E2E_GATEWAY_LOG="${log_file}"
  export OPENSHELL_E2E_GATEWAY_PID_FILE="${pid_file}"
}

e2e_stop_gateway() {
  local gateway_pid=$1
  local gateway_pid_file=$2

  if [ -f "${gateway_pid_file}" ]; then
    gateway_pid="$(cat "${gateway_pid_file}" 2>/dev/null || true)"
  fi
  if [ -n "${gateway_pid}" ] && kill -0 "${gateway_pid}" 2>/dev/null; then
    echo "Stopping openshell-gateway (pid ${gateway_pid})..."
    kill "${gateway_pid}" 2>/dev/null || true

    # A Rust E2E test may have restarted the gateway and updated the PID file.
    # That replacement process is not a child of this shell, so `wait` returns
    # immediately even though gateway shutdown (including its sandbox stop
    # sweep) is still in progress. Poll until either the process exits or a
    # child process becomes a zombie that the final `wait` can reap.
    local attempts=0
    local process_state=""
    while kill -0 "${gateway_pid}" 2>/dev/null && [ "${attempts}" -lt 120 ]; do
      process_state="$(ps -p "${gateway_pid}" -o stat= 2>/dev/null || true)"
      case "${process_state}" in
        *Z*) break ;;
      esac
      sleep 0.5
      attempts=$((attempts + 1))
    done
    if kill -0 "${gateway_pid}" 2>/dev/null; then
      process_state="$(ps -p "${gateway_pid}" -o stat= 2>/dev/null || true)"
      case "${process_state}" in
        *Z*) ;;
        *) kill -KILL "${gateway_pid}" 2>/dev/null || true ;;
      esac
    fi
    wait "${gateway_pid}" 2>/dev/null || true
  fi
}

e2e_print_gateway_log_on_failure() {
  local exit_code=$1
  local gateway_log=$2

  if [ "${exit_code}" -ne 0 ] && [ -f "${gateway_log}" ]; then
    echo "=== gateway log (preserved for debugging) ==="
    cat "${gateway_log}"
    echo "=== end gateway log ==="
  fi
}