#!/usr/bin/env bash
# Run a saturated rolling-deploy experiment end-to-end with reliable capture.
#
# Sequence: scale up → wait for steady → start tail_state → force deploy →
# wait for convergence → wait for CloudWatch ingestion → pull complete logs.
#
# Usage: ./run_saturated_deploy.sh <desired_count> <results_subdir>
# Example: ./run_saturated_deploy.sh 3 experiment-1-repeat-B

set -euo pipefail

COUNT="${1:?Usage: $0 <desired_count> <results_subdir>}"
RESULTS_SUBDIR="${2:?Usage: $0 <desired_count> <results_subdir>}"

CLUSTER="claude-ecs-task-protection-experiment"
SERVICE="claude-ecs-task-protection-experiment"
LOG_GROUP="/ecs/${SERVICE}"
RESULTS_DIR="$(pwd)/results/${RESULTS_SUBDIR}"

mkdir -p "${RESULTS_DIR}"

echo "[$(date -u +%H:%M:%S)] scaling up to ${COUNT}"
./scripts/set_desired.sh "${COUNT}"

echo "[$(date -u +%H:%M:%S)] waiting for ${COUNT} running..."
until [ "$(aws ecs describe-services --cluster "${CLUSTER}" --services "${SERVICE}" --query 'services[0].runningCount' --output text)" = "${COUNT}" ]; do
    sleep 5
done

# Additional steady-state wait so workers have time to start cycling before
# the deploy trigger. Avoids measuring cold-start noise.
sleep 15

echo "[$(date -u +%H:%M:%S)] ${COUNT} running, starting state capture"
./scripts/tail_state.sh "${RESULTS_DIR}" > "${RESULTS_DIR}/tail_state.log" 2>&1 &
TAIL_STATE_PID=$!

sleep 5

TRIGGER_EPOCH_MS=$(($(date +%s) * 1000))
TRIGGER_ISO=$(date -u +%Y-%m-%dT%H:%M:%SZ)
echo "trigger_time=${TRIGGER_ISO}" > "${RESULTS_DIR}/trigger.txt"
echo "trigger_epoch_ms=${TRIGGER_EPOCH_MS}" >> "${RESULTS_DIR}/trigger.txt"

echo "[$(date -u +%H:%M:%S)] triggering deploy"
./scripts/force_deploy.sh

START_EPOCH=$(date +%s)

echo "[$(date -u +%H:%M:%S)] polling for convergence..."
while true; do
    STATE=$(aws ecs describe-services --cluster "${CLUSTER}" --services "${SERVICE}" --query 'services[0].{d:length(deployments), r:runningCount, dc:desiredCount, rs:deployments[?status==`PRIMARY`].rolloutState | [0]}' --output json)
    D=$(echo "${STATE}" | python3 -c "import sys, json; print(json.load(sys.stdin)['d'])")
    R=$(echo "${STATE}" | python3 -c "import sys, json; print(json.load(sys.stdin)['r'])")
    DC=$(echo "${STATE}" | python3 -c "import sys, json; print(json.load(sys.stdin)['dc'])")
    RS=$(echo "${STATE}" | python3 -c "import sys, json; print(json.load(sys.stdin)['rs'])")
    if [ "${D}" = "1" ] && [ "${R}" = "${DC}" ] && [ "${RS}" = "COMPLETED" ]; then
        ELAPSED=$(($(date +%s) - START_EPOCH))
        echo "convergence_seconds=${ELAPSED}" >> "${RESULTS_DIR}/trigger.txt"
        echo "[$(date -u +%H:%M:%S)] converged in ${ELAPSED}s"
        break
    fi
    sleep 5
done

kill "${TAIL_STATE_PID}" 2>/dev/null || true

# CloudWatch has ingestion lag for ECS awslogs. Wait for the agent flush and
# a buffer so filter-log-events returns the final exit events from old tasks.
echo "[$(date -u +%H:%M:%S)] waiting 45s for CloudWatch ingestion..."
sleep 45

END_EPOCH_MS=$(($(date +%s) * 1000))
START_EPOCH_MS=$((TRIGGER_EPOCH_MS - 30000))

echo "[$(date -u +%H:%M:%S)] pulling worker logs from ${START_EPOCH_MS} to ${END_EPOCH_MS}"
aws logs filter-log-events \
    --log-group-name "${LOG_GROUP}" \
    --start-time "${START_EPOCH_MS}" \
    --end-time "${END_EPOCH_MS}" \
    --query 'events[*].message' \
    --output text | tr '\t' '\n' > "${RESULTS_DIR}/worker-logs.jsonl"

echo "[$(date -u +%H:%M:%S)] done"
echo "--- event counts ---"
grep -oE '"event": "[^"]+"' "${RESULTS_DIR}/worker-logs.jsonl" | sort | uniq -c