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

# Wait until Docker Compose services pass their health checks.
# Read-only: does not start, stop, or recreate any containers.
# Used by phrocs processes to block until infrastructure is ready.
#
# Dozens of phrocs panes run this loop concurrently, so each poll must be
# cheap: readiness is read via plain `docker ps` (label-filtered), not
# `docker compose ps`, which re-parses and re-interpolates every compose file
# on each invocation (~0.5 CPU-seconds per call). While the docker daemon is
# unreachable the loop backs off to a slow poll instead of spinning.
#
# Usage:
#   bin/wait-for-docker                     # wait for core (db redis7 kafka clickhouse)
#   bin/wait-for-docker temporal            # wait for core + temporal
#   bin/wait-for-docker --only db redis7    # wait only for db and redis7 (skip core)
#
# Stays bash 3.2 compatible (the version macOS ships as /bin/bash): no
# associative arrays, no empty-array expansion under `set -u`. Per-service
# state is kept in parallel indexed arrays scanned by hand — the service
# count is tiny, so the linear scans are free.

# In a sandbox container, infra is managed by docker compose externally
# and is guaranteed healthy via depends_on health checks.
if [[ "${POSTHOG_SANDBOX:-}" == "1" ]]; then
    echo "Sandbox: infrastructure already running"
    exit 0
fi

SCRIPT_DIR="$(CDPATH='' cd -- "$(dirname -- "$0")" && pwd)"
REPO_ROOT="$(dirname "$SCRIPT_DIR")"
DEFAULT_COMPOSE_FILE="${REPO_ROOT}/docker-compose.dev.yml"
TIMEOUT="${WAIT_FOR_DOCKER_TIMEOUT:-300}"
DAEMON_DOWN_POLL_SECONDS=2

# Pin the project so polls target the same stack the up command created under
# (see generator.py). Without -p, compose derives it from COMPOSE_PROJECT_NAME
# or the worktree folder name, which drifts from "posthog" when direnv is off.
PROJECT="${COMPOSE_PROJECT_NAME:-posthog}"

only_listed=0
if [ "${1:-}" = "--only" ]; then
    only_listed=1
    shift
fi

if [ "$only_listed" -eq 1 ]; then
    if [ "$#" -eq 0 ]; then
        echo "wait-for-docker: --only requires at least one service argument" >&2
        exit 2
    fi
    SERVICES=("$@")
else
    SERVICES=(db redis7 kafka clickhouse)
    SERVICES+=("$@")
fi

# Dedup while preserving order via a linear membership scan.
unique_services=()
for svc in "${SERVICES[@]}"; do
    already_seen=0
    for kept in ${unique_services[@]+"${unique_services[@]}"}; do
        if [ "$kept" = "$svc" ]; then
            already_seen=1
            break
        fi
    done
    if [ "$already_seen" -eq 0 ]; then
        unique_services+=("$svc")
    fi
done
SERVICES=("${unique_services[@]}")

compose() {
    local compose_files=()
    local compose_args=()
    local compose_file
    local separator

    if [ -n "${COMPOSE_FILE:-}" ]; then
        separator="${COMPOSE_PATH_SEPARATOR:-:}"
        IFS="$separator" read -r -a compose_files <<< "$COMPOSE_FILE"
        for compose_file in "${compose_files[@]}"; do
            if [ -n "$compose_file" ]; then
                compose_args+=(-f "$compose_file")
            fi
        done
        env -u COMPOSE_FILE docker compose -p "$PROJECT" ${compose_args[@]+"${compose_args[@]}"} "$@"
    else
        docker compose -p "$PROJECT" -f "$DEFAULT_COMPOSE_FILE" "$@"
    fi
}

# One readiness poll. The label filters reproduce compose's view of the
# project (oneoff=False excludes `compose run` containers, whose rows would
# shadow the real service in the last-write-wins scan below). `-a` keeps
# exited containers visible so timeout diagnostics can show their logs.
# Fails when the daemon is unreachable, which the main loop uses to back off.
# stderr is captured by the caller: a failure isn't necessarily a stopped
# daemon (bad context, socket permissions), so the real error is reported.
poll_containers() {
    docker ps -a \
        --filter "label=com.docker.compose.project=${PROJECT}" \
        --filter "label=com.docker.compose.oneoff=False" \
        --format '{{.Label "com.docker.compose.service"}}|{{.Status}}|{{.State}}|{{.ID}}'
}

readiness_status() {
    local health="${1:-}"
    local state="${2:-}"

    if [ -n "$health" ]; then
        echo "$health"
    elif [ -n "$state" ]; then
        echo "$state"
    else
        echo "missing"
    fi
}

service_is_ready() {
    [ "$1" = "healthy" ] || [ "$1" = "running" ]
}

