#!/bin/bash

# find-ownership-ingestion-arn.sh
#
# A script to find the ARN of an AWS Step Function execution based on a value
# within its input JSON, optionally filtered by a specific day.
#
# Usage:
#   ./find-ownership-ingestion-arn.sh <state_machine_arn> <substring_to_match> [date YYYY-MM-DD]
#
# Example (without date):
#   ./find-ownership-ingestion-arn.sh "arn:..." "unique-file"
#
# Example (with date filter):
#   ./find-ownership-ingestion-arn.sh "arn:..." "unique-file" "2025-07-04"
#
# Dependencies:
#   - AWS CLI (v2 recommended): https://aws.amazon.com/cli/
#   - jq: https://stedolan.github.io/jq/

set -e # Exit immediately if a command exits with a non-zero status.
set -o pipefail # The return value of a pipeline is the status of the last command to exit with a non-zero status.

# --- Argument and Dependency Validation ---

# Check for required dependencies
if ! command -v aws &> /dev/null; then
    echo "Error: The 'aws' CLI is required but could not be found." >&2
    echo "Please install it and configure it before running this script." >&2
    exit 1
fi

if ! command -v jq &> /dev/null; then
    echo "Error: 'jq' is required but could not be found." >&2
    echo "Please install it (e.g., 'brew install jq' or 'sudo apt-get install jq')." >&2
    exit 1
fi

# Check for the correct number of arguments
if [ "$#" -lt 2 ] || [ "$#" -gt 3 ]; then
    echo "Usage: $0 <state_machine_arn> <string_to_match> [date YYYY-MM-DD]"
    echo "Example: $0 \"arn:aws:states:us-east-1:123456789012:stateMachine:my-machine\" \"my-file.zip\" \"2025-07-04\""
    exit 1
fi

STATE_MACHINE_ARN="$1"
SEARCH_STRING="$2"
SEARCH_DATE="$3" # This will be empty if not provided

# --- Main Logic ---

echo "🔍 Searching for an execution of '${STATE_MACHINE_ARN}'"
echo "   where input's 'detail.object.key' contains '${SEARCH_STRING}'..."
if [ -n "$SEARCH_DATE" ]; then
    echo "   ...and started on date (UTC): ${SEARCH_DATE}"
fi
echo "---"

# Get a list of all executions. We will filter this list using jq.
# The AWS CLI v2 handles pagination automatically.
all_executions_json=$(aws stepfunctions list-executions \
  --state-machine-arn "${STATE_MACHINE_ARN}" \
  --output json)

# Prepare the jq query.
# Start with the base array.
jq_query=".executions[]"

# If a date is provided, add a filter.
# This checks if the 'startDate' (which is a full timestamp in UTC) starts with the given date string.
if [ -n "$SEARCH_DATE" ]; then
    jq_query+=" | select(.startDate | startswith(\"${SEARCH_DATE}\"))"
fi

# Finally, extract the executionArn from the filtered results.
jq_query+=" | .executionArn"

# Execute the jq query and get the ARNs to check.
execution_arns_to_check=$(echo "$all_executions_json" | jq -r "${jq_query}")


if [ -z "$execution_arns_to_check" ]; then
    echo "No executions found for the specified state machine and date."
    exit 0
fi

# This variable will hold the ARN once we find it.
found_arn=""

# Use process substitution (< <(...)) to feed the while loop.
# This avoids creating a subshell for the loop, allowing `break` to work as expected.
while read -r execution_arn; do
    echo "   Checking execution: ${execution_arn##*:}" # Print just the UUID part for brevity

    # For the current execution, get its full description, which includes the input.
    # The '.input' field is a JSON string, so we must parse it twice.
    key_from_input=$(aws stepfunctions describe-execution \
      --execution-arn "${execution_arn}" \
      --query "input" \
      --output json | jq -r 'fromjson | .detail.object.key // "null"')

    # Check if the extracted key contains the search string (partial match).
    if [[ "${key_from_input}" == *"${SEARCH_STRING}"* ]]; then
        found_arn="${execution_arn}"
        # We found our match, so we can exit the loop.
        break
    fi
done < <(echo "$execution_arns_to_check")

# After the loop, check if we found a match.
if [ -n "$found_arn" ]; then
    echo ""
    echo "✅ Match Found!"
    echo "Execution ARN: ${found_arn}"
    exit 0
else
    # If the loop completes without finding a match.
    echo "---"
    echo "❌ No execution found with a matching input."
    exit 1
fi