#!/bin/sh
set -eu

compose_file=${COMPOSE_FILE:-deploy/docker/docker-compose.yml}
duration=${KESTRA_SOAK_SECONDS:-1800}
interval=${KESTRA_SOAK_INTERVAL_SECONDS:-15}
base_url=${KESTRA_SOAK_BASE_URL:-http://127.0.0.1:18080/api/v1/main}
username=${KESTRA_USERNAME:-admin@dataops.local}
password=${KESTRA_PASSWORD:-DataOpsKestra1!}
evidence=${KESTRA_SOAK_EVIDENCE:-/tmp/dataops-kestra-soak-$(date +%Y%m%dT%H%M%S).jsonl}
namespace=dataops_validation
flow_id=queue_restart_soak_$(date +%s)
deadline=$(( $(date +%s) + duration ))
restart_at=$(( $(date +%s) + duration / 3 ))
restarted=false
post_restart_since=
soak_started_at=$(date -u +%Y-%m-%dT%H:%M:%SZ)
checks=0
executions=0
postgres_container=$(docker compose -f "$compose_file" ps -q postgres)
kestra_container=$(docker compose -f "$compose_file" ps -q kestra)
postgres_restart_count_before=$(docker inspect -f '{{.RestartCount}}' "$postgres_container")
kestra_restart_count_before=$(docker inspect -f '{{.RestartCount}}' "$kestra_container")

jq -cn \
  --arg started_at "$soak_started_at" \
  --argjson duration_seconds "$duration" \
  --argjson postgres_restart_count "$postgres_restart_count_before" \
  --argjson kestra_restart_count "$kestra_restart_count_before" \
  '{kind:"soak_start",started_at:$started_at,
    duration_seconds:$duration_seconds,
    postgres_restart_count:$postgres_restart_count,
    kestra_restart_count:$kestra_restart_count}' >"$evidence"

flow_file=$(mktemp "${TMPDIR:-/tmp}/dataops-kestra-soak.XXXXXX.yml")
trap 'rm -f "$flow_file"' EXIT
printf '%s\n' \
  "id: ${flow_id}" \
  "namespace: ${namespace}" \
  "description: DataOps controlled queue and restart soak" \
  "disabled: false" \
  "tasks:" \
  "  - id: queue_probe" \
  "    type: io.kestra.plugin.core.log.Log" \
  "    message: DataOps queue soak probe" >"$flow_file"

api() {
  method=$1
  path=$2
  shift 2
  curl -fsS --user "${username}:${password}" \
    -X "$method" "${base_url}/${path}" "$@"
}

api POST flows -H 'Content-Type: application/x-yaml' \
  --data-binary "@${flow_file}" >/dev/null

while [ "$(date +%s)" -lt "$deadline" ]; do
  sample_time=$(date -u +%Y-%m-%dT%H:%M:%SZ)
  execution=$(api POST "executions/${namespace}/${flow_id}")
  execution_id=$(printf '%s' "$execution" | jq -er '.id')
  state=UNKNOWN
  poll=0
  while [ "$poll" -lt 30 ]; do
    state=$(api GET "executions/${execution_id}" |
      jq -r '.state.current // .state // "UNKNOWN"')
    case "$state" in
      SUCCESS|FAILED|KILLED|CANCELLED|WARNING) break ;;
    esac
    poll=$((poll + 1))
    sleep 1
  done
  [ "$state" = "SUCCESS" ] || {
    printf 'Kestra soak execution %s ended as %s\n' "$execution_id" "$state" >&2
    exit 1
  }
  executions=$((executions + 1))

  docker compose -f "$compose_file" exec -T kestra \
    curl -fsS http://127.0.0.1:8081/health >/dev/null
  health='{"status":"UP"}'
  jq -cn \
    --arg kind "sample" \
    --arg sampled_at "$sample_time" \
    --arg execution_id "$execution_id" \
    --arg state "$state" \
    --argjson health "$health" \
    --argjson postgres_restarted "$restarted" \
    '{kind:$kind,sampled_at:$sampled_at,execution_id:$execution_id,state:$state,
      health:$health,controlled_restart_completed:$postgres_restarted}' \
    >>"$evidence"
  checks=$((checks + 1))

  if [ "$restarted" = false ] && [ "$(date +%s)" -ge "$restart_at" ]; then
    docker compose -f "$compose_file" restart postgres
    ready_attempt=0
    until docker compose -f "$compose_file" exec -T postgres \
      pg_isready -U dataops -d dataops >/dev/null 2>&1; do
      ready_attempt=$((ready_attempt + 1))
      [ "$ready_attempt" -lt 30 ] || {
        echo "PostgreSQL did not recover during controlled restart" >&2
        exit 1
      }
      sleep 2
    done
    docker compose -f "$compose_file" run --rm -T kestra-db-init
    docker compose -f "$compose_file" restart kestra
    docker compose -f "$compose_file" up -d --wait kestra
    post_restart_since=$(date -u +%Y-%m-%dT%H:%M:%SZ)
    restarted=true
  fi
  sleep "$interval"
