#!/usr/bin/env bash

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
COMPOSE_FILE="${COMPOSE_FILE:-${SCRIPT_DIR}/../docker-compose.yml}"

BASE_URL="${BASE_URL:-http://localhost:5270}"
PGHOST="${PGHOST:-localhost}"
PGPORT="${PGPORT:-5436}"
PGDATABASE="${PGDATABASE:-vigilcare}"
PGUSER="${PGUSER:-postgres}"
PGPASSWORD="${PGPASSWORD:-password}"

TOPIC_OBSERVATION="${TOPIC_OBSERVATION:-observation.recorded}"
TOPIC_ENCOUNTER="${TOPIC_ENCOUNTER:-encounter.status.changed}"
RELAY_WAIT_SECS="${RELAY_WAIT_SECS:-45}"
KAFKA_READY_WAIT_SECS="${KAFKA_READY_WAIT_SECS:-30}"

RECORDED_AT="$(date -u +"%Y-%m-%dT%H:%M:%SZ")"
SCRIPT_RUN_ID="$(date -u +"%Y%m%d%H%M%S")"

TMP_FILES=()

cleanup() {
  local f
  for f in "${TMP_FILES[@]}"; do
    rm -f "${f}" "${f}.status" 2>/dev/null || true
  done
}
trap cleanup EXIT

if ! command -v curl >/dev/null 2>&1; then
  echo "Missing dependency: curl"
  exit 1
fi

if ! command -v jq >/dev/null 2>&1; then
  echo "Missing dependency: jq"
  exit 1
fi

if ! command -v docker >/dev/null 2>&1 || [[ ! -f "${COMPOSE_FILE}" ]]; then
  echo "Missing dependency: docker compose (${COMPOSE_FILE})"
  exit 1
fi

compose() {
  docker compose -f "${COMPOSE_FILE}" "$@"
}

kafka_exec() {
  compose exec -T kafka "$@"
}

psql_cmd() {
  local sql="$1"
  if command -v psql >/dev/null 2>&1; then
    PGPASSWORD="${PGPASSWORD}" psql -h "${PGHOST}" -p "${PGPORT}" -U "${PGUSER}" -d "${PGDATABASE}" -tAc "${sql}"
  else
    compose exec -T postgres psql -U "${PGUSER}" -d "${PGDATABASE}" -tAc "${sql}"
  fi
}

request() {
  local method="$1"
  local url="$2"
  local body="${3:-}"
  local tmp_body
  tmp_body="$(mktemp)"
  TMP_FILES+=("${tmp_body}")
  local status

  if [[ -n "${body}" ]]; then
    status="$(curl -sS -o "${tmp_body}" -w "%{http_code}" -X "${method}" "${url}" \
      -H "Content-Type: application/json" -d "${body}")"
  else
    status="$(curl -sS -o "${tmp_body}" -w "%{http_code}" -X "${method}" "${url}")"
  fi

  echo "${status}" > "${tmp_body}.status"
  echo "${tmp_body}"
}

assert_status() {
  local expected="$1"
  local body_file="$2"
  local status
  status="$(cat "${body_file}.status")"
  if [[ "${status}" != "${expected}" ]]; then
    echo "Expected HTTP ${expected}, got ${status}"
    echo "Response body:"
    cat "${body_file}"
    echo
    return 1
  fi
}

wait_for_kafka() {
  local elapsed=0
  while (( elapsed < KAFKA_READY_WAIT_SECS )); do
    if kafka_exec /opt/kafka/bin/kafka-broker-api-versions.sh --bootstrap-server localhost:9092 >/dev/null 2>&1; then
      return 0
    fi
    sleep 2
    elapsed=$((elapsed + 2))
  done
  echo "Kafka did not become ready within ${KAFKA_READY_WAIT_SECS}s"
  return 1
}

wait_for_outbox_processed() {
  local outbox_id="$1"
  local elapsed=0
  while (( elapsed < RELAY_WAIT_SECS )); do
    local processed
    processed="$(psql_cmd "SELECT processed_at IS NOT NULL FROM outbox_events WHERE id = '${outbox_id}'")"
    if [[ "${processed}" == "t" ]]; then
      return 0
    fi
    sleep 1
    elapsed=$((elapsed + 1))
  done
  echo "Outbox row ${outbox_id} was not processed within ${RELAY_WAIT_SECS}s"
  return 1
}

