#!/bin/bash
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Marketplace model package deployment script.
# Deploys a pre-built AWS Marketplace model package using CreateModel with ModelPackageName.
# No build, push, or submit steps — the vendor provides the container and weights.

set -e
set -u
set -o pipefail

# Parse flags
FORCE_NEW=false
FORCE_IC=false
while [ $# -gt 0 ]; do
    case "$1" in
        --force) FORCE_NEW=true; shift ;;
        --force-ic) FORCE_IC=true; shift ;;
        --help|-h)
            echo "Usage: ./do/deploy [--force] [--force-ic]"
            echo ""
            echo "Options:"
            echo "  --force      Create a new deployment, even if one already exists."
            echo "  --force-ic   Recreate the endpoint configuration on the existing endpoint."
            echo ""
            echo "Without flags, deploy resumes from the last run."
            exit 0
            ;;
        *)
            echo "❌ Unknown option: $1"
            echo "   Run ./do/deploy --help for usage."
            exit 1
            ;;
    esac
done

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

echo "🚀 Deploying Marketplace Model Package"
echo "   Project: ${PROJECT_NAME}"
echo "   Deployment config: marketplace"
echo "   Region: ${AWS_REGION}"
echo "   Model package: ${MODEL_PACKAGE_ARN}"
echo "   Deployment target: ${DEPLOYMENT_TARGET}"
echo "   Instance type: ${INSTANCE_TYPE}"
<% if (deploymentTarget === 'async-inference') { %>
echo "   S3 output: ${ASYNC_S3_OUTPUT_PATH}"
echo "   SNS success: ${ASYNC_SNS_SUCCESS_TOPIC}"
echo "   SNS error: ${ASYNC_SNS_ERROR_TOPIC}"
<% if (asyncMaxConcurrentInvocations) { %>
echo "   Max concurrent: ${ASYNC_MAX_CONCURRENT_INVOCATIONS}"
<% } %>
<% } else if (deploymentTarget === 'batch-transform') { %>
echo "   Instance count: ${BATCH_INSTANCE_COUNT}"
echo "   S3 input: ${BATCH_INPUT_PATH}"
echo "   S3 output: ${BATCH_OUTPUT_PATH}"
echo "   Split type: ${BATCH_SPLIT_TYPE}"
echo "   Strategy: ${BATCH_STRATEGY}"
<% } %>

# Check AWS credentials
echo "🔍 Validating AWS credentials..."
if ! aws sts get-caller-identity &> /dev/null; then
    echo "❌ AWS credentials not configured"
    echo "   Run: aws configure"
    echo "   Or set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables"
    exit 4
fi

AWS_ACCOUNT_ID=$(aws sts get-caller-identity --query Account --output text)
echo "✅ AWS credentials validated (Account: ${AWS_ACCOUNT_ID})"

# Source shared helpers
source "${SCRIPT_DIR}/lib/wait.sh"
source "${SCRIPT_DIR}/lib/endpoint-config.sh"

# Validate execution role ARN
if [ -z "${ROLE_ARN:-}" ]; then
    echo "❌ Execution role ARN not provided"
    echo ""
    echo "Usage:"
    echo "  export ROLE_ARN=arn:aws:iam::ACCOUNT_ID:role/YOUR_ROLE"
    echo "  ./do/deploy"
    echo ""
    echo "Or set ROLE_ARN in do/config"
    echo ""
    echo "The execution role must have permissions for:"
    echo "  • SageMaker model and endpoint management"
    echo "  • Access to the Marketplace model package"
    echo "  • CloudWatch Logs"
    exit 3
fi

echo "   Using execution role: ${ROLE_ARN}"

<% if (deploymentTarget === 'realtime-inference') { %>
# ============================================================
# SageMaker Real-Time Inference Deployment (Model-Based)
# Marketplace packages use the classic model-based flow:
# CreateModel(ModelPackageName) → CreateEndpointConfig → CreateEndpoint
# ============================================================

# ============================================================
# Idempotency: check for existing deployment from a previous run
# ============================================================
SKIP_TO=""

if [ "${FORCE_NEW}" = true ]; then
    echo "🔄 --force: ignoring previous deployment, creating new resources."
elif [ -n "${ENDPOINT_NAME:-}" ]; then
    echo "🔍 Checking for existing deployment: ${ENDPOINT_NAME}"

    EP_STATUS=$(_get_endpoint_status "${ENDPOINT_NAME}")

    case "${EP_STATUS}" in
        InService)
            echo "✅ Endpoint already InService: ${ENDPOINT_NAME}"
            echo ""
            echo "📋 Deployment is already live. Nothing to do."
            echo "   Endpoint: ${ENDPOINT_NAME}"
            echo ""
            echo "🧪 Test your endpoint:"
            echo "   ./do/test"
            echo ""
            echo "🧹 Clean up when done:"
            echo "   ./do/clean endpoint"
            exit 0
            ;;
        Creating|Updating)
            echo "⏳ Endpoint still ${EP_STATUS}: ${ENDPOINT_NAME}"
            SKIP_TO="wait_endpoint"
            ;;
        Failed)
            echo "⚠️  Previous endpoint failed: ${ENDPOINT_NAME}"
            echo "   Creating a new deployment. Clean up the failed endpoint with:"
            echo "   ./do/clean endpoint"
            echo ""
            ;;
        "")
            echo "   Previous endpoint not found (may have been cleaned up). Creating new deployment."
            ;;
        *)
            echo "   Endpoint in unexpected state: ${EP_STATUS}. Creating new deployment."
            ;;
    esac
