#!/usr/bin/env bash
# Continuously snapshot ECS service + per-task state to JSONL files.
# Run in one terminal while an experiment executes in another.
#
# Usage: ./tail_state.sh [output_dir]
# Example: ./tail_state.sh ./results/2026-04-17-saturated-deploy

set -euo pipefail

OUTPUT_DIR="${1:-./state-$(date +%Y%m%d-%H%M%S)}"
CLUSTER="claude-ecs-task-protection-experiment"
SERVICE="claude-ecs-task-protection-experiment"

mkdir -p "${OUTPUT_DIR}"
STATE_FILE="${OUTPUT_DIR}/service-state.jsonl"
EVENTS_FILE="${OUTPUT_DIR}/service-events.jsonl"
TASKS_FILE="${OUTPUT_DIR}/task-state.jsonl"
DEPLOYMENTS_FILE="${OUTPUT_DIR}/service-deployments.jsonl"
TASK_DEFS_FILE="${OUTPUT_DIR}/task-definitions.jsonl"

# Ensure the task def snapshot file exists so set+jq works even on first tick.
touch "${TASK_DEFS_FILE}"

echo "Snapshotting every 2s to ${OUTPUT_DIR}/"
echo "  service-state.jsonl — desired/running/pending + deployments[] summary"
echo "  service-events.jsonl — ECS scheduler events"
echo "  task-state.jsonl — per-task lifecycle incl. stopCode/stoppedReason"
echo "  service-deployments.jsonl — detailed rollout state via DescribeServiceDeployments (ECS 2024+ API)"
echo "  task-definitions.jsonl — content of each task def revision we encounter (captured once per revision)"
echo "Ctrl-C to stop."

# Memoize task definition ARNs we've already recorded content for — each
# revision's content is constant, so recording it once per experiment is enough.
# File-based set so this works on bash 3 (macOS default).
SEEN_TASK_DEFS_FILE="${OUTPUT_DIR}/.seen-task-defs"
touch "${SEEN_TASK_DEFS_FILE}"

