#!/bin/bash
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0

set -e
set -u
set -o pipefail

# Source configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "${SCRIPT_DIR}/config"
source "${SCRIPT_DIR}/lib/profile.sh"

<% if (deploymentTarget === 'realtime-inference') { %>
# ============================================================
# SageMaker Real-Time Inference Testing
# ============================================================

# Parse arguments: ./do/test [<ic-name>]
IC_ARG="${1:-}"

# Determine test mode based on ENDPOINT_NAME in config
if [ -z "${ENDPOINT_NAME:-}" ]; then
    echo "🧪 Testing local container at localhost:8080"
    echo "   Project: ${PROJECT_NAME}"
    echo "   Framework: ${FRAMEWORK}"
    echo "   Model server: ${MODEL_SERVER}"
    TARGET_URL="http://localhost:8080"
    TEST_MODE="local"
else
    echo "🧪 Testing SageMaker endpoint: ${ENDPOINT_NAME}"
    echo "   Project: ${PROJECT_NAME}"
    echo "   Framework: ${FRAMEWORK}"
    echo "   Model server: ${MODEL_SERVER}"
    echo "   Region: ${AWS_REGION}"
    TEST_MODE="sagemaker"
fi

echo ""

# Test 1: Health check (/ping)
echo "🔍 Test 1: Health check"
if [ "${TEST_MODE}" = "local" ]; then
    echo "   Sending GET request to ${TARGET_URL}/ping"
    
    if ! PING_RESPONSE=$(curl -s -w "\n%{http_code}" -X GET "${TARGET_URL}/ping" 2>&1); then
        echo "❌ Health check failed: Could not connect to local container"
        echo "   Make sure the container is running: ./do/run"
        exit 1
    fi
    
    HTTP_CODE=$(echo "${PING_RESPONSE}" | tail -n1)
    RESPONSE_BODY=$(echo "${PING_RESPONSE}" | sed '$d')
    
    if [ "${HTTP_CODE}" = "200" ]; then
        echo "✅ Health check passed (HTTP ${HTTP_CODE})"
    else
        echo "❌ Health check failed (HTTP ${HTTP_CODE})"
        echo "   Response: ${RESPONSE_BODY}"
        exit 1
    fi
else
    # For SageMaker endpoints, /ping is not directly accessible
    # We'll verify the endpoint exists and is InService
    echo "   Checking endpoint status..."
    
    if ! ENDPOINT_STATUS=$(aws sagemaker describe-endpoint \
        --endpoint-name "${ENDPOINT_NAME}" \
        --region "${AWS_REGION}" \
        --query 'EndpointStatus' \
        --output text 2>&1); then
        echo "❌ Endpoint not found or not accessible"
        echo "   Error: ${ENDPOINT_STATUS}"
        exit 1
    fi
    
    if [ "${ENDPOINT_STATUS}" = "InService" ]; then
        echo "✅ Endpoint is InService"
    else
        echo "❌ Endpoint is not InService (Status: ${ENDPOINT_STATUS})"
        exit 1
    fi
fi

echo ""

# Test 2: Inference request (/invocations)
echo "🔍 Test 2: Inference request"

# Create framework-specific test payload
case "${FRAMEWORK}" in
    sklearn|xgboost)
        # Traditional ML: JSON with instances array
        TEST_PAYLOAD='{"instances": [[1.0, 2.0, 3.0, 4.0]]}'
        echo "   Payload: Sample feature vector"
        ;;
    tensorflow)
        # TensorFlow: JSON with instances array
        TEST_PAYLOAD='{"instances": [[1.0, 2.0, 3.0, 4.0]]}'
        echo "   Payload: Sample feature vector"
        ;;
    transformers)
        # Transformers: payload format depends on model server
        case "${MODEL_SERVER}" in
            vllm|sglang)
                # OpenAI-compatible chat completions format
                # For S3/registry models, vLLM registers the model under the local path
                VLLM_MODEL_NAME="${MODEL_NAME}"
                if [[ "${MODEL_NAME}" == s3://* ]] || [[ "${MODEL_NAME}" == /opt/ml/* ]]; then
                    VLLM_MODEL_NAME="/opt/ml/model"
                fi
                TEST_PAYLOAD='{"model": "'"${VLLM_MODEL_NAME}"'", "messages": [{"role": "user", "content": "What is machine learning?"}], "max_tokens": 50, "temperature": 0.7}'
                echo "   Payload: OpenAI-compatible chat completion request"
                echo "   Model: ${VLLM_MODEL_NAME}"
                ;;
            *)
                # HuggingFace-style format for LMI, DJL, TensorRT-LLM
                TEST_PAYLOAD='{"inputs": "What is machine learning?", "parameters": {"max_new_tokens": 50, "temperature": 0.7}}'
                echo "   Payload: HuggingFace-style text generation request"
                ;;
        esac
        ;;
    diffusors)
        # Diffusors: OpenAI DALL-E compatible image generation request
        TEST_PAYLOAD='{"prompt": "A white cat", "n": 1, "size": "512x512"}'
        echo "   Payload: OpenAI DALL-E compatible image generation request"
        ;;
    *)
        echo "❌ Unknown framework: ${FRAMEWORK}"
        exit 3
        ;;
esac

if [ "${TEST_MODE}" = "local" ]; then
    if [ "${FRAMEWORK}" = "diffusors" ]; then
        echo "   Sending POST request to ${TARGET_URL}/v1/images/generations"
        
        if ! INVOKE_RESPONSE=$(curl -s -m 120 -w "\n%{http_code}" -X POST "${TARGET_URL}/v1/images/generations" \
            -H "Content-Type: application/json" \
            -d "${TEST_PAYLOAD}" 2>&1); then
            echo "❌ Inference request failed: Could not connect to local container"
            exit 1
        fi
        
        HTTP_CODE=$(echo "${INVOKE_RESPONSE}" | tail -n1)
        RESPONSE_BODY=$(echo "${INVOKE_RESPONSE}" | sed '$d')
        
        if [ "${HTTP_CODE}" = "200" ]; then
            # Validate response contains data array with b64_json
            if echo "${RESPONSE_BODY}" | python3 -c "
import sys, json
resp = json.load(sys.stdin)
assert 'data' in resp, 'Missing data array'
assert len(resp['data']) > 0, 'Empty data array'
assert 'b64_json' in resp['data'][0], 'Missing b64_json in data entry'
" 2>/dev/null; then
                echo "✅ Image generation successful (HTTP ${HTTP_CODE})"
                echo "   Response contains valid data array with b64_json image"
                
                # Save generated image to file
                OUTPUT_IMAGE="${SCRIPT_DIR}/../test_output.png"
                if echo "${RESPONSE_BODY}" | python3 -c "