fi

# ============================================================
# Create resources (skip if resuming from wait)
# ============================================================
if [ -z "${SKIP_TO}" ]; then
    TIMESTAMP=$(date +%s)
    MODEL_NAME_SM="${PROJECT_NAME}-mkt-model-${TIMESTAMP}"
    ENDPOINT_CONFIG_NAME="${PROJECT_NAME}-mkt-epc-${TIMESTAMP}"
    ENDPOINT_NAME="${PROJECT_NAME}-mkt-ep-${TIMESTAMP}"

    _update_config_var "ENDPOINT_NAME" "${ENDPOINT_NAME}"
    _update_config_var "ENDPOINT_CONFIG_NAME" "${ENDPOINT_CONFIG_NAME}"
    _update_config_var "SAGEMAKER_MODEL_NAME" "${MODEL_NAME_SM}"

    # Step 1: Create SageMaker model from Marketplace model package
    echo "📦 Creating SageMaker model from Marketplace package: ${MODEL_NAME_SM}"
    if ! aws sagemaker create-model \
        --model-name "${MODEL_NAME_SM}" \
        --primary-container "{\"ModelPackageName\":\"${MODEL_PACKAGE_ARN}\"}" \
        --execution-role-arn "${ROLE_ARN}" \
        --region "${AWS_REGION}"; then

        echo "❌ Failed to create model from package ARN. Check IAM permissions and subscription status."
        echo "   Check that:"
        echo "   • The model package ARN is correct: ${MODEL_PACKAGE_ARN}"
        echo "   • Your Marketplace subscription is active"
        echo "   • The execution role has permission to access the model package"
        exit 4
    fi

    echo "✅ SageMaker model created: ${MODEL_NAME_SM}"

    # Record model in manifest (non-blocking)
    MODEL_ARN="arn:aws:sagemaker:${AWS_REGION}:${AWS_ACCOUNT_ID}:model/${MODEL_NAME_SM}"
    ./do/manifest add \
        --type sagemaker-model \
        --id "${MODEL_ARN}" \
        --project "${PROJECT_NAME}" \
        --meta "{\"modelName\":\"${MODEL_NAME_SM}\",\"modelPackageArn\":\"${MODEL_PACKAGE_ARN}\",\"region\":\"${AWS_REGION}\"}" \
        2>/dev/null || true

    # Step 2: Create endpoint configuration
    # Set MODEL_NAME_SM so endpoint-config.sh uses model-based flow (no --execution-role-arn on epc)
    VARIANT_JSON="[{\"VariantName\":\"AllTraffic\",\"ModelName\":\"${MODEL_NAME_SM}\",\"InstanceType\":\"${INSTANCE_TYPE}\",\"InitialInstanceCount\":1}]"

    echo "⚙️  Creating endpoint configuration: ${ENDPOINT_CONFIG_NAME}"
    if ! aws sagemaker create-endpoint-config \
        --endpoint-config-name "${ENDPOINT_CONFIG_NAME}" \
        --production-variants "${VARIANT_JSON}" \
        --region "${AWS_REGION}"; then

        echo "❌ Failed to create endpoint configuration"
        echo "   Check that:"
        echo "   • The instance type is valid: ${INSTANCE_TYPE}"
        echo "   • The instance type is available in region: ${AWS_REGION}"
        echo "   • You have sufficient service quota for the instance type"
        exit 4
    fi

    echo "✅ Endpoint configuration created: ${ENDPOINT_CONFIG_NAME}"

    # Record endpoint config in manifest (non-blocking)
    ENDPOINT_CONFIG_ARN="arn:aws:sagemaker:${AWS_REGION}:${AWS_ACCOUNT_ID}:endpoint-config/${ENDPOINT_CONFIG_NAME}"
    ./do/manifest add \
        --type sagemaker-endpoint-config \
        --id "${ENDPOINT_CONFIG_ARN}" \
        --project "${PROJECT_NAME}" \
        --meta "{\"endpointConfigName\":\"${ENDPOINT_CONFIG_NAME}\",\"instanceType\":\"${INSTANCE_TYPE}\",\"region\":\"${AWS_REGION}\"}" \
        2>/dev/null || true

    # Step 3: Create endpoint
    echo "🚀 Creating endpoint: ${ENDPOINT_NAME}"
    if ! aws sagemaker create-endpoint \
        --endpoint-name "${ENDPOINT_NAME}" \
        --endpoint-config-name "${ENDPOINT_CONFIG_NAME}" \
        --region "${AWS_REGION}"; then

        echo "❌ Failed to create endpoint"
        echo "   Check that:"
        echo "   • Your IAM credentials have sagemaker:CreateEndpoint permission"
        echo "   • You have sufficient service quota in region: ${AWS_REGION}"
        exit 4
    fi

    echo "✅ Endpoint creation initiated: ${ENDPOINT_NAME}"

    # Record endpoint in manifest (non-blocking)
    ENDPOINT_ARN="arn:aws:sagemaker:${AWS_REGION}:${AWS_ACCOUNT_ID}:endpoint/${ENDPOINT_NAME}"
    ./do/manifest add \
        --type sagemaker-endpoint \
        --id "${ENDPOINT_ARN}" \
        --project "${PROJECT_NAME}" \
        --meta "{\"endpointName\":\"${ENDPOINT_NAME}\",\"instanceType\":\"${INSTANCE_TYPE}\",\"region\":\"${AWS_REGION}\"}" \
        2>/dev/null || true
