#!/usr/bin/env bash
# Register a new task definition with updated WORK_DURATION_MS/IDLE_MS env
# vars (plus optional jitter) and point the service at it. Used to change
# experiment profile without rebuilding the image.
#
# Usage: ./update_env.sh <work_duration_ms> <idle_ms> [expires_minutes] [startup_jitter_ms] [idle_jitter_ms]
# Example: ./update_env.sh 200 50 30                    # saturated, no jitter
# Example: ./update_env.sh 200 50 30 250 100            # saturated + phase-jitter
# Example: ./update_env.sh 20000 20000 30               # idle-ish simulation

set -euo pipefail

WORK_MS="${1:?Usage: $0 <work_duration_ms> <idle_ms> [expires_minutes] [startup_jitter_ms] [idle_jitter_ms]}"
IDLE_MS="${2:?Usage: $0 <work_duration_ms> <idle_ms> [expires_minutes] [startup_jitter_ms] [idle_jitter_ms]}"
EXPIRES_MIN="${3:-5}"
STARTUP_JITTER_MS="${4:-0}"
IDLE_JITTER_MS="${5:-0}"

CLUSTER="claude-ecs-task-protection-experiment"
SERVICE="claude-ecs-task-protection-experiment"
FAMILY="claude-ecs-task-protection-experiment"

CURRENT=$(aws ecs describe-task-definition --task-definition "${FAMILY}" --output json)

NEW_DEF=$(echo "${CURRENT}" | jq \
    --arg work "${WORK_MS}" \
    --arg idle "${IDLE_MS}" \
    --arg expires "${EXPIRES_MIN}" \
    --arg startup_jitter "${STARTUP_JITTER_MS}" \
    --arg idle_jitter "${IDLE_JITTER_MS}" \
    '
    .taskDefinition
    | .containerDefinitions[0].environment = (
        [.containerDefinitions[0].environment[]
         | select(.name | test("^(WORK_DURATION_MS|IDLE_MS|EXPIRES_MINUTES|STARTUP_JITTER_MS|IDLE_JITTER_MS)$") | not)]
        + [{name: "WORK_DURATION_MS", value: $work},
           {name: "IDLE_MS", value: $idle},
           {name: "EXPIRES_MINUTES", value: $expires},
           {name: "STARTUP_JITTER_MS", value: $startup_jitter},
           {name: "IDLE_JITTER_MS", value: $idle_jitter}]
    )
    | {family, networkMode, requiresCompatibilities, cpu, memory, executionRoleArn, taskRoleArn, containerDefinitions}
    ')

echo "${NEW_DEF}" > /tmp/experiment-task-def.json
REGISTERED=$(aws ecs register-task-definition --cli-input-json file:///tmp/experiment-task-def.json --output json)
NEW_ARN=$(echo "${REGISTERED}" | jq -r '.taskDefinition.taskDefinitionArn')

echo "Registered ${NEW_ARN}"

aws ecs update-service \
    --cluster "${CLUSTER}" \
    --service "${SERVICE}" \
    --task-definition "${NEW_ARN}" \
    --force-new-deployment \
    --output json > /dev/null

echo "Service pointed at new task def (work=${WORK_MS}ms idle=${IDLE_MS}ms expires=${EXPIRES_MIN}min)"