wait_for_pending_outbox_count() {
  local expected="$1"
  local elapsed=0
  while (( elapsed < RELAY_WAIT_SECS )); do
    local pending
    pending="$(psql_cmd "SELECT COUNT(*) FROM outbox_events WHERE processed_at IS NULL")"
    if [[ "${pending}" == "${expected}" ]]; then
      return 0
    fi
    sleep 1
    elapsed=$((elapsed + 1))
  done
  echo "Expected ${expected} pending outbox row(s), timed out after ${RELAY_WAIT_SECS}s"
  return 1
}

consume_topic() {
  local topic="$1"
  local group="$2"
  local max_messages="$3"
  local tmp_out
  tmp_out="$(mktemp)"
  TMP_FILES+=("${tmp_out}")

  kafka_exec /opt/kafka/bin/kafka-console-consumer.sh \
    --bootstrap-server localhost:9092 \
    --topic "${topic}" \
    --group "${group}" \
    --from-beginning \
    --property print.key=true \
    --property print.partition=true \
    --timeout-ms 15000 \
    --max-messages "${max_messages}" > "${tmp_out}" 2>/dev/null || true

  echo "${tmp_out}"
}

TOTAL_STEPS=11

echo "Running Kafka + outbox verification against ${BASE_URL}"
echo "Script run id: ${SCRIPT_RUN_ID}"

echo ""
echo "[0/${TOTAL_STEPS}] Preflight — API, Postgres, and Kafka reachable"
preflight_status="$(curl -sS -o /dev/null -w "%{http_code}" "${BASE_URL}/api/v1/alert-thresholds" || true)"
if [[ "${preflight_status}" != "200" ]]; then
  echo "API not reachable at ${BASE_URL} (HTTP ${preflight_status})."
  echo "Start the API with: dotnet run --project VigilCareClinicalAPI"
  exit 1
fi

if ! psql_cmd "SELECT 1" >/dev/null 2>&1; then
  echo "Postgres not reachable on ${PGHOST}:${PGPORT}."
  echo "Start the stack with: docker compose up -d"
  exit 1
fi

if ! wait_for_kafka; then
  exit 1
fi
echo "OK: API, Postgres, and Kafka are up"

echo ""
echo "[1/${TOTAL_STEPS}] Verifying Kafka topics were provisioned"
topic_list="$(kafka_exec /opt/kafka/bin/kafka-topics.sh --bootstrap-server localhost:9092 --list)"
for topic in alert.generated alert.acknowledged encounter.status.changed observation.recorded; do
  if ! grep -qx "${topic}" <<< "${topic_list}"; then
    echo "Missing Kafka topic: ${topic}"
    echo "Topics found:"
    echo "${topic_list}"
    exit 1
  fi
done
echo "OK: observation.recorded, alert.generated, alert.acknowledged, encounter.status.changed exist"

echo ""
echo "[2/${TOTAL_STEPS}] Creating patient and two active encounters"
patient_payload='{"firstName":"Kafka","lastName":"Verifier","dateOfBirth":"1990-06-01","gender":"M"}'
resp="$(request POST "${BASE_URL}/api/v1/patients" "${patient_payload}")"
assert_status "201" "${resp}"
patient_id="$(jq -r '.data.id' "${resp}")"

enc_payload_a='{"encounterType":"Inpatient","department":"ICU","attendingPhysician":"Dr. Kafka"}'
resp="$(request POST "${BASE_URL}/api/v1/patients/${patient_id}/encounters" "${enc_payload_a}")"
assert_status "201" "${resp}"
encounter_a="$(jq -r '.data.id' "${resp}")"

enc_payload_b='{"encounterType":"Outpatient","department":"Clinic","attendingPhysician":"Dr. Kafka"}'
resp="$(request POST "${BASE_URL}/api/v1/patients/${patient_id}/encounters" "${enc_payload_b}")"
assert_status "201" "${resp}"
encounter_b="$(jq -r '.data.id' "${resp}")"

enc_payload_discharge='{"encounterType":"Emergency","department":"ED","attendingPhysician":"Dr. Kafka"}'
resp="$(request POST "${BASE_URL}/api/v1/patients/${patient_id}/encounters" "${enc_payload_discharge}")"
assert_status "201" "${resp}"
encounter_discharge="$(jq -r '.data.id' "${resp}")"

echo "OK: encounters ${encounter_a}, ${encounter_b}, ${encounter_discharge}"