fi

# ============================================================
# Wait for endpoint
# ============================================================
if [ -z "${SKIP_TO}" ] || [ "${SKIP_TO}" = "wait_endpoint" ]; then
    echo "⏳ Waiting for endpoint to reach InService status..."
    echo "   This may take several minutes..."
    echo "   If this times out, re-run ./do/deploy to resume."

    wait_endpoint "${ENDPOINT_NAME}"
fi

echo "✅ Deployment complete!"
echo ""
echo "📋 Deployment Details:"
echo "   Endpoint: ${ENDPOINT_NAME}"
echo "   Endpoint Config: ${ENDPOINT_CONFIG_NAME}"
echo "   Model: ${SAGEMAKER_MODEL_NAME:-${MODEL_NAME_SM:-N/A}}"
echo "   Model Package: ${MODEL_PACKAGE_ARN}"
echo "   Region: ${AWS_REGION}"
echo "   Instance Type: ${INSTANCE_TYPE}"
echo ""
echo "📋 What's next?"
echo "   • Test your endpoint:         ./do/test"
<% if (typeof includeBenchmark !== 'undefined' && includeBenchmark) { %>
echo "   • Benchmark performance:      ./do/benchmark"
<% } %>
echo "   • View endpoint status:       ./do/status"
echo "   • Register this deployment:   ./do/register"
echo "   • View logs:                  ./do/logs"
echo "   • Clean up when done:         ./do/clean endpoint"

<% } else if (deploymentTarget === 'async-inference') { %>
# ============================================================
# SageMaker Async Inference Deployment (Model-Based)
# Marketplace packages use: CreateModel(ModelPackageName) → CreateEndpointConfig(AsyncInferenceConfig) → CreateEndpoint
# ============================================================

# ============================================================
# Bootstrap async infrastructure (S3 bucket + SNS topics)
# ============================================================

# Extract bucket name from S3 output path
ASYNC_S3_BUCKET=$(echo "${ASYNC_S3_OUTPUT_PATH}" | sed 's|s3://||' | cut -d'/' -f1)

<% if (!asyncS3OutputPath) { %>
# Bootstrap default S3 bucket (check-and-create)
echo "🔍 Checking if S3 bucket exists: ${ASYNC_S3_BUCKET}"
if ! aws s3api head-bucket --bucket "${ASYNC_S3_BUCKET}" --region "${AWS_REGION}" 2>/dev/null; then
    echo "📦 Creating S3 bucket: ${ASYNC_S3_BUCKET}"
    if [ "${AWS_REGION}" = "us-east-1" ]; then
        if ! aws s3api create-bucket \
            --bucket "${ASYNC_S3_BUCKET}" \
            --region "${AWS_REGION}"; then
            echo "❌ Failed to create S3 bucket: ${ASYNC_S3_BUCKET}"
            exit 4
        fi
    else
        if ! aws s3api create-bucket \
            --bucket "${ASYNC_S3_BUCKET}" \
            --region "${AWS_REGION}" \
            --create-bucket-configuration LocationConstraint="${AWS_REGION}"; then
            echo "❌ Failed to create S3 bucket: ${ASYNC_S3_BUCKET}"
            exit 4
        fi
    fi
    echo "✅ S3 bucket created: ${ASYNC_S3_BUCKET}"
else
    echo "✅ S3 bucket exists: ${ASYNC_S3_BUCKET}"
fi
<% } else { %>
# Custom S3 output path provided — skip bucket creation
echo "✅ Using custom S3 output path: ${ASYNC_S3_OUTPUT_PATH}"
<% } %>

# Extract topic name from SNS success topic ARN
ASYNC_SNS_SUCCESS_TOPIC_NAME=$(echo "${ASYNC_SNS_SUCCESS_TOPIC}" | awk -F: '{print $NF}')

<% if (!asyncSnsSuccessTopic) { %>
# Bootstrap default SNS success topic (check-and-create)
echo "🔍 Checking if SNS success topic exists: ${ASYNC_SNS_SUCCESS_TOPIC_NAME}"
if ! aws sns get-topic-attributes --topic-arn "${ASYNC_SNS_SUCCESS_TOPIC}" --region "${AWS_REGION}" 2>/dev/null; then
    echo "📦 Creating SNS success topic: ${ASYNC_SNS_SUCCESS_TOPIC_NAME}"
    if ! aws sns create-topic \
        --name "${ASYNC_SNS_SUCCESS_TOPIC_NAME}" \
        --region "${AWS_REGION}" > /dev/null; then
        echo "❌ Failed to create SNS success topic"
        exit 4
    fi
    echo "✅ SNS success topic created: ${ASYNC_SNS_SUCCESS_TOPIC_NAME}"
else
    echo "✅ SNS success topic exists: ${ASYNC_SNS_SUCCESS_TOPIC_NAME}"
fi
<% } else { %>
# Custom SNS success topic ARN provided — skip topic creation
echo "✅ Using custom SNS success topic: ${ASYNC_SNS_SUCCESS_TOPIC}"
<% } %>

# Extract topic name from SNS error topic ARN
ASYNC_SNS_ERROR_TOPIC_NAME=$(echo "${ASYNC_SNS_ERROR_TOPIC}" | awk -F: '{print $NF}')

<% if (!asyncSnsErrorTopic) { %>
# Bootstrap default SNS error topic (check-and-create)
echo "🔍 Checking if SNS error topic exists: ${ASYNC_SNS_ERROR_TOPIC_NAME}"
if ! aws sns get-topic-attributes --topic-arn "${ASYNC_SNS_ERROR_TOPIC}" --region "${AWS_REGION}" 2>/dev/null; then
    echo "📦 Creating SNS error topic: ${ASYNC_SNS_ERROR_TOPIC_NAME}"
    if ! aws sns create-topic \
        --name "${ASYNC_SNS_ERROR_TOPIC_NAME}" \
        --region "${AWS_REGION}" > /dev/null; then
        echo "❌ Failed to create SNS error topic"
        exit 4
    fi
    echo "✅ SNS error topic created: ${ASYNC_SNS_ERROR_TOPIC_NAME}"
else
    echo "✅ SNS error topic exists: ${ASYNC_SNS_ERROR_TOPIC_NAME}"
fi
<% } else { %>
# Custom SNS error topic ARN provided — skip topic creation
echo "✅ Using custom SNS error topic: ${ASYNC_SNS_ERROR_TOPIC}"
<% } %>

# ============================================================
# Idempotency: check for existing deployment from a previous run
# ============================================================
SKIP_TO=""

if [ "${FORCE_NEW}" = true ]; then
    echo "🔄 --force: ignoring previous deployment, creating new resources."
elif [ -n "${ENDPOINT_NAME:-}" ]; then
    echo "🔍 Checking for existing deployment: ${ENDPOINT_NAME}"

    EP_STATUS=$(_get_endpoint_status "${ENDPOINT_NAME}")

    case "${EP_STATUS}" in
        InService)
            echo "✅ Async endpoint already InService: ${ENDPOINT_NAME}"
            echo ""
            echo "📋 Deployment is already live. Nothing to do."
            echo "   Endpoint: ${ENDPOINT_NAME}"
            echo ""
            echo "🧪 Test your async endpoint:"
            echo "   ./do/test"
            echo ""
            echo "🧹 Clean up when done:"
            echo "   ./do/clean endpoint"
            exit 0
            ;;
        Creating|Updating)
            echo "⏳ Endpoint still ${EP_STATUS}: ${ENDPOINT_NAME}"
            SKIP_TO="wait_endpoint"
            ;;
        Failed)
            echo "⚠️  Previous endpoint failed: ${ENDPOINT_NAME}"
            echo "   Creating a new deployment. Clean up the failed endpoint with:"
            echo "   ./do/clean endpoint"
            echo ""
            ;;
        "")
            echo "   Previous endpoint not found (may have been cleaned up). Creating new deployment."
            ;;
        *)
            echo "   Endpoint in unexpected state: ${EP_STATUS}. Creating new deployment."
            ;;
    esac