print_service_diagnostics() {
    local svc="$1"
    local elapsed="$2"
    local status="$3"
    local container_id="$4"

    echo ""
    echo "=========================================="
    echo "TIMEOUT: ${TIMEOUT}s waiting for $svc"
    echo "Last observed state after ${elapsed}s: ${status}"
    echo "=========================================="
    if [ -z "$container_id" ]; then
        echo "Container for '$svc' does not exist yet."
        echo ""
        echo "docker compose ps $svc:"
        compose ps "$svc" 2>&1 || true
    else
        echo "Container exists (id=${container_id:0:12}) but readiness state is: $status"
        echo ""
        echo "Last 100 lines of container logs:"
        docker logs --tail 100 "$container_id" 2>&1 || true
        echo ""
        echo "Health check log:"
        docker inspect --format='{{if .State.Health}}{{range .State.Health.Log}}{{.Output}}{{end}}{{else}}no health check configured{{end}}' "$container_id" 2>&1 | tail -5 || true
    fi
    echo "=========================================="
}

POLL_ERR_FILE="$(mktemp "${TMPDIR:-/tmp}/wait-for-docker-err.XXXXXX")"
trap 'rm -f "$POLL_ERR_FILE"' EXIT

SECONDS=0
last_progress_log=-5
daemon_down=0
poll_error=""
pending=("${SERVICES[@]}")
# Services that time out, captured at the moment of timeout as
# "svc|status|container_id" (this iteration's status is the last observed).
timed_out=()

while [ ${#pending[@]} -gt 0 ]; do
    ps_output=""
    if ! ps_output=$(poll_containers 2>"$POLL_ERR_FILE"); then
        poll_error="$(head -n 1 "$POLL_ERR_FILE" 2>/dev/null || true)"
        # Nothing can converge until docker is back, so poll slowly.
        if [ "$SECONDS" -ge "$TIMEOUT" ]; then
            echo ""
            echo "=========================================="
            echo "TIMEOUT: ${TIMEOUT}s waiting for the Docker daemon"
            echo "=========================================="
            echo "docker ps kept failing with: ${poll_error:-no error output}"
            echo "If the daemon isn't running, start Docker (OrbStack, Docker Desktop, colima, ...) and retry."
            exit 1
        fi
        if [ "$daemon_down" -eq 0 ] || [ $(( SECONDS - last_progress_log )) -ge 10 ]; then
            echo "Docker is not reachable: ${poll_error:-no error output}"
            echo "If the daemon isn't running, start it (OrbStack, Docker Desktop, colima, ...); retrying... (${SECONDS}s)"
            last_progress_log=$SECONDS
            daemon_down=1
        fi
        sleep "$DAEMON_DOWN_POLL_SECONDS"
        continue
    fi
    if [ "$daemon_down" -eq 1 ]; then
        echo "Docker daemon is reachable after ${SECONDS}s."
        daemon_down=0
    fi

    # Parse this poll into parallel arrays keyed by ps_keys.
    ps_keys=()
    ps_health=()
    ps_state=()
    ps_id=()
    while IFS='|' read -r svc status_text state container_id; do
        [ -z "$svc" ] && continue
        # docker ps has no Health column; health is embedded in the status
        # text, e.g. "Up 3 minutes (healthy)" / "Up 2 seconds (health: starting)".
        case "$status_text" in
            *"(healthy)"*) health="healthy" ;;
            *"(health: starting)"*) health="starting" ;;
            *"(unhealthy)"*) health="unhealthy" ;;
            *) health="" ;;
        esac
        ps_keys+=("$svc")
        ps_health+=("$health")
        ps_state+=("$state")
        ps_id+=("$container_id")
    done <<< "$ps_output"

    still_pending=()
    pending_summary=()
    elapsed=$SECONDS

    for svc in "${pending[@]}"; do
        # Look up this service's row from the parsed poll output.
        health=""
        state=""
        container_id=""
        key_count=${#ps_keys[@]}
        i=0
        while [ "$i" -lt "$key_count" ]; do
            # No early break: a scaled service emits one ps row per replica, so
            # keep the last match to preserve the prior map's last-write-wins.
            if [ "${ps_keys[$i]}" = "$svc" ]; then
                health="${ps_health[$i]}"
                state="${ps_state[$i]}"
                container_id="${ps_id[$i]}"
            fi
            i=$((i + 1))
        done

        status=$(readiness_status "$health" "$state")

        if service_is_ready "$status"; then
            echo "$svc is ready after ${elapsed}s."
            continue
        fi

        if [ "$elapsed" -ge "$TIMEOUT" ]; then
            timed_out+=("$svc|$status|$container_id")
            continue
        fi

        still_pending+=("$svc")
        pending_summary+=("$svc=${status}(${elapsed}s)")
    done

    pending=(${still_pending[@]+"${still_pending[@]}"})

    if [ ${#pending[@]} -eq 0 ]; then
        break
    fi

    if [ $(( SECONDS - last_progress_log )) -ge 5 ]; then
        echo "Waiting for: ${pending_summary[*]}"
        last_progress_log=$SECONDS
    fi

    # Fast polling keeps convergence latency low for every pane gated on this
    # script, and is affordable because `docker ps` is a thin API call.
    sleep 0.2
done

if [ ${#timed_out[@]} -gt 0 ]; then
    for entry in "${timed_out[@]}"; do
        IFS='|' read -r svc status container_id <<< "$entry"
        print_service_diagnostics "$svc" "$SECONDS" "$status" "$container_id"
    done
    exit 1
fi

echo "All services ready in ${SECONDS}s."