echo ""
echo "[3/${TOTAL_STEPS}] Ingesting observations for partition-key verification"
for encounter_id in "${encounter_a}" "${encounter_b}"; do
  for i in 1 2; do
    obs_payload="$(jq -nc \
      --arg recordedAt "${RECORDED_AT}" \
      --arg key "kafka-script-${SCRIPT_RUN_ID}-${encounter_id}-${i}" \
      '{observations:[{observationCode:"HEART_RATE",value:80,unit:"bpm",source:"DEVICE",recordedAt:$recordedAt,idempotencyKey:$key}]}')"
    resp="$(request POST "${BASE_URL}/api/v1/encounters/${encounter_id}/observations" "${obs_payload}")"
    assert_status "201" "${resp}"
  done
done
echo "OK: four observation.recorded outbox rows written"

echo ""
echo "[4/${TOTAL_STEPS}] Waiting for relay to publish observation events"
latest_outbox="$(psql_cmd "SELECT id FROM outbox_events WHERE topic = '${TOPIC_OBSERVATION}' AND partition_key = '${encounter_a}' ORDER BY created_at DESC LIMIT 1")"
if [[ -z "${latest_outbox}" ]]; then
  echo "Could not find outbox row for encounter ${encounter_a}"
  exit 1
fi
wait_for_outbox_processed "${latest_outbox}"
echo "OK: relay marked observation outbox rows processed"

echo ""
echo "[5/${TOTAL_STEPS}] Verifying partition ordering (same encounter key → same partition)"
consumer_out="$(consume_topic "${TOPIC_OBSERVATION}" "partition-check-${SCRIPT_RUN_ID}" 50)"
for encounter_id in "${encounter_a}" "${encounter_b}"; do
  partitions="$(grep -F "${encounter_id}" "${consumer_out}" | sed -n 's/^Partition:\([0-9]*\).*/\1/p' | sort -u)"
  partition_count="$(echo "${partitions}" | grep -c . || true)"
  if [[ "${partition_count}" -lt 1 ]]; then
    echo "No Kafka messages found for encounter ${encounter_id}"
    exit 1
  fi
  if [[ "${partition_count}" -gt 1 ]]; then
    echo "Encounter ${encounter_id} landed on multiple partitions: ${partitions}"
    exit 1
  fi
done
echo "OK: each encounter's messages share a single partition"

echo ""
echo "[6/${TOTAL_STEPS}] Verifying consumer group independence"
group_a="es-indexer-${SCRIPT_RUN_ID}"
group_b="sepsis-engine-${SCRIPT_RUN_ID}"
out_a="$(consume_topic "${TOPIC_OBSERVATION}" "${group_a}" 20)"
out_b="$(consume_topic "${TOPIC_OBSERVATION}" "${group_b}" 20)"
count_a="$(grep -c '^Partition:' "${out_a}" || true)"
count_b="$(grep -c '^Partition:' "${out_b}" || true)"
if [[ "${count_a}" -lt 4 || "${count_b}" -lt 4 ]]; then
  echo "Expected each consumer group to read at least 4 messages, got ${count_a} and ${count_b}"
  exit 1
fi

offset_a="$(kafka_exec /opt/kafka/bin/kafka-consumer-groups.sh \
  --bootstrap-server localhost:9092 \
  --describe --group "${group_a}" 2>/dev/null | awk '/observation.recorded/ {sum += $4} END {print sum+0}')"
offset_b="$(kafka_exec /opt/kafka/bin/kafka-consumer-groups.sh \
  --bootstrap-server localhost:9092 \
  --describe --group "${group_b}" 2>/dev/null | awk '/observation.recorded/ {sum += $4} END {print sum+0}')"
if [[ "${offset_a}" -lt 4 || "${offset_b}" -lt 4 ]]; then
  echo "Consumer groups did not commit expected offsets (es-indexer=${offset_a}, sepsis-engine=${offset_b})"
  exit 1
fi
echo "OK: two consumer groups read independently (offsets es-indexer=${offset_a}, sepsis-engine=${offset_b})"

echo ""
echo "[7/${TOTAL_STEPS}] Verifying encounter.status.changed outbox write and Kafka publish"
discharge_payload='{"status":"Discharged"}'
resp="$(request PATCH "${BASE_URL}/api/v1/encounters/${encounter_discharge}/status" "${discharge_payload}")"
assert_status "200" "${resp}"