fi

# ============================================================
# Create async resources (skip if resuming from wait)
# ============================================================
if [ -z "${SKIP_TO}" ]; then
    TIMESTAMP=$(date +%s)
    MODEL_NAME_SM="${PROJECT_NAME}-mkt-async-model-${TIMESTAMP}"
    ENDPOINT_CONFIG_NAME="${PROJECT_NAME}-mkt-async-epc-${TIMESTAMP}"
    ENDPOINT_NAME="${PROJECT_NAME}-mkt-async-ep-${TIMESTAMP}"

    _update_config_var "ENDPOINT_NAME" "${ENDPOINT_NAME}"
    _update_config_var "ENDPOINT_CONFIG_NAME" "${ENDPOINT_CONFIG_NAME}"
    _update_config_var "SAGEMAKER_MODEL_NAME" "${MODEL_NAME_SM}"

    # Step 1: Create SageMaker model from Marketplace model package
    echo "📦 Creating SageMaker model from Marketplace package: ${MODEL_NAME_SM}"
    if ! aws sagemaker create-model \
        --model-name "${MODEL_NAME_SM}" \
        --primary-container "{\"ModelPackageName\":\"${MODEL_PACKAGE_ARN}\"}" \
        --execution-role-arn "${ROLE_ARN}" \
        --region "${AWS_REGION}"; then

        echo "❌ Failed to create model from package ARN. Check IAM permissions and subscription status."
        echo "   Check that:"
        echo "   • The model package ARN is correct: ${MODEL_PACKAGE_ARN}"
        echo "   • Your Marketplace subscription is active"
        echo "   • The execution role has permission to access the model package"
        exit 4
    fi

    echo "✅ SageMaker model created: ${MODEL_NAME_SM}"

    # Record model in manifest (non-blocking)
    MODEL_ARN="arn:aws:sagemaker:${AWS_REGION}:${AWS_ACCOUNT_ID}:model/${MODEL_NAME_SM}"
    ./do/manifest add \
        --type sagemaker-model \
        --id "${MODEL_ARN}" \
        --project "${PROJECT_NAME}" \
        --meta "{\"modelName\":\"${MODEL_NAME_SM}\",\"modelPackageArn\":\"${MODEL_PACKAGE_ARN}\",\"region\":\"${AWS_REGION}\"}" \
        2>/dev/null || true

    # Step 2: Build production variant and AsyncInferenceConfig
    VARIANT_JSON="[{\"VariantName\":\"AllTraffic\",\"ModelName\":\"${MODEL_NAME_SM}\",\"InstanceType\":\"${INSTANCE_TYPE}\",\"InitialInstanceCount\":1}]"

    ASYNC_CONFIG="{\"OutputConfig\":{\"S3OutputPath\":\"${ASYNC_S3_OUTPUT_PATH}\",\"NotificationConfig\":{\"SuccessTopic\":\"${ASYNC_SNS_SUCCESS_TOPIC}\",\"ErrorTopic\":\"${ASYNC_SNS_ERROR_TOPIC}\"}}"
    if [ -n "${ASYNC_MAX_CONCURRENT_INVOCATIONS:-}" ]; then
        ASYNC_CONFIG="${ASYNC_CONFIG},\"ClientConfig\":{\"MaxConcurrentInvocationsPerInstance\":${ASYNC_MAX_CONCURRENT_INVOCATIONS}}"
    fi
    ASYNC_CONFIG="${ASYNC_CONFIG}}"

    # Step 3: Create endpoint configuration with AsyncInferenceConfig
    echo "⚙️  Creating async endpoint configuration: ${ENDPOINT_CONFIG_NAME}"
    if ! aws sagemaker create-endpoint-config \
        --endpoint-config-name "${ENDPOINT_CONFIG_NAME}" \
        --production-variants "${VARIANT_JSON}" \
        --async-inference-config "${ASYNC_CONFIG}" \
        --region "${AWS_REGION}"; then

        echo "❌ Failed to create async endpoint configuration"
        echo "   Check that:"
        echo "   • The S3 output path is accessible: ${ASYNC_S3_OUTPUT_PATH}"
        echo "   • The IAM role has s3:PutObject permission on the output path"
        echo "   • The instance type is valid: ${INSTANCE_TYPE}"
        echo "   • You have sufficient service quota for the instance type"
        exit 4
    fi

    echo "✅ Async endpoint configuration created: ${ENDPOINT_CONFIG_NAME}"

    # Record endpoint config in manifest (non-blocking)
    ENDPOINT_CONFIG_ARN="arn:aws:sagemaker:${AWS_REGION}:${AWS_ACCOUNT_ID}:endpoint-config/${ENDPOINT_CONFIG_NAME}"
    ./do/manifest add \
        --type sagemaker-endpoint-config \
        --id "${ENDPOINT_CONFIG_ARN}" \
        --project "${PROJECT_NAME}" \
        --meta "{\"endpointConfigName\":\"${ENDPOINT_CONFIG_NAME}\",\"instanceType\":\"${INSTANCE_TYPE}\",\"region\":\"${AWS_REGION}\"}" \
        2>/dev/null || true

    # Step 4: Create endpoint
    echo "🚀 Creating async endpoint: ${ENDPOINT_NAME}"
    if ! aws sagemaker create-endpoint \
        --endpoint-name "${ENDPOINT_NAME}" \
        --endpoint-config-name "${ENDPOINT_CONFIG_NAME}" \
        --region "${AWS_REGION}"; then

        echo "❌ Failed to create async endpoint"
        echo "   Check that:"
        echo "   • Your IAM credentials have sagemaker:CreateEndpoint permission"
        echo "   • You have sufficient service quota in region: ${AWS_REGION}"
        exit 4
    fi

    echo "✅ Async endpoint creation initiated: ${ENDPOINT_NAME}"

    # Record endpoint in manifest (non-blocking)
    ENDPOINT_ARN="arn:aws:sagemaker:${AWS_REGION}:${AWS_ACCOUNT_ID}:endpoint/${ENDPOINT_NAME}"
    ./do/manifest add \
        --type sagemaker-endpoint \
        --id "${ENDPOINT_ARN}" \
        --project "${PROJECT_NAME}" \
        --meta "{\"endpointName\":\"${ENDPOINT_NAME}\",\"instanceType\":\"${INSTANCE_TYPE}\",\"region\":\"${AWS_REGION}\"}" \
        2>/dev/null || true