import sys, json, base64
resp = json.load(sys.stdin)
img_data = base64.b64decode(resp['data'][0]['b64_json'])
with open('${OUTPUT_IMAGE}', 'wb') as f:
    f.write(img_data)
" 2>/dev/null; then
                    echo "   🖼️  Image saved to: ${OUTPUT_IMAGE}"
                fi
            else
                echo "⚠️  Image generation returned HTTP ${HTTP_CODE} but response format unexpected"
                echo "   Response preview: ${RESPONSE_BODY:0:200}"
            fi
        else
            echo "❌ Image generation failed (HTTP ${HTTP_CODE})"
            echo "   Response: ${RESPONSE_BODY}"
            exit 1
        fi
    else
        echo "   Sending POST request to ${TARGET_URL}/invocations"
        
        if ! INVOKE_RESPONSE=$(curl -s -w "\n%{http_code}" -X POST "${TARGET_URL}/invocations" \
            -H "Content-Type: application/json" \
            -d "${TEST_PAYLOAD}" 2>&1); then
            echo "❌ Inference request failed: Could not connect to local container"
            exit 1
        fi
        
        HTTP_CODE=$(echo "${INVOKE_RESPONSE}" | tail -n1)
        RESPONSE_BODY=$(echo "${INVOKE_RESPONSE}" | sed '$d')
        
        if [ "${HTTP_CODE}" = "200" ]; then
            echo "✅ Inference request successful (HTTP ${HTTP_CODE})"
            echo "   Response preview: ${RESPONSE_BODY:0:200}"
            if [ ${#RESPONSE_BODY} -gt 200 ]; then
                echo "   (truncated, full response is ${#RESPONSE_BODY} characters)"
            fi
        else
            echo "❌ Inference request failed (HTTP ${HTTP_CODE})"
            echo "   Response: ${RESPONSE_BODY}"
            exit 1
        fi
    fi
else
    echo "   Invoking SageMaker endpoint..."
    
    # Create temporary file for payload
    TEMP_PAYLOAD=$(mktemp)
    echo "${TEST_PAYLOAD}" > "${TEMP_PAYLOAD}"
    
    # Create temporary file for response
    TEMP_RESPONSE=$(mktemp)
    
    # Resolve inference component name
    # Precedence: do/adapters/ → do/ic/ → legacy config
    IC_NAME=""
    ADAPTER_MODEL_NAME=""
    if [ -n "${IC_ARG}" ] && [ -f "${SCRIPT_DIR}/adapters/${IC_ARG}.conf" ]; then
        # Argument matches an adapter name — use adapter IC
        ADAPTER_IC_NAME=""
        source "${SCRIPT_DIR}/adapters/${IC_ARG}.conf"
        if [ -z "${ADAPTER_IC_NAME}" ]; then
            echo "❌ Adapter '${IC_ARG}' conf is missing ADAPTER_IC_NAME."
            exit 1
        fi
        IC_NAME="${ADAPTER_IC_NAME}"
        ADAPTER_MODEL_NAME="${IC_ARG}"
    elif [ -n "${IC_ARG}" ]; then
        # Explicit IC name provided as argument
        IC_CONF="${SCRIPT_DIR}/ic/${IC_ARG}.conf"
        if [ ! -f "${IC_CONF}" ]; then
            echo "❌ IC config not found: do/ic/${IC_ARG}.conf"
            exit 1
        fi
        IC_DEPLOYED_NAME=""
        source "${IC_CONF}"
        if [ -z "${IC_DEPLOYED_NAME}" ]; then
            echo "❌ IC '${IC_ARG}' has not been deployed yet. Run ./do/deploy --ic ${IC_ARG} first."
            exit 1
        fi
        IC_NAME="${IC_DEPLOYED_NAME}"
    elif [ -d "${SCRIPT_DIR}/ic" ]; then
        # No argument, but do/ic/ exists — use first IC alphabetically
        IC_NAME=""
        for conf in "${SCRIPT_DIR}"/ic/*.conf; do
            [ -f "${conf}" ] || continue
            IC_DEPLOYED_NAME=""
            source "${conf}"
            if [ -n "${IC_DEPLOYED_NAME}" ]; then
                IC_NAME="${IC_DEPLOYED_NAME}"
                break
            fi
        done
        if [ -z "${IC_NAME}" ]; then
            echo "❌ No ICs deployed. Run ./do/deploy first."
            exit 1
        fi
    else
        # Legacy: no do/ic/ directory, use INFERENCE_COMPONENT_NAME from do/config
        IC_NAME="${INFERENCE_COMPONENT_NAME:-}"
    fi

    # If testing an adapter, override the model name in the payload
    if [ -n "${ADAPTER_MODEL_NAME}" ] && [ "${FRAMEWORK}" = "transformers" ]; then
        case "${MODEL_SERVER}" in
            vllm|sglang)
                TEST_PAYLOAD='{"model": "'"${ADAPTER_MODEL_NAME}"'", "messages": [{"role": "user", "content": "What is machine learning?"}], "max_tokens": 50, "temperature": 0.7}'
                echo "${TEST_PAYLOAD}" > "${TEMP_PAYLOAD}"
                echo "   (Using adapter model name: ${ADAPTER_MODEL_NAME})"
                ;;
        esac
    fi

    INVOKE_ARGS=(
        --endpoint-name "${ENDPOINT_NAME}"
        --region "${AWS_REGION}"
        --content-type "application/json"
        --body "fileb://${TEMP_PAYLOAD}"
    )
    if [ -n "${IC_NAME}" ]; then
        INVOKE_ARGS+=(--inference-component-name "${IC_NAME}")
        echo "   Inference component: ${IC_NAME}"
    fi

    INVOKE_ERROR=$(mktemp)
    if ! aws sagemaker-runtime invoke-endpoint \
        "${INVOKE_ARGS[@]}" \
        "${TEMP_RESPONSE}" 2>"${INVOKE_ERROR}"; then
        echo "❌ Inference request failed"
        echo "   Error: $(cat "${INVOKE_ERROR}")"
        rm -f "${TEMP_PAYLOAD}" "${TEMP_RESPONSE}" "${INVOKE_ERROR}"
        exit 1
    fi
    rm -f "${INVOKE_ERROR}"
    
    # Read response
    RESPONSE_BODY=$(cat "${TEMP_RESPONSE}")
    
    # Clean up temp files
    rm -f "${TEMP_PAYLOAD}" "${TEMP_RESPONSE}"
    
    if [ "${FRAMEWORK}" = "diffusors" ]; then
        # Validate response contains data array with b64_json
        if echo "${RESPONSE_BODY}" | python3 -c "
import sys, json
resp = json.load(sys.stdin)
assert 'data' in resp, 'Missing data array'
assert len(resp['data']) > 0, 'Empty data array'
assert 'b64_json' in resp['data'][0], 'Missing b64_json in data entry'
" 2>/dev/null; then
            echo "✅ Image generation successful"
            echo "   Response contains valid data array with b64_json image"
            
            # Save generated image to file
            OUTPUT_IMAGE="${SCRIPT_DIR}/../test_output.png"
            if echo "${RESPONSE_BODY}" | python3 -c "
import sys, json, base64
resp = json.load(sys.stdin)
img_data = base64.b64decode(resp['data'][0]['b64_json'])
with open('${OUTPUT_IMAGE}', 'wb') as f:
    f.write(img_data)
" 2>/dev/null; then
                echo "   🖼️  Image saved to: ${OUTPUT_IMAGE}"
            fi
        else
            echo "⚠️  Image generation returned but response format unexpected"
            echo "   Response preview: ${RESPONSE_BODY:0:200}"
        fi
    else
        echo "✅ Inference request successful"
        echo "   Response preview: ${RESPONSE_BODY:0:200}"
        if [ ${#RESPONSE_BODY} -gt 200 ]; then
            echo "   (truncated, full response is ${#RESPONSE_BODY} characters)"
        fi
    fi
fi

echo ""
echo "✅ All tests passed!"
echo ""

if [ "${TEST_MODE}" = "local" ]; then
    echo "Next steps:"
    echo "  • Push to ECR: ./do/push"
    echo "  • Deploy to SageMaker: ./do/deploy"
else
    echo "📋 What's next?"
<% if (typeof includeBenchmark !== 'undefined' && includeBenchmark) { %>
    echo "   • Benchmark performance:      ./do/benchmark"
<% } %>
<% if (typeof enableLora !== 'undefined' && enableLora) { %>
    echo "   • Add a LoRA adapter:         ./do/adapter add <name> --weights s3://..."
<% } %>
    echo "   • Register this deployment:   ./do/register"
    echo "   • View logs:                  ./do/logs"
fi

<% } else if (deploymentTarget === 'async-inference') { %>
# ============================================================
# SageMaker Async Inference Testing
# ============================================================

# Parse arguments
ENDPOINT_NAME="${1:-${ENDPOINT_NAME:-}}"

if [ -z "${ENDPOINT_NAME}" ]; then
    echo "🧪 Testing local container at localhost:8080"
    echo "   Project: ${PROJECT_NAME}"
    echo "   Framework: ${FRAMEWORK}"
    echo "   Model server: ${MODEL_SERVER}"
    TARGET_URL="http://localhost:8080"
    TEST_MODE="local"
else
    echo "🧪 Testing SageMaker async endpoint: ${ENDPOINT_NAME}"
    echo "   Project: ${PROJECT_NAME}"
    echo "   Framework: ${FRAMEWORK}"
    echo "   Model server: ${MODEL_SERVER}"
    echo "   Region: ${AWS_REGION}"
    echo "   S3 output: ${ASYNC_S3_OUTPUT_PATH}"
    TEST_MODE="sagemaker"
fi

echo ""

# Test 1: Health check (/ping)
echo "🔍 Test 1: Health check"
if [ "${TEST_MODE}" = "local" ]; then
    echo "   Sending GET request to ${TARGET_URL}/ping"
    
    if ! PING_RESPONSE=$(curl -s -w "\n%{http_code}" -X GET "${TARGET_URL}/ping" 2>&1); then
        echo "❌ Health check failed: Could not connect to local container"
        echo "   Make sure the container is running: ./do/run"
        exit 1
    fi
    
    HTTP_CODE=$(echo "${PING_RESPONSE}" | tail -n1)
    RESPONSE_BODY=$(echo "${PING_RESPONSE}" | sed '$d')
    
    if [ "${HTTP_CODE}" = "200" ]; then
        echo "✅ Health check passed (HTTP ${HTTP_CODE})"
    else
        echo "❌ Health check failed (HTTP ${HTTP_CODE})"
        echo "   Response: ${RESPONSE_BODY}"
        exit 1
    fi
else
    # For SageMaker endpoints, verify the endpoint exists and is InService
    echo "   Checking endpoint status..."
    
    if ! ENDPOINT_STATUS=$(aws sagemaker describe-endpoint \
        --endpoint-name "${ENDPOINT_NAME}" \
        --region "${AWS_REGION}" \
        --query 'EndpointStatus' \
        --output text 2>&1); then
        echo "❌ Endpoint not found or not accessible"
        echo "   Error: ${ENDPOINT_STATUS}"
        exit 1
    fi
    
    if [ "${ENDPOINT_STATUS}" = "InService" ]; then
        echo "✅ Endpoint is InService"
    else
        echo "❌ Endpoint is not InService (Status: ${ENDPOINT_STATUS})"
        exit 1
    fi
fi

echo ""

# Test 2: Inference request
echo "🔍 Test 2: Inference request"

# Create framework-specific test payload
case "${FRAMEWORK}" in
    sklearn|xgboost)
        TEST_PAYLOAD='{"instances": [[1.0, 2.0, 3.0, 4.0]]}'
        echo "   Payload: Sample feature vector"
        ;;
    tensorflow)
        TEST_PAYLOAD='{"instances": [[1.0, 2.0, 3.0, 4.0]]}'
        echo "   Payload: Sample feature vector"
        ;;
    transformers)
        case "${MODEL_SERVER}" in
            vllm|sglang)
                VLLM_MODEL_NAME="${MODEL_NAME}"
                if [[ "${MODEL_NAME}" == s3://* ]] || [[ "${MODEL_NAME}" == /opt/ml/* ]]; then
                    VLLM_MODEL_NAME="/opt/ml/model"
                fi
                TEST_PAYLOAD='{"model": "'"${VLLM_MODEL_NAME}"'", "messages": [{"role": "user", "content": "What is machine learning?"}], "max_tokens": 50, "temperature": 0.7}'
                echo "   Payload: OpenAI-compatible chat completion request"
                echo "   Model: ${VLLM_MODEL_NAME}"
                ;;
            *)
                TEST_PAYLOAD='{"inputs": "What is machine learning?", "parameters": {"max_new_tokens": 50, "temperature": 0.7}}'
                echo "   Payload: HuggingFace-style text generation request"
                ;;
        esac
        ;;
    diffusors)
        TEST_PAYLOAD='{"prompt": "A white cat", "n": 1, "size": "512x512"}'
        echo "   Payload: OpenAI DALL-E compatible image generation request"
        ;;
    *)
        echo "❌ Unknown framework: ${FRAMEWORK}"
        exit 3
        ;;
esac

if [ "${TEST_MODE}" = "local" ]; then
    if [ "${FRAMEWORK}" = "diffusors" ]; then
        echo "   Sending POST request to ${TARGET_URL}/v1/images/generations"
        
        if ! INVOKE_RESPONSE=$(curl -s -m 120 -w "\n%{http_code}" -X POST "${TARGET_URL}/v1/images/generations" \
            -H "Content-Type: application/json" \
            -d "${TEST_PAYLOAD}" 2>&1); then
            echo "❌ Inference request failed: Could not connect to local container"
            exit 1
        fi
        
        HTTP_CODE=$(echo "${INVOKE_RESPONSE}" | tail -n1)
        RESPONSE_BODY=$(echo "${INVOKE_RESPONSE}" | sed '$d')
        
        if [ "${HTTP_CODE}" = "200" ]; then
            if echo "${RESPONSE_BODY}" | python3 -c "
import sys, json
resp = json.load(sys.stdin)
assert 'data' in resp, 'Missing data array'
assert len(resp['data']) > 0, 'Empty data array'
assert 'b64_json' in resp['data'][0], 'Missing b64_json in data entry'
" 2>/dev/null; then
                echo "✅ Image generation successful (HTTP ${HTTP_CODE})"
                echo "   Response contains valid data array with b64_json image"
                
                # Save generated image to file
                OUTPUT_IMAGE="${SCRIPT_DIR}/../test_output.png"
                if echo "${RESPONSE_BODY}" | python3 -c "
import sys, json, base64
resp = json.load(sys.stdin)
img_data = base64.b64decode(resp['data'][0]['b64_json'])
with open('${OUTPUT_IMAGE}', 'wb') as f:
    f.write(img_data)
" 2>/dev/null; then
                    echo "   🖼️  Image saved to: ${OUTPUT_IMAGE}"
                fi
            else
                echo "⚠️  Image generation returned HTTP ${HTTP_CODE} but response format unexpected"
                echo "   Response preview: ${RESPONSE_BODY:0:200}"
            fi
        else
            echo "❌ Image generation failed (HTTP ${HTTP_CODE})"
            echo "   Response: ${RESPONSE_BODY}"
            exit 1
        fi
    else
        echo "   Sending POST request to ${TARGET_URL}/invocations"
        
        if ! INVOKE_RESPONSE=$(curl -s -w "\n%{http_code}" -X POST "${TARGET_URL}/invocations" \
            -H "Content-Type: application/json" \
            -d "${TEST_PAYLOAD}" 2>&1); then
            echo "❌ Inference request failed: Could not connect to local container"
            exit 1
        fi
        
        HTTP_CODE=$(echo "${INVOKE_RESPONSE}" | tail -n1)
        RESPONSE_BODY=$(echo "${INVOKE_RESPONSE}" | sed '$d')
        
        if [ "${HTTP_CODE}" = "200" ]; then
            echo "✅ Inference request successful (HTTP ${HTTP_CODE})"
            echo "   Response preview: ${RESPONSE_BODY:0:200}"
            if [ ${#RESPONSE_BODY} -gt 200 ]; then
                echo "   (truncated, full response is ${#RESPONSE_BODY} characters)"
            fi
        else
            echo "❌ Inference request failed (HTTP ${HTTP_CODE})"
            echo "   Response: ${RESPONSE_BODY}"
            exit 1
        fi
    fi
else
    # SageMaker async invocation: upload payload to S3, invoke async, poll for result
    echo "   Uploading test payload to S3..."
    
    # Create temporary file for payload
    TEMP_PAYLOAD=$(mktemp)
    echo "${TEST_PAYLOAD}" > "${TEMP_PAYLOAD}"
    
    # Upload payload to S3 input location
    ASYNC_INPUT_KEY="${PROJECT_NAME}/input/test-payload-$(date +%s).json"
    ASYNC_S3_BUCKET=$(echo "${ASYNC_S3_OUTPUT_PATH}" | sed 's|s3://||' | cut -d'/' -f1)
    S3_INPUT_LOCATION="s3://${ASYNC_S3_BUCKET}/${ASYNC_INPUT_KEY}"
    
    if ! aws s3 cp "${TEMP_PAYLOAD}" "${S3_INPUT_LOCATION}" --region "${AWS_REGION}" &> /dev/null; then
        echo "❌ Failed to upload test payload to S3"
        echo "   Location: ${S3_INPUT_LOCATION}"
        echo "   Check that your IAM credentials have s3:PutObject permission"
        rm -f "${TEMP_PAYLOAD}"
        exit 1
    fi
    rm -f "${TEMP_PAYLOAD}"
    echo "✅ Test payload uploaded to: ${S3_INPUT_LOCATION}"
    
    # Invoke endpoint asynchronously (no inference components for async)
    echo "   Invoking async endpoint..."
    
    INVOKE_ARGS=(
        --endpoint-name "${ENDPOINT_NAME}"
        --input-location "${S3_INPUT_LOCATION}"
        --region "${AWS_REGION}"
        --content-type "application/json"
    )
    
    if ! INVOKE_RESULT=$(aws sagemaker-runtime invoke-endpoint-async \
        "${INVOKE_ARGS[@]}" 2>&1); then
        echo "❌ Async invocation failed"
        echo "   Error: ${INVOKE_RESULT}"
        exit 1
    fi
    
    # Extract output location from response
    OUTPUT_LOCATION=$(echo "${INVOKE_RESULT}" | python3 -c "import sys,json; print(json.load(sys.stdin).get('OutputLocation',''))" 2>/dev/null || echo "")
    
    if [ -z "${OUTPUT_LOCATION}" ]; then
        echo "⚠️  Async invocation accepted but no output location returned"
        echo "   Check the S3 output path for results: ${ASYNC_S3_OUTPUT_PATH}"
        echo ""
        echo "✅ Async invocation submitted successfully"
    else
        echo "✅ Async invocation accepted"
        echo "   Output location: ${OUTPUT_LOCATION}"
        
        # Poll S3 output location for result
        POLL_TIMEOUT=300
        POLL_INTERVAL=10
        ELAPSED=0
        
        echo "⏳ Polling for async result (timeout: ${POLL_TIMEOUT}s)..."
        
        while [ ${ELAPSED} -lt ${POLL_TIMEOUT} ]; do
            if aws s3 ls "${OUTPUT_LOCATION}" --region "${AWS_REGION}" &> /dev/null; then
                echo "✅ Async inference result available"
                
                # Download and display result
                TEMP_RESULT=$(mktemp)
                if aws s3 cp "${OUTPUT_LOCATION}" "${TEMP_RESULT}" --region "${AWS_REGION}" &> /dev/null; then
                    RESPONSE_BODY=$(cat "${TEMP_RESULT}")
                    rm -f "${TEMP_RESULT}"
                    
                    echo "   Response preview: ${RESPONSE_BODY:0:200}"
                    if [ ${#RESPONSE_BODY} -gt 200 ]; then
                        echo "   (truncated, full response is ${#RESPONSE_BODY} characters)"
                    fi
                    
                    # For diffusors, extract and save the generated image
                    if [ "${FRAMEWORK}" = "diffusors" ]; then
                        OUTPUT_IMAGE="${SCRIPT_DIR}/../test_output.png"
                        if echo "${RESPONSE_BODY}" | python3 -c "
import sys, json, base64
resp = json.load(sys.stdin)
if 'data' in resp and len(resp['data']) > 0 and 'b64_json' in resp['data'][0]:
    img_data = base64.b64decode(resp['data'][0]['b64_json'])
    with open('${OUTPUT_IMAGE}', 'wb') as f:
        f.write(img_data)
    print('ok')
else:
    print('skip')
" 2>/dev/null | grep -q "ok"; then
                            echo "   🖼️  Image saved to: ${OUTPUT_IMAGE}"
                        fi
                    fi
                else
                    rm -f "${TEMP_RESULT}"
                    echo "⚠️  Result exists but could not be downloaded"
                fi
                break
            fi
            
            sleep ${POLL_INTERVAL}
            ELAPSED=$((ELAPSED + POLL_INTERVAL))
            echo "   ⏳ Waiting... (${ELAPSED}s / ${POLL_TIMEOUT}s)"
        done
        
        if [ ${ELAPSED} -ge ${POLL_TIMEOUT} ]; then
            echo "❌ Async inference timed out after ${POLL_TIMEOUT}s"
            echo ""
            echo "   The request may still be processing. Check:"
            echo "   • S3 output path: ${OUTPUT_LOCATION}"
            echo "   • CloudWatch Logs: /aws/sagemaker/Endpoints/${ENDPOINT_NAME}"
            echo "   • Endpoint status: aws sagemaker describe-endpoint --endpoint-name ${ENDPOINT_NAME} --region ${AWS_REGION}"
            exit 1
        fi
    fi
fi

echo ""
echo "✅ All tests passed!"
echo ""

if [ "${TEST_MODE}" = "local" ]; then
    echo "Next steps:"
    echo "  • Push to ECR: ./do/push"
    echo "  • Deploy to SageMaker: ./do/deploy"
else
    echo "📋 What's next?"
<% if (typeof includeBenchmark !== 'undefined' && includeBenchmark) { %>
    echo "   • Benchmark performance:      ./do/benchmark"
<% } %>
    echo "   • Check async output:         aws s3 ls ${ASYNC_S3_OUTPUT_PATH}"
    echo "   • Register this deployment:   ./do/register"
    echo "   • View logs:                  ./do/logs"
fi

<% } else if (deploymentTarget === 'hyperpod-eks') { %>
# ============================================================
# HyperPod EKS Testing
# ============================================================

# Parse arguments: local or hyperpod test mode
# Default to hyperpod if no argument given (deployment target is hyperpod-eks)
TEST_TARGET="${1:-hyperpod}"

case "${TEST_TARGET}" in
    local)
        echo "🧪 Testing local container at localhost:8080"
        echo "   Project: ${PROJECT_NAME}"
        echo "   Framework: ${FRAMEWORK}"
        echo "   Model server: ${MODEL_SERVER}"
        TARGET_URL="http://localhost:8080"
        ;;
    hyperpod)
        echo "🧪 Testing HyperPod EKS deployment"
        echo "   Project: ${PROJECT_NAME}"
        echo "   Framework: ${FRAMEWORK}"
        echo "   Model server: ${MODEL_SERVER}"
        echo "   Cluster: ${HYPERPOD_CLUSTER_NAME}"
        echo "   Namespace: ${HYPERPOD_NAMESPACE}"
        echo "   Region: ${AWS_REGION}"
        echo ""
        
        # Get kubeconfig for HyperPod cluster
        echo "🔑 Configuring kubectl for HyperPod cluster..."
        KUBECONFIG_PATH="${HOME}/.kube/hyperpod-${HYPERPOD_CLUSTER_NAME}"
        
        EKS_CLUSTER_ARN=$(aws sagemaker describe-cluster \
            --cluster-name "${HYPERPOD_CLUSTER_NAME}" \
            --region "${AWS_REGION}" \
            --query "Orchestrator.Eks.ClusterArn" \
            --output text 2>&1) || {
            echo "❌ Failed to describe HyperPod cluster: ${HYPERPOD_CLUSTER_NAME}"
            echo ""
            echo "   Check that:"
            echo "   • The cluster name is correct"
            echo "   • The cluster exists in region: ${AWS_REGION}"
            echo "   • Your IAM user/role has permission to access the cluster"
            exit 4
        }
        
        EKS_CLUSTER_NAME=$(echo "${EKS_CLUSTER_ARN}" | awk -F'/' '{print $NF}')
        
        if ! aws eks update-kubeconfig \
            --name "${EKS_CLUSTER_NAME}" \
            --region "${AWS_REGION}" \
            --kubeconfig "${KUBECONFIG_PATH}" 2>&1; then
            echo "❌ Failed to configure kubectl for EKS cluster: ${EKS_CLUSTER_NAME}"
            exit 4
        fi
        
        export KUBECONFIG="${KUBECONFIG_PATH}"
        
        # Verify cluster connectivity
        if ! kubectl cluster-info &> /dev/null; then
            echo "❌ Cannot connect to HyperPod cluster"
            exit 4
        fi
        echo "✅ Connected to HyperPod cluster"
        
        # Port-forward the service to a local port
        LOCAL_PORT=8080
        echo ""
        echo "🔌 Port-forwarding svc/${PROJECT_NAME} to localhost:${LOCAL_PORT}..."
        kubectl port-forward "svc/${PROJECT_NAME}" "${LOCAL_PORT}:8080" \
            -n "${HYPERPOD_NAMESPACE}" &
        PF_PID=$!
        
        # Wait for port-forward to establish
        sleep 3
        
        # Ensure cleanup on exit
        trap "kill ${PF_PID} 2>/dev/null || true" EXIT
        
        # Verify port-forward is running
        if ! kill -0 ${PF_PID} 2>/dev/null; then
            echo "❌ Port-forward failed to start"
            echo ""
            echo "   Check that:"
            echo "   • The service exists: kubectl get svc ${PROJECT_NAME} -n ${HYPERPOD_NAMESPACE}"
            echo "   • The deployment is running: kubectl get pods -n ${HYPERPOD_NAMESPACE}"
            exit 1
        fi
        echo "✅ Port-forward established"
        
        TARGET_URL="http://localhost:${LOCAL_PORT}"
        ;;
    *)
        echo "Usage: ./do/test [local|hyperpod]"
        echo ""
        echo "Test modes:"
        echo "  local    - Test local container at localhost:8080"
        echo "  hyperpod - Test HyperPod EKS deployment via port-forward"
        exit 1
        ;;
esac

echo ""

# Test 1: Health check (/ping)
echo "🔍 Test 1: Health check"
echo "   Sending GET request to ${TARGET_URL}/ping"

if ! PING_RESPONSE=$(curl -s -w "\n%{http_code}" -X GET "${TARGET_URL}/ping" 2>&1); then
    echo "❌ Health check failed: Could not connect"
    if [ "${TEST_TARGET}" = "local" ]; then
        echo "   Make sure the container is running: ./do/run"
    else
        echo "   Check that the port-forward is working and pods are running"
    fi
    exit 1
fi

HTTP_CODE=$(echo "${PING_RESPONSE}" | tail -n1)
RESPONSE_BODY=$(echo "${PING_RESPONSE}" | sed '$d')

if [ "${HTTP_CODE}" = "200" ]; then
    echo "✅ Health check passed (HTTP ${HTTP_CODE})"
else
    echo "❌ Health check failed (HTTP ${HTTP_CODE})"
    echo "   Response: ${RESPONSE_BODY}"
    exit 1
fi

echo ""

# Test 2: Inference request (/invocations)
echo "🔍 Test 2: Inference request"

# Create framework-specific test payload
case "${FRAMEWORK}" in
    sklearn|xgboost)
        # Traditional ML: JSON with instances array
        TEST_PAYLOAD='{"instances": [[1.0, 2.0, 3.0, 4.0]]}'
        echo "   Payload: Sample feature vector"
        ;;
    tensorflow)
        # TensorFlow: JSON with instances array
        TEST_PAYLOAD='{"instances": [[1.0, 2.0, 3.0, 4.0]]}'
        echo "   Payload: Sample feature vector"
        ;;
    transformers)
        # Transformers: payload format depends on model server
        case "${MODEL_SERVER}" in
            vllm|sglang)
                # OpenAI-compatible chat completions format
                VLLM_MODEL_NAME="${MODEL_NAME}"
                if [[ "${MODEL_NAME}" == s3://* ]] || [[ "${MODEL_NAME}" == /opt/ml/* ]]; then
                    VLLM_MODEL_NAME="/opt/ml/model"
                fi
                TEST_PAYLOAD='{"model": "'"${VLLM_MODEL_NAME}"'", "messages": [{"role": "user", "content": "What is machine learning?"}], "max_tokens": 50, "temperature": 0.7}'
                echo "   Payload: OpenAI-compatible chat completion request"
                echo "   Model: ${VLLM_MODEL_NAME}"
                ;;
            *)
                # HuggingFace-style format for LMI, DJL, TensorRT-LLM
                TEST_PAYLOAD='{"inputs": "What is machine learning?", "parameters": {"max_new_tokens": 50, "temperature": 0.7}}'
                echo "   Payload: HuggingFace-style text generation request"
                ;;
        esac
        ;;
    diffusors)
        # Diffusors: OpenAI DALL-E compatible image generation request
        TEST_PAYLOAD='{"prompt": "A white cat", "n": 1, "size": "512x512"}'
        echo "   Payload: OpenAI DALL-E compatible image generation request"
        ;;
    *)
        echo "❌ Unknown framework: ${FRAMEWORK}"
        exit 3
        ;;
esac

if [ "${FRAMEWORK}" = "diffusors" ]; then
    echo "   Sending POST request to ${TARGET_URL}/v1/images/generations"

    if ! INVOKE_RESPONSE=$(curl -s -m 120 -w "\n%{http_code}" -X POST "${TARGET_URL}/v1/images/generations" \
        -H "Content-Type: application/json" \
        -d "${TEST_PAYLOAD}" 2>&1); then
        echo "❌ Inference request failed: Could not connect"
        exit 1
    fi

    HTTP_CODE=$(echo "${INVOKE_RESPONSE}" | tail -n1)
    RESPONSE_BODY=$(echo "${INVOKE_RESPONSE}" | sed '$d')

    if [ "${HTTP_CODE}" = "200" ]; then
        # Validate response contains data array with b64_json
        if echo "${RESPONSE_BODY}" | python3 -c "
import sys, json
resp = json.load(sys.stdin)
assert 'data' in resp, 'Missing data array'
assert len(resp['data']) > 0, 'Empty data array'
assert 'b64_json' in resp['data'][0], 'Missing b64_json in data entry'
" 2>/dev/null; then
            echo "✅ Image generation successful (HTTP ${HTTP_CODE})"
            echo "   Response contains valid data array with b64_json image"
            
            # Save generated image to file
            OUTPUT_IMAGE="${SCRIPT_DIR}/../test_output.png"
            if echo "${RESPONSE_BODY}" | python3 -c "
import sys, json, base64
resp = json.load(sys.stdin)
img_data = base64.b64decode(resp['data'][0]['b64_json'])
with open('${OUTPUT_IMAGE}', 'wb') as f:
    f.write(img_data)
" 2>/dev/null; then
                echo "   🖼️  Image saved to: ${OUTPUT_IMAGE}"
            fi
        else
            echo "⚠️  Image generation returned HTTP ${HTTP_CODE} but response format unexpected"
            echo "   Response preview: ${RESPONSE_BODY:0:200}"
        fi
    else
        echo "❌ Image generation failed (HTTP ${HTTP_CODE})"
        echo "   Response: ${RESPONSE_BODY}"
        exit 1
    fi
else
    echo "   Sending POST request to ${TARGET_URL}/invocations"

    if ! INVOKE_RESPONSE=$(curl -s -w "\n%{http_code}" -X POST "${TARGET_URL}/invocations" \
        -H "Content-Type: application/json" \
        -d "${TEST_PAYLOAD}" 2>&1); then
        echo "❌ Inference request failed: Could not connect"
        exit 1
    fi

    HTTP_CODE=$(echo "${INVOKE_RESPONSE}" | tail -n1)
    RESPONSE_BODY=$(echo "${INVOKE_RESPONSE}" | sed '$d')

    if [ "${HTTP_CODE}" = "200" ]; then
        echo "✅ Inference request successful (HTTP ${HTTP_CODE})"
        echo "   Response preview: ${RESPONSE_BODY:0:200}"
        if [ ${#RESPONSE_BODY} -gt 200 ]; then
            echo "   (truncated, full response is ${#RESPONSE_BODY} characters)"
        fi
    else
        echo "❌ Inference request failed (HTTP ${HTTP_CODE})"
        echo "   Response: ${RESPONSE_BODY}"
        exit 1
    fi
fi

echo ""
echo "✅ All tests passed!"
echo ""

if [ "${TEST_TARGET}" = "local" ]; then
    echo "Next steps:"
    echo "  • Push to ECR: ./do/push"
    echo "  • Deploy to HyperPod: ./do/deploy"
else
    echo "📋 What's next?"
<% if (typeof includeBenchmark !== 'undefined' && includeBenchmark) { %>
    echo "   • Benchmark performance:      ./do/benchmark"
<% } %>
    echo "   • Check pod status:           kubectl get pods -n ${HYPERPOD_NAMESPACE}"
    echo "   • View pod logs:              kubectl logs -n ${HYPERPOD_NAMESPACE} -l app=${PROJECT_NAME}"
    echo "   • Register this deployment:   ./do/register"
    echo "   • View logs:                  ./do/logs"
fi

<% } else if (deploymentTarget === 'batch-transform') { %>
# ============================================================
# SageMaker Batch Transform Testing
# ============================================================

# Parse arguments: local or batch test mode
# Default to batch if no argument given (deployment target is batch-transform)
TEST_TARGET="${1:-batch}"

case "${TEST_TARGET}" in
    local)
        echo "🧪 Testing local container at localhost:8080"
        echo "   Project: ${PROJECT_NAME}"
        echo "   Framework: ${FRAMEWORK}"
        echo "   Model server: ${MODEL_SERVER}"
        TARGET_URL="http://localhost:8080"
        ;;
    batch)
        echo "🧪 Checking batch transform job status"
        echo "   Project: ${PROJECT_NAME}"
        echo "   Framework: ${FRAMEWORK}"
        echo "   Region: ${AWS_REGION}"
        echo "   S3 input: ${BATCH_INPUT_PATH}"
        echo "   S3 output: ${BATCH_OUTPUT_PATH}"
        echo ""

        # Get transform job name from config
        TRANSFORM_JOB_NAME="${TRANSFORM_JOB_NAME:-}"
        if [ -z "${TRANSFORM_JOB_NAME}" ]; then
            echo "❌ No transform job name found"
            echo "   Run ./do/deploy first to create a transform job"
            exit 1
        fi

        echo "🔍 Checking transform job: ${TRANSFORM_JOB_NAME}"

        if ! JOB_STATUS_JSON=$(aws sagemaker describe-transform-job \
            --transform-job-name "${TRANSFORM_JOB_NAME}" \
            --region "${AWS_REGION}" 2>&1); then
            echo "❌ Failed to describe transform job"
            echo "   Error: ${JOB_STATUS_JSON}"
            exit 1
        fi

        JOB_STATUS=$(echo "${JOB_STATUS_JSON}" | python3 -c "import sys,json; print(json.load(sys.stdin).get('TransformJobStatus','Unknown'))" 2>/dev/null || echo "Unknown")

        case "${JOB_STATUS}" in
            Completed)
                echo "✅ Transform job completed successfully"
                echo ""

                # Download results locally
                LOCAL_OUTPUT_DIR="${SCRIPT_DIR}/../batch-output"
                mkdir -p "${LOCAL_OUTPUT_DIR}"
                echo "📥 Downloading results to ${LOCAL_OUTPUT_DIR}/"
                if aws s3 sync "${BATCH_OUTPUT_PATH}" "${LOCAL_OUTPUT_DIR}/" --region "${AWS_REGION}"; then
                    DOWNLOADED=$(ls -1 "${LOCAL_OUTPUT_DIR}" 2>/dev/null | wc -l | tr -d ' ')
                    echo "✅ Downloaded ${DOWNLOADED} file(s) to ${LOCAL_OUTPUT_DIR}/"
                    echo ""

                    # Display first output file preview
                    FIRST_FILE=$(ls -1 "${LOCAL_OUTPUT_DIR}" 2>/dev/null | head -1)
                    if [ -n "${FIRST_FILE}" ]; then
                        echo "📄 Sample output (${FIRST_FILE}):"
                        head -5 "${LOCAL_OUTPUT_DIR}/${FIRST_FILE}"
                        LINES=$(wc -l < "${LOCAL_OUTPUT_DIR}/${FIRST_FILE}" | tr -d ' ')
                        if [ "${LINES}" -gt 5 ]; then
                            echo "   ... (${LINES} total lines)"
                        fi
                    fi
                else
                    echo "⚠️  Could not download output files"
                fi

                echo ""
                echo "✅ All tests passed!"
                echo ""
                echo "📋 What's next?"
                echo "   • View results:               cat batch-output/"
                echo "   • Register this deployment:   ./do/register"
                echo "   • View logs:                  ./do/logs"
                ;;
            InProgress)
                echo "⏳ Transform job is still in progress"

                # Extract progress details
                CREATION_TIME=$(echo "${JOB_STATUS_JSON}" | python3 -c "import sys,json; print(json.load(sys.stdin).get('CreationTime','Unknown'))" 2>/dev/null || echo "Unknown")
                echo "   Started: ${CREATION_TIME}"
                echo "   Status: InProgress"
                echo ""
                echo "   The job is still running. Check again later:"
                echo "   ./do/test"
                echo ""
                echo "   View logs:"
                echo "   ./do/logs"
                ;;
            Failed)
                echo "❌ Transform job failed"

                FAILURE_REASON=$(echo "${JOB_STATUS_JSON}" | python3 -c "import sys,json; print(json.load(sys.stdin).get('FailureReason','Unknown'))" 2>/dev/null || echo "Unknown")
                echo "   Reason: ${FAILURE_REASON}"
                echo ""
                echo "   View logs for more details:"
                echo "   ./do/logs"
                exit 1
                ;;
            Stopped)
                echo "⚠️  Transform job was stopped"
                echo "   The job was manually stopped before completion"
                echo ""
                echo "   To start a new job, run:"
                echo "   ./do/deploy"
                ;;
            *)
                echo "⚠️  Transform job status: ${JOB_STATUS}"
                echo "   Check again later: ./do/test"
                ;;
        esac
        exit 0
        ;;
    *)
        echo "Usage: ./do/test [local|batch]"
        echo ""
        echo "Test modes:"
        echo "  local - Test local container at localhost:8080"
        echo "  batch - Check transform job status and view results"
        exit 1
        ;;
esac

echo ""

# Test 1: Health check (/ping)
echo "🔍 Test 1: Health check"
echo "   Sending GET request to ${TARGET_URL}/ping"

if ! PING_RESPONSE=$(curl -s -w "\n%{http_code}" -X GET "${TARGET_URL}/ping" 2>&1); then
    echo "❌ Health check failed: Could not connect to local container"
    echo "   Make sure the container is running: ./do/run"
    exit 1
fi

HTTP_CODE=$(echo "${PING_RESPONSE}" | tail -n1)
RESPONSE_BODY=$(echo "${PING_RESPONSE}" | sed '$d')

if [ "${HTTP_CODE}" = "200" ]; then
    echo "✅ Health check passed (HTTP ${HTTP_CODE})"
else
    echo "❌ Health check failed (HTTP ${HTTP_CODE})"
    echo "   Response: ${RESPONSE_BODY}"
    exit 1
fi

echo ""

# Test 2: Inference request (/invocations)
echo "🔍 Test 2: Inference request"

# Create framework-specific test payload
case "${FRAMEWORK}" in
    sklearn|xgboost)
        TEST_PAYLOAD='{"instances": [[1.0, 2.0, 3.0, 4.0]]}'
        echo "   Payload: Sample feature vector"
        ;;
    tensorflow)
        TEST_PAYLOAD='{"instances": [[1.0, 2.0, 3.0, 4.0]]}'
        echo "   Payload: Sample feature vector"
        ;;
    transformers)
        case "${MODEL_SERVER}" in
            vllm|sglang)
                VLLM_MODEL_NAME="${MODEL_NAME}"
                if [[ "${MODEL_NAME}" == s3://* ]] || [[ "${MODEL_NAME}" == /opt/ml/* ]]; then
                    VLLM_MODEL_NAME="/opt/ml/model"
                fi
                TEST_PAYLOAD='{"model": "'"${VLLM_MODEL_NAME}"'", "messages": [{"role": "user", "content": "What is machine learning?"}], "max_tokens": 50, "temperature": 0.7}'
                echo "   Payload: OpenAI-compatible chat completion request"
                echo "   Model: ${VLLM_MODEL_NAME}"
                ;;
            *)
                TEST_PAYLOAD='{"inputs": "What is machine learning?", "parameters": {"max_new_tokens": 50, "temperature": 0.7}}'
                echo "   Payload: HuggingFace-style text generation request"
                ;;
        esac
        ;;
    diffusors)
        TEST_PAYLOAD='{"prompt": "A white cat", "n": 1, "size": "512x512"}'
        echo "   Payload: OpenAI DALL-E compatible image generation request"
        ;;
    *)
        echo "❌ Unknown framework: ${FRAMEWORK}"
        exit 3
        ;;
esac

if [ "${FRAMEWORK}" = "diffusors" ]; then
    echo "   Sending POST request to ${TARGET_URL}/v1/images/generations"

    if ! INVOKE_RESPONSE=$(curl -s -m 120 -w "\n%{http_code}" -X POST "${TARGET_URL}/v1/images/generations" \
        -H "Content-Type: application/json" \
        -d "${TEST_PAYLOAD}" 2>&1); then
        echo "❌ Inference request failed: Could not connect to local container"
        exit 1
    fi

    HTTP_CODE=$(echo "${INVOKE_RESPONSE}" | tail -n1)
    RESPONSE_BODY=$(echo "${INVOKE_RESPONSE}" | sed '$d')

    if [ "${HTTP_CODE}" = "200" ]; then
        if echo "${RESPONSE_BODY}" | python3 -c "
import sys, json
resp = json.load(sys.stdin)
assert 'data' in resp, 'Missing data array'
assert len(resp['data']) > 0, 'Empty data array'
assert 'b64_json' in resp['data'][0], 'Missing b64_json in data entry'
" 2>/dev/null; then
            echo "✅ Image generation successful (HTTP ${HTTP_CODE})"
            echo "   Response contains valid data array with b64_json image"

            # Save generated image to file
            OUTPUT_IMAGE="${SCRIPT_DIR}/../test_output.png"
            if echo "${RESPONSE_BODY}" | python3 -c "
import sys, json, base64
resp = json.load(sys.stdin)
img_data = base64.b64decode(resp['data'][0]['b64_json'])
with open('${OUTPUT_IMAGE}', 'wb') as f:
    f.write(img_data)
" 2>/dev/null; then
                echo "   🖼️  Image saved to: ${OUTPUT_IMAGE}"
            fi
        else
            echo "⚠️  Image generation returned HTTP ${HTTP_CODE} but response format unexpected"
            echo "   Response preview: ${RESPONSE_BODY:0:200}"
        fi
    else
        echo "❌ Image generation failed (HTTP ${HTTP_CODE})"
        echo "   Response: ${RESPONSE_BODY}"
        exit 1
    fi
else
    echo "   Sending POST request to ${TARGET_URL}/invocations"

    if ! INVOKE_RESPONSE=$(curl -s -w "\n%{http_code}" -X POST "${TARGET_URL}/invocations" \
        -H "Content-Type: application/json" \
        -d "${TEST_PAYLOAD}" 2>&1); then
        echo "❌ Inference request failed: Could not connect to local container"
        exit 1
    fi

    HTTP_CODE=$(echo "${INVOKE_RESPONSE}" | tail -n1)
    RESPONSE_BODY=$(echo "${INVOKE_RESPONSE}" | sed '$d')

    if [ "${HTTP_CODE}" = "200" ]; then
        echo "✅ Inference request successful (HTTP ${HTTP_CODE})"
        echo "   Response preview: ${RESPONSE_BODY:0:200}"
        if [ ${#RESPONSE_BODY} -gt 200 ]; then
            echo "   (truncated, full response is ${#RESPONSE_BODY} characters)"
        fi
    else
        echo "❌ Inference request failed (HTTP ${HTTP_CODE})"
        echo "   Response: ${RESPONSE_BODY}"
        exit 1
    fi
fi

echo ""
echo "✅ All tests passed!"
echo ""

echo "Next steps:"
echo "  • Push to ECR: ./do/push"
echo "  • Deploy batch transform: ./do/deploy"

<% } %>