outbox_status="$(psql_cmd "SELECT payload->>'newStatus' FROM outbox_events WHERE topic = '${TOPIC_ENCOUNTER}' AND partition_key = '${encounter_discharge}' ORDER BY created_at DESC LIMIT 1")"
if [[ "${outbox_status}" != "DISCHARGED" ]]; then
  echo "Expected encounter.status.changed outbox row with newStatus=DISCHARGED, got '${outbox_status}'"
  exit 1
fi

discharge_outbox_id="$(psql_cmd "SELECT id FROM outbox_events WHERE topic = '${TOPIC_ENCOUNTER}' AND partition_key = '${encounter_discharge}' ORDER BY created_at DESC LIMIT 1")"
wait_for_outbox_processed "${discharge_outbox_id}"

encounter_consumer_out="$(consume_topic "${TOPIC_ENCOUNTER}" "encounter-check-${SCRIPT_RUN_ID}" 5)"
if ! grep -Fq "${encounter_discharge}" "${encounter_consumer_out}"; then
  echo "encounter.status.changed message not found on Kafka for ${encounter_discharge}"
  exit 1
fi
echo "OK: discharge produced outbox row and Kafka message"

echo ""
echo "[8/${TOTAL_STEPS}] Verifying relay survives Kafka restart (no rows lost)"
compose stop kafka >/dev/null

restart_obs_payload="$(jq -nc \
  --arg recordedAt "${RECORDED_AT}" \
  --arg key "kafka-restart-${SCRIPT_RUN_ID}" \
  '{observations:[{observationCode:"HEART_RATE",value:83,unit:"bpm",source:"DEVICE",recordedAt:$recordedAt,idempotencyKey:$key}]}')"
resp="$(request POST "${BASE_URL}/api/v1/encounters/${encounter_a}/observations" "${restart_obs_payload}")"
assert_status "201" "${resp}"

restart_outbox_id="$(psql_cmd "SELECT id FROM outbox_events WHERE topic = '${TOPIC_OBSERVATION}' AND partition_key = '${encounter_a}' ORDER BY created_at DESC LIMIT 1")"
restart_pending="$(psql_cmd "SELECT processed_at IS NULL FROM outbox_events WHERE id = '${restart_outbox_id}'")"
if [[ "${restart_pending}" != "t" ]]; then
  echo "Expected outbox row ${restart_outbox_id} to remain pending while Kafka is down"
  exit 1
fi
echo "OK: outbox row pending while Kafka is stopped"

compose start kafka >/dev/null
wait_for_kafka
wait_for_outbox_processed "${restart_outbox_id}"
echo "OK: relay caught up after Kafka restart"

echo ""
echo "[9/${TOTAL_STEPS}] Verifying partial index and relay poll plan"
index_count="$(psql_cmd "SELECT COUNT(*) FROM pg_indexes WHERE tablename = 'outbox_events' AND indexdef LIKE '%processed_at IS NULL%'")"
if [[ "${index_count}" -lt 1 ]]; then
  echo "Expected partial index on outbox_events (processed_at IS NULL), found ${index_count}"
  exit 1
fi

psql_cmd "ANALYZE outbox_events" >/dev/null
explain_out="$(psql_cmd "BEGIN; SET LOCAL enable_seqscan = off; EXPLAIN SELECT id, topic, payload, partition_key FROM outbox_events WHERE processed_at IS NULL ORDER BY created_at ASC LIMIT 100 FOR UPDATE SKIP LOCKED; ROLLBACK;")"
if ! grep -qi 'Index Scan' <<< "${explain_out}"; then
  echo "Expected Index Scan when seqscan is disabled, got:"
  echo "${explain_out}"
  exit 1
fi
echo "OK: partial index exists and backs the relay poll query"

echo ""
echo "[10/${TOTAL_STEPS}] Verifying partition_key is set on clinical outbox rows"
missing_keys="$(psql_cmd "SELECT COUNT(*) FROM outbox_events WHERE created_at >= NOW() - INTERVAL '10 minutes' AND topic IN ('${TOPIC_OBSERVATION}', '${TOPIC_ENCOUNTER}', 'alert.generated', 'alert.acknowledged') AND (partition_key IS NULL OR partition_key = '')")"
if [[ "${missing_keys}" != "0" ]]; then
  echo "Found ${missing_keys} clinical outbox row(s) without partition_key"
  exit 1
fi
echo "OK: all clinical outbox rows have partition_key"

echo ""
echo "All ${TOTAL_STEPS} Kafka + outbox checks passed."