fi

# ============================================================
# Wait for endpoint
# ============================================================
if [ -z "${SKIP_TO}" ] || [ "${SKIP_TO}" = "wait_endpoint" ]; then
    echo "⏳ Waiting for async endpoint to reach InService status..."
    echo "   This may take several minutes..."
    echo "   If this times out, re-run ./do/deploy to resume."

    wait_endpoint "${ENDPOINT_NAME}"
fi

echo "✅ Async deployment complete!"
echo ""
echo "📋 Deployment Details:"
echo "   Endpoint: ${ENDPOINT_NAME}"
echo "   Endpoint Config: ${ENDPOINT_CONFIG_NAME}"
echo "   Model: ${SAGEMAKER_MODEL_NAME:-${MODEL_NAME_SM:-N/A}}"
echo "   Model Package: ${MODEL_PACKAGE_ARN}"
echo "   Region: ${AWS_REGION}"
echo "   Instance Type: ${INSTANCE_TYPE}"
echo "   S3 Output: ${ASYNC_S3_OUTPUT_PATH}"
echo "   SNS Success: ${ASYNC_SNS_SUCCESS_TOPIC}"
echo "   SNS Error: ${ASYNC_SNS_ERROR_TOPIC}"
echo ""
echo "📋 What's next?"
echo "   • Test your async endpoint:   ./do/test"
echo "   • Check async output:         aws s3 ls ${ASYNC_S3_OUTPUT_PATH}"
<% if (typeof includeBenchmark !== 'undefined' && includeBenchmark) { %>
echo "   • Benchmark performance:      ./do/benchmark"
<% } %>
echo "   • Register this deployment:   ./do/register"
echo "   • View logs:                  ./do/logs"
echo "   • Clean up when done:         ./do/clean endpoint"