while true; do
    TS=$(date -u +%Y-%m-%dT%H:%M:%S.%3NZ)
    SERVICE_SNAPSHOT=$(aws ecs describe-services \
        --cluster "${CLUSTER}" \
        --services "${SERVICE}" \
        --output json 2>/dev/null)

    echo "${SERVICE_SNAPSHOT}" | jq -c \
        --arg ts "${TS}" \
        '{ts: $ts, runningCount: .services[0].runningCount, desiredCount: .services[0].desiredCount, pendingCount: .services[0].pendingCount, deployments: [.services[0].deployments[] | {id, status, taskDefinition, desiredCount, runningCount, pendingCount, failedTasks, createdAt, updatedAt, rolloutState, rolloutStateReason}]}' \
        >> "${STATE_FILE}"

    # Dump the most recent 3 service events each tick. Dedup later.
    echo "${SERVICE_SNAPSHOT}" | jq -c \
        --arg ts "${TS}" \
        '.services[0].events[0:3][] | {snapshot_ts: $ts, id, createdAt, message}' \
        >> "${EVENTS_FILE}"

    # Per-task state. Includes RUNNING/STOPPED tasks so we can see the full
    # lifecycle: creation, task def revision, stopCode + stoppedReason when
    # ECS terminates. JSON output ensures we get one ARN per line, which
    # is needed because describe-tasks caps at 100 ARNs per call.
    RUNNING_ARNS=$(aws ecs list-tasks \
        --cluster "${CLUSTER}" --service-name "${SERVICE}" \
        --desired-status RUNNING --output json 2>/dev/null \
        | jq -r '.taskArns[]')
    # Cap stopped tasks at 20 most-recent — old revisions accumulate from
    # repeated update_env registrations and would otherwise exceed the API cap.
    STOPPED_ARNS=$(aws ecs list-tasks \
        --cluster "${CLUSTER}" --service-name "${SERVICE}" \
        --desired-status STOPPED --output json 2>/dev/null \
        | jq -r '.taskArns[]' | tail -20)

    ALL_ARNS=$(printf '%s\n%s\n' "${RUNNING_ARNS}" "${STOPPED_ARNS}" | tr '\n' ' ' | xargs)
    if [[ -n "${ALL_ARNS}" ]]; then
        # shellcheck disable=SC2086
        TASK_SNAPSHOT=$(aws ecs describe-tasks \
            --cluster "${CLUSTER}" --tasks ${ALL_ARNS} \
            --output json 2>/dev/null)

        echo "${TASK_SNAPSHOT}" | jq -c --arg ts "${TS}" \
            '.tasks[] | {snapshot_ts: $ts, taskArn, taskDefinitionArn, lastStatus, desiredStatus, createdAt, startedAt, stoppingAt, stoppedAt, stopCode, stoppedReason, healthStatus, imageDigest: (.containers[0].imageDigest // null), image: (.containers[0].image // null)}' \
            >> "${TASKS_FILE}"

        # Capture task definition CONTENT the first time we see each revision.
        # Content doesn't change within a revision, so we avoid re-reading every
        # tick. This tells us what env vars etc. differed between revisions.
        NEW_DEF_ARNS=$(echo "${TASK_SNAPSHOT}" | jq -r '[.tasks[].taskDefinitionArn] | unique | .[]')
        for DEF_ARN in ${NEW_DEF_ARNS}; do
            if ! grep -qxF "${DEF_ARN}" "${SEEN_TASK_DEFS_FILE}"; then
                echo "${DEF_ARN}" >> "${SEEN_TASK_DEFS_FILE}"
                aws ecs describe-task-definition \
                    --task-definition "${DEF_ARN}" \
                    --output json 2>/dev/null \
                    | jq -c --arg ts "${TS}" \
                        '{first_seen_ts: $ts, taskDefinitionArn: .taskDefinition.taskDefinitionArn, family: .taskDefinition.family, revision: .taskDefinition.revision, registeredAt: .taskDefinition.registeredAt, containerDefinitions: .taskDefinition.containerDefinitions}' \
                    >> "${TASK_DEFS_FILE}"
            fi
        done
    fi

    # Newer ECS API — detailed service-deployment records (rollout state,
    # alarm breaches, circuit breaker stages). Safe to call even if the
    # account/region doesn't support it yet — we swallow errors.
    DEPLOYMENT_ARNS=$(echo "${SERVICE_SNAPSHOT}" | jq -r '.services[0].deployments[].id // empty' | head -10)
    if [[ -n "${DEPLOYMENT_ARNS}" ]]; then
        SERVICE_ARN=$(echo "${SERVICE_SNAPSHOT}" | jq -r '.services[0].serviceArn')
        # Prefix deployment IDs with the service ARN to form full deployment ARNs.
        DEPLOYMENT_FULL_ARNS=""
        for ID in ${DEPLOYMENT_ARNS}; do
            # ID format: "ecs-svc/1234567890" — deployment ARN replaces service/ with service-deployment/
            DEPLOY_ARN=$(echo "${SERVICE_ARN}" | sed 's|:service/|:service-deployment/|')/${ID}
            DEPLOYMENT_FULL_ARNS="${DEPLOYMENT_FULL_ARNS} ${DEPLOY_ARN}"
        done
        # shellcheck disable=SC2086
        aws ecs describe-service-deployments \
            --service-deployments ${DEPLOYMENT_FULL_ARNS} \
            --output json 2>/dev/null \
            | jq --arg ts "${TS}" \
                '.serviceDeployments[]? | {snapshot_ts: $ts, serviceDeploymentArn, status, statusReason, createdAt, finishedAt, stoppedAt, targetServiceRevisionArn, sourceServiceRevisions, rollout, alarms: (.alarms // null), deploymentCircuitBreaker: (.deploymentCircuitBreaker // null), deploymentConfiguration: (.deploymentConfiguration // null)}' \
            >> "${DEPLOYMENTS_FILE}" 2>/dev/null || true
    fi

    sleep 2
done