#!/usr/bin/bash

# Default values
OUTPUT_DIR="."
TIMEOUT=""

# Function to show usage
usage() {
    echo "Usage: $0 [-o|--output <dir>] [-t|--timeout <seconds>] -- <cmd2optimize args>"
    exit 1
}

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        -o|--output)
            OUTPUT_DIR="$2"
            shift 2
            ;;
        -t|--timeout)
            TIMEOUT="$2"
            shift 2
            ;;
        --)
            shift
            CMD_ARGS=("$@")
            break
            ;;
        *)
            echo "Unknown option: $1"
            usage
            ;;
    esac
done

# Check dependencies
if ! command -v cmd2optimize &> /dev/null; then
    echo "Error: cmd2optimize command not found."
    exit 1
fi

if ! command -v jq &> /dev/null; then
    echo "Error: jq command not found."
    exit 1
fi

if [ ${#CMD_ARGS[@]} -eq 0 ]; then
    echo "Error: No arguments provided for cmd2optimize."
    usage
fi

CMD_PREFIX="cmd2optimize"
if [ -n "$TIMEOUT" ]; then
    CMD_PREFIX="$CMD_PREFIX -t $TIMEOUT"
fi

# Execute cmd2optimize to get the optimization command
# We pass the arguments captured after --
# Note: We need to handle the arguments carefully to preserve quoting if possible, 
# but "$@" does that for the array.
OPTIMIZE_CMD=$($CMD_PREFIX -- "${CMD_ARGS[@]}")

if [ $? -ne 0 ]; then
    echo "Error: cmd2optimize failed."
    exit 1
fi

# Create temporary file for record-out
TMP_JSON=$(mktemp)
trap 'rm -f "$TMP_JSON"' EXIT

# Modify the command:
# 1. Replace 'ctr-remote' with 'ctr-analyze'
# 2. Inject '--record-out <tmp_file>'
# The cmd2optimize output starts with "ctr-remote image optimize"
# We replace it with "ctr-analyze image optimize --record-out <tmp_file>"
MODIFIED_CMD=$(echo "$OPTIMIZE_CMD" | sed "s|ctr-remote image optimize|ctr-analyze image optimize --record-out $TMP_JSON|")

echo "Running: "
echo "$MODIFIED_CMD"

eval "$MODIFIED_CMD"

if [ $? -ne 0 ]; then
    echo "Error: ctr-analyze failed."
    exit 1
fi

# Process the output with jq
if [ ! -d "$OUTPUT_DIR" ]; then
    mkdir -p "$OUTPUT_DIR"
fi
OUTPUT_FILE="${OUTPUT_DIR}/prefetch.json"

if [ -s "$TMP_JSON" ]; then
    # Convert stream of JSON objects to array of paths
    # Input: {"path":"...", ...}\n{"path":"...", ...}
    # Output: ["path1", "path2", ...]
    jq -s '[.[] | .path]' "$TMP_JSON" > "$OUTPUT_FILE"
    echo "Prefetch list saved to $OUTPUT_FILE"
else
    echo "Warning: No prefetch data recorded."
fi