<% } else if (deploymentTarget === 'batch-transform') { %>
# ============================================================
# SageMaker Batch Transform Deployment
# Marketplace packages use: CreateModel(ModelPackageName) → CreateTransformJob
# ============================================================

# Validate S3 input path
if [ -z "${BATCH_INPUT_PATH:-}" ]; then
    echo "❌ S3 input path not provided"
    echo ""
    echo "Set BATCH_INPUT_PATH in do/config or provide via CLI:"
    echo "  export BATCH_INPUT_PATH=s3://my-bucket/input/"
    echo "  ./do/deploy"
    exit 3
fi

if [[ "${BATCH_INPUT_PATH}" != s3://* ]]; then
    echo "❌ S3 input path must start with s3://"
    echo "   Current value: ${BATCH_INPUT_PATH}"
    exit 3
fi

# Validate S3 output path
if [ -z "${BATCH_OUTPUT_PATH:-}" ]; then
    echo "❌ S3 output path not provided"
    echo ""
    echo "Set BATCH_OUTPUT_PATH in do/config or provide via CLI:"
    echo "  export BATCH_OUTPUT_PATH=s3://my-bucket/output/"
    echo "  ./do/deploy"
    exit 3
fi

if [[ "${BATCH_OUTPUT_PATH}" != s3://* ]]; then
    echo "❌ S3 output path must start with s3://"
    echo "   Current value: ${BATCH_OUTPUT_PATH}"
    exit 3
fi

# ============================================================
# Bootstrap S3 buckets for batch transform
# ============================================================

BATCH_INPUT_BUCKET=$(echo "${BATCH_INPUT_PATH}" | sed 's|s3://||' | cut -d'/' -f1)
BATCH_OUTPUT_BUCKET=$(echo "${BATCH_OUTPUT_PATH}" | sed 's|s3://||' | cut -d'/' -f1)

<% if (!batchInputPath) { %>
# Bootstrap default S3 input bucket (check-and-create)
echo "🔍 Checking if S3 input bucket exists: ${BATCH_INPUT_BUCKET}"
if ! aws s3api head-bucket --bucket "${BATCH_INPUT_BUCKET}" --region "${AWS_REGION}" 2>/dev/null; then
    echo "📦 Creating S3 input bucket: ${BATCH_INPUT_BUCKET}"
    if [ "${AWS_REGION}" = "us-east-1" ]; then
        if ! aws s3api create-bucket \
            --bucket "${BATCH_INPUT_BUCKET}" \
            --region "${AWS_REGION}"; then
            echo "❌ Failed to create S3 input bucket: ${BATCH_INPUT_BUCKET}"
            exit 4
        fi
    else
        if ! aws s3api create-bucket \
            --bucket "${BATCH_INPUT_BUCKET}" \
            --region "${AWS_REGION}" \
            --create-bucket-configuration LocationConstraint="${AWS_REGION}"; then
            echo "❌ Failed to create S3 input bucket: ${BATCH_INPUT_BUCKET}"
            exit 4
        fi
    fi
    echo "✅ S3 input bucket created: ${BATCH_INPUT_BUCKET}"
else
    echo "✅ S3 input bucket exists: ${BATCH_INPUT_BUCKET}"
fi

# Upload sample input file if the input prefix is empty
EXISTING_OBJECTS=$(aws s3 ls "${BATCH_INPUT_PATH}" --region "${AWS_REGION}" 2>/dev/null | head -1 || true)
if [ -z "${EXISTING_OBJECTS}" ]; then
    echo "📄 Uploading sample input file to ${BATCH_INPUT_PATH}"
    echo '{"inputs": "What is machine learning?", "parameters": {"max_new_tokens": 50}}' | aws s3 cp - "${BATCH_INPUT_PATH}sample.jsonl" --region "${AWS_REGION}"
    echo "✅ Sample input uploaded: ${BATCH_INPUT_PATH}sample.jsonl"
    echo "   ⚠️  Replace this with your actual input data before running production jobs"
fi
<% } else { %>
# Custom S3 input path provided — skip bucket creation
echo "✅ Using custom S3 input path: ${BATCH_INPUT_PATH}"
<% } %>

<% if (!batchOutputPath) { %>
# Bootstrap default S3 output bucket (check-and-create, may be same as input)
if [ "${BATCH_OUTPUT_BUCKET}" != "${BATCH_INPUT_BUCKET}" ]; then
    echo "🔍 Checking if S3 output bucket exists: ${BATCH_OUTPUT_BUCKET}"
    if ! aws s3api head-bucket --bucket "${BATCH_OUTPUT_BUCKET}" --region "${AWS_REGION}" 2>/dev/null; then
        echo "📦 Creating S3 output bucket: ${BATCH_OUTPUT_BUCKET}"
        if [ "${AWS_REGION}" = "us-east-1" ]; then
            if ! aws s3api create-bucket \
                --bucket "${BATCH_OUTPUT_BUCKET}" \
                --region "${AWS_REGION}"; then
                echo "❌ Failed to create S3 output bucket: ${BATCH_OUTPUT_BUCKET}"
                exit 4
            fi
        else
            if ! aws s3api create-bucket \
                --bucket "${BATCH_OUTPUT_BUCKET}" \
                --region "${AWS_REGION}" \
                --create-bucket-configuration LocationConstraint="${AWS_REGION}"; then
                echo "❌ Failed to create S3 output bucket: ${BATCH_OUTPUT_BUCKET}"
                exit 4
            fi
        fi
        echo "✅ S3 output bucket created: ${BATCH_OUTPUT_BUCKET}"
    else
        echo "✅ S3 output bucket exists: ${BATCH_OUTPUT_BUCKET}"
    fi
else
    echo "✅ S3 output bucket same as input: ${BATCH_OUTPUT_BUCKET}"
fi
<% } else { %>
# Custom S3 output path provided — skip bucket creation
echo "✅ Using custom S3 output path: ${BATCH_OUTPUT_PATH}"
<% } %>

# ============================================================
# Check for previous transform job still running
# ============================================================
if [ "${FORCE_NEW}" != true ] && [ -n "${TRANSFORM_JOB_NAME:-}" ]; then
    echo "🔍 Checking previous transform job: ${TRANSFORM_JOB_NAME}"
    PREV_JOB_STATUS=$(aws sagemaker describe-transform-job \
        --transform-job-name "${TRANSFORM_JOB_NAME}" \
        --region "${AWS_REGION}" \
        --query "TransformJobStatus" \
        --output text 2>/dev/null || echo "")

    case "${PREV_JOB_STATUS}" in
        InProgress)
            echo "⚠️  Previous transform job is still running: ${TRANSFORM_JOB_NAME}"
            echo "   Wait for it to complete, or stop it with:"
            echo "   aws sagemaker stop-transform-job --transform-job-name ${TRANSFORM_JOB_NAME} --region ${AWS_REGION}"
            echo ""
            echo "   Use --force to create a new job anyway."
            exit 4
            ;;
        Completed)
            echo "✅ Previous transform job completed: ${TRANSFORM_JOB_NAME}"
            echo "   Creating a new job. Results from the previous job are in:"
            echo "   ${BATCH_OUTPUT_PATH}"
            echo ""
            ;;
        *)
            # Failed, Stopped, or not found — proceed with new job
            ;;
    esac
fi

# Generate unique names with timestamp
TIMESTAMP=$(date +%s)
MODEL_NAME_SM="${PROJECT_NAME}-mkt-batch-model-${TIMESTAMP}"
TRANSFORM_JOB_NAME="${PROJECT_NAME}-mkt-batch-job-${TIMESTAMP}"

_update_config_var "TRANSFORM_JOB_NAME" "${TRANSFORM_JOB_NAME}"
_update_config_var "SAGEMAKER_MODEL_NAME" "${MODEL_NAME_SM}"

# Step 1: Create SageMaker model from Marketplace model package
echo "📦 Creating SageMaker model from Marketplace package: ${MODEL_NAME_SM}"
if ! aws sagemaker create-model \
    --model-name "${MODEL_NAME_SM}" \
    --primary-container "{\"ModelPackageName\":\"${MODEL_PACKAGE_ARN}\"}" \
    --execution-role-arn "${ROLE_ARN}" \
    --region "${AWS_REGION}"; then

    echo "❌ Failed to create model from package ARN. Check IAM permissions and subscription status."
    echo "   Check that:"
    echo "   • The model package ARN is correct: ${MODEL_PACKAGE_ARN}"
    echo "   • Your Marketplace subscription is active"
    echo "   • The execution role has permission to access the model package"
    exit 4
fi

echo "✅ SageMaker model created: ${MODEL_NAME_SM}"

# Record model in manifest (non-blocking)
MODEL_ARN="arn:aws:sagemaker:${AWS_REGION}:${AWS_ACCOUNT_ID}:model/${MODEL_NAME_SM}"
./do/manifest add \
    --type sagemaker-model \
    --id "${MODEL_ARN}" \
    --project "${PROJECT_NAME}" \
    --meta "{\"modelName\":\"${MODEL_NAME_SM}\",\"modelPackageArn\":\"${MODEL_PACKAGE_ARN}\",\"region\":\"${AWS_REGION}\"}" \
    2>/dev/null || true

# Step 2: Build transform job JSON
TRANSFORM_JOB_JSON="{
    \"TransformJobName\": \"${TRANSFORM_JOB_NAME}\",
    \"ModelName\": \"${MODEL_NAME_SM}\",
    \"TransformInput\": {
        \"DataSource\": {
            \"S3DataSource\": {
                \"S3DataType\": \"S3Prefix\",
                \"S3Uri\": \"${BATCH_INPUT_PATH}\"
            }
        },
        \"ContentType\": \"application/json\",
        \"SplitType\": \"${BATCH_SPLIT_TYPE}\"
    },
    \"TransformOutput\": {
        \"S3OutputPath\": \"${BATCH_OUTPUT_PATH}\"
        $([ "${BATCH_JOIN_SOURCE:-None}" = "Input" ] && echo ",\"Accept\": \"application/json\", \"AssembleWith\": \"${BATCH_SPLIT_TYPE}\"")
    },
    \"TransformResources\": {
        \"InstanceType\": \"${INSTANCE_TYPE}\",
        \"InstanceCount\": ${BATCH_INSTANCE_COUNT}
    },
    \"MaxConcurrentTransforms\": ${BATCH_MAX_CONCURRENT_TRANSFORMS:-1},
    \"MaxPayloadInMB\": ${BATCH_MAX_PAYLOAD_IN_MB:-6},
    \"BatchStrategy\": \"${BATCH_STRATEGY}\"
    $([ "${BATCH_JOIN_SOURCE:-None}" = "Input" ] && echo ",\"DataProcessing\": { \"JoinSource\": \"Input\" }")
}"