done

[ "$restarted" = true ] || {
  echo "Controlled PostgreSQL/Kestra restart did not run" >&2
  exit 1
}
postgres_restart_count_after=$(docker inspect -f '{{.RestartCount}}' "$postgres_container")
kestra_restart_count_after=$(docker inspect -f '{{.RestartCount}}' "$kestra_container")
# Docker's explicit `docker restart` does not consume RestartCount; that field
# counts daemon/restart-policy recoveries. Any delta therefore proves an
# uncontrolled extra restart during the observation window.
if [ "$postgres_restart_count_after" != "$postgres_restart_count_before" ] ||
  [ "$kestra_restart_count_after" != "$kestra_restart_count_before" ]; then
  echo "Unexpected automatic container restart during soak" >&2
  exit 1
fi
log_file=$(mktemp "${TMPDIR:-/tmp}/dataops-kestra-soak-logs.XXXXXX")
docker compose -f "$compose_file" logs --since "$soak_started_at" kestra >"$log_file"
if grep -Eiq 'outofmemory' "$log_file"; then
  echo "Kestra emitted a fatal database/queue signature during soak" >&2
  exit 1
fi
# PostgreSQL emits SQLSTATE 57P01 / "terminating connection due to
# administrator command" for the one deliberate database restart above. It is
# retained and counted as controlled evidence, while every other FATAL line
# across the full observation window remains a release-gate failure.
if grep -Ei '(^|[^a-z])fatal([^a-z]|$)' "$log_file" |
  grep -Eiv 'fatal: terminating connection due to administrator command' \
    >/dev/null; then
  echo "Kestra emitted an unexpected fatal database/queue signature" >&2
  exit 1
fi
known_restart_disconnects=$(grep -Eic \
  'broken pipe|connection is closed|connection reset|hikaripool.*failed|terminating connection due to administrator command' \
  "$log_file" || true)
soak_completed_at=$(date -u +%Y-%m-%dT%H:%M:%SZ)
jq -cn \
  --arg completed_at "$soak_completed_at" \
  --arg log_since "$soak_started_at" \
  --argjson executions "$executions" \
  --argjson samples "$checks" \
  --argjson known_restart_disconnects "$known_restart_disconnects" \
  --argjson postgres_restart_count "$postgres_restart_count_after" \
  --argjson kestra_restart_count "$kestra_restart_count_after" \
  '{kind:"soak_complete",completed_at:$completed_at,log_since:$log_since,
    executions:$executions,samples:$samples,
    known_controlled_restart_disconnects:$known_restart_disconnects,
    postgres_restart_count:$postgres_restart_count,
    kestra_restart_count:$kestra_restart_count}' >>"$evidence"
rm -f "$log_file"

printf 'Kestra soak passed: %s executions, %s samples, evidence=%s\n' \
  "$executions" "$checks" "$evidence"