# Step 3: Create transform job
echo "🚀 Creating transform job: ${TRANSFORM_JOB_NAME}"
if ! aws sagemaker create-transform-job \
    --cli-input-json "${TRANSFORM_JOB_JSON}" \
    --region "${AWS_REGION}"; then

    echo "❌ Failed to create transform job"
    echo "   Check that:"
    echo "   • The S3 input path exists and is accessible: ${BATCH_INPUT_PATH}"
    echo "   • The S3 output path is writable: ${BATCH_OUTPUT_PATH}"
    echo "   • The instance type is valid: ${INSTANCE_TYPE}"
    echo "   • You have sufficient service quota for the instance type"
    exit 4
fi

echo "✅ Transform job created: ${TRANSFORM_JOB_NAME}"

# Record transform job in manifest (non-blocking)
TRANSFORM_JOB_ARN="arn:aws:sagemaker:${AWS_REGION}:${AWS_ACCOUNT_ID}:transform-job/${TRANSFORM_JOB_NAME}"
./do/manifest add \
    --type sagemaker-transform-job \
    --id "${TRANSFORM_JOB_ARN}" \
    --project "${PROJECT_NAME}" \
    --meta "{\"transformJobName\":\"${TRANSFORM_JOB_NAME}\",\"modelName\":\"${MODEL_NAME_SM}\",\"instanceType\":\"${INSTANCE_TYPE}\",\"region\":\"${AWS_REGION}\"}" \
    2>/dev/null || true

# Step 4: Poll transform job status until completion or failure
echo "⏳ Waiting for transform job to complete..."
echo "   This may take several minutes depending on dataset size..."
echo "   If this times out, check status with:"
echo "   aws sagemaker describe-transform-job --transform-job-name ${TRANSFORM_JOB_NAME} --region ${AWS_REGION}"
echo ""

while true; do
    JOB_STATUS=$(aws sagemaker describe-transform-job \
        --transform-job-name "${TRANSFORM_JOB_NAME}" \
        --region "${AWS_REGION}" \
        --query "TransformJobStatus" \
        --output text 2>&1) || {
        if echo "${JOB_STATUS}" | grep -qi "expired\|token"; then
            echo ""
            echo "⚠️  Credentials expired, but the transform job is still running."
            echo "   Refresh your credentials and check status with:"
            echo "   aws sagemaker describe-transform-job --transform-job-name ${TRANSFORM_JOB_NAME} --region ${AWS_REGION} --query TransformJobStatus"
            exit 4
        fi
        echo "❌ Failed to describe transform job: ${TRANSFORM_JOB_NAME}"
        echo "   Error: ${JOB_STATUS}"
        exit 4
    }

    case "${JOB_STATUS}" in
        Completed)
            echo "✅ Transform job completed successfully!"
            break
            ;;
        Failed)
            FAILURE_REASON=$(aws sagemaker describe-transform-job \
                --transform-job-name "${TRANSFORM_JOB_NAME}" \
                --region "${AWS_REGION}" \
                --query "FailureReason" \
                --output text 2>/dev/null || echo "Unknown")
            echo "❌ Transform job failed"
            echo "   Reason: ${FAILURE_REASON}"
            echo ""
            echo "   Check CloudWatch Logs for details:"
            echo "   https://console.aws.amazon.com/cloudwatch/home?region=${AWS_REGION}#logsV2:log-groups/log-group//aws/sagemaker/TransformJobs"
            exit 4
            ;;
        Stopped)
            echo "⚠️  Transform job was stopped"
            exit 4
            ;;
        InProgress)
            echo "   $(date +%H:%M:%S) Job status: InProgress..."
            sleep 30
            ;;
        *)
            echo "   $(date +%H:%M:%S) Job status: ${JOB_STATUS}..."
            sleep 30
            ;;
    esac
done

echo ""
echo "📋 Deployment Details:"
echo "   Transform Job: ${TRANSFORM_JOB_NAME}"
echo "   Model: ${MODEL_NAME_SM}"
echo "   Model Package: ${MODEL_PACKAGE_ARN}"
echo "   Region: ${AWS_REGION}"
echo "   Instance Type: ${INSTANCE_TYPE}"
echo "   Instance Count: ${BATCH_INSTANCE_COUNT}"
echo "   S3 Input: ${BATCH_INPUT_PATH}"
echo "   S3 Output: ${BATCH_OUTPUT_PATH}"
echo "   Split Type: ${BATCH_SPLIT_TYPE}"
echo "   Strategy: ${BATCH_STRATEGY}"
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 "📋 What's next?"
echo "   • View results:               cat batch-output/"
echo "   • Review results:             ./do/test"
echo "   • Register this deployment:   ./do/register"
echo "   • View logs:                  ./do/logs"
echo "   • Clean up when done:         ./do/clean"

<% } %>
