#!/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"

# ── Profile-resolved variables (env var > profile > default) ──────────────────
ROLE_ARN="${ROLE_ARN:-${_PROFILE[roleArn]:-}}"
ECR_REPOSITORY_NAME="${ECR_REPOSITORY_NAME:-${_PROFILE[ecrRepositoryName]:-ml-container-creator}}"

# ============================================================
# Register deployment to the deployment registry
# ============================================================

# Parse command line arguments
NOTES=""
STATUS="success"
PROJECT_FLAG=""
JSON_OUTPUT=false
CI_MODE=false
CI_TABLE_NAME="${CI_TABLE_NAME:-mlcc-ci-table}"
CI_BUILD_STRATEGY="codebuild-submit"
BENCHMARK_STATUS=""
BENCHMARK_RUN_ID=""

while [[ $# -gt 0 ]]; do
    case "$1" in
        --notes)
            NOTES="$2"
            shift 2
            ;;
        --notes=*)
            NOTES="${1#*=}"
            shift
            ;;
        --status)
            STATUS="$2"
            shift 2
            ;;
        --status=*)
            STATUS="${1#*=}"
            shift
            ;;
        --project)
            PROJECT_FLAG="--project"
            shift
            ;;
        --json)
            JSON_OUTPUT=true
            shift
            ;;
        --ci)
            CI_MODE=true
            JSON_OUTPUT=true
            shift
            ;;
        --ci-table)
            CI_TABLE_NAME="$2"
            shift 2
            ;;
        --ci-table=*)
            CI_TABLE_NAME="${1#*=}"
            shift
            ;;
        --build-strategy)
            CI_BUILD_STRATEGY="$2"
            shift 2
            ;;
        --build-strategy=*)
            CI_BUILD_STRATEGY="${1#*=}"
            shift
            ;;
        --benchmark-status)
            BENCHMARK_STATUS="$2"
            shift 2
            ;;
        --benchmark-status=*)
            BENCHMARK_STATUS="${1#*=}"
            shift
            ;;
        --benchmark-run-id)
            BENCHMARK_RUN_ID="$2"
            shift 2
            ;;
        --benchmark-run-id=*)
            BENCHMARK_RUN_ID="${1#*=}"
            shift
            ;;
        *)
            echo "⚠️  Unknown option: $1"
            echo ""
            echo "Usage: ./do/register [--notes \"text\"] [--status success|partial|failed] [--project] [--json] [--ci] [--ci-table <name>] [--build-strategy <strategy>]"
            echo "       ./do/register --benchmark-status <completed|failed|in-progress> --benchmark-run-id <run-id>"
            exit 1
            ;;
    esac
done

# ============================================================
# Handle --benchmark-status mode (Stage 2 failure isolation)
# This mode ONLY writes benchmark fields to DynamoDB.
# It uses an UpdateExpression that sets ONLY:
#   lastBenchmarkRunId, lastBenchmarkTimestamp, lastBenchmarkStatus
# It does NOT touch testStatus, configJson, or any other fields.
# This ensures Stage 2 failure does not change the CI gate result.
# Requirements: 1.4, 7.3
# ============================================================

if [ -n "${BENCHMARK_STATUS}" ]; then
    # Validate benchmark status
    case "${BENCHMARK_STATUS}" in
        completed|failed|in-progress) ;;
        *)
            echo "❌ Invalid benchmark status: ${BENCHMARK_STATUS}"
            echo "   Valid values: completed, failed, in-progress"
            exit 1
            ;;
    esac

    if [ -z "${BENCHMARK_RUN_ID}" ]; then
        echo "❌ --benchmark-run-id is required when using --benchmark-status"
        exit 1
    fi

    # Compute configId (same hash logic used by write_ci_record below)
    _benchmark_input="${DEPLOYMENT_CONFIG}:${MODEL_NAME:-none}:${INSTANCE_TYPE}:${AWS_REGION}:${DEPLOYMENT_TARGET}:ic${IC_COUNT:-1}:adapt${ADAPTER_COUNT:-0}"
    if command -v sha256sum &> /dev/null; then
        CONFIG_ID=$(echo -n "$_benchmark_input" | sha256sum | cut -c1-16)
    else
        CONFIG_ID=$(echo -n "$_benchmark_input" | shasum -a 256 | cut -c1-16)
    fi
    BENCHMARK_TIMESTAMP=$(date -u +%Y-%m-%dT%H:%M:%SZ)

    echo "📊 Recording benchmark status for configId: ${CONFIG_ID}"
    echo "   Benchmark run ID:  ${BENCHMARK_RUN_ID}"
    echo "   Benchmark status:  ${BENCHMARK_STATUS}"
    echo "   Timestamp:         ${BENCHMARK_TIMESTAMP}"
    echo ""

    # Check if CI_Table exists
    if ! aws dynamodb describe-table --table-name "${CI_TABLE_NAME}" --region "${AWS_REGION}" &>/dev/null; then
        echo "⚠️  CI infrastructure not provisioned. Run 'ml-container-creator bootstrap' with CI enabled."
        echo "   Skipping benchmark status write."
        exit 0
    fi

    # UpdateExpression ONLY sets the 3 benchmark fields.
    # This is a SET operation — it does NOT overwrite the entire record (PutItem).
    # testStatus, configJson, and all other fields remain unchanged.
    if aws dynamodb update-item \
        --region "${AWS_REGION}" \
        --table-name "${CI_TABLE_NAME}" \
        --key "{\"configId\": {\"S\": \"${CONFIG_ID}\"}}" \
        --update-expression "SET lastBenchmarkRunId = :rid, lastBenchmarkTimestamp = :ts, lastBenchmarkStatus = :bs" \
        --expression-attribute-values "{
            \":rid\": {\"S\": \"${BENCHMARK_RUN_ID}\"},
            \":ts\": {\"S\": \"${BENCHMARK_TIMESTAMP}\"},
            \":bs\": {\"S\": \"${BENCHMARK_STATUS}\"}
        }" 2>/dev/null; then
        echo "✅ Benchmark status recorded: ${BENCHMARK_STATUS}"
    else
        echo "❌ Failed to write benchmark status for configId: ${CONFIG_ID}"
        exit 1
    fi

    exit 0
fi

# Validate status
case "${STATUS}" in
    success|partial|failed) ;;
    *)
        echo "❌ Invalid status: ${STATUS}"
        echo "   Valid values: success, partial, failed"
        exit 1
        ;;
esac

# ============================================================
# Build IC list from do/ic/ directory (multi-IC support)
# ============================================================

<% if (deploymentTarget === 'realtime-inference') { %>
IC_LIST_JSON="[]"
if [ -d "${SCRIPT_DIR}/ic" ]; then
    # Build IC list from all conf files (alphabetical order)
    IC_ENTRIES=""
    IC_COUNT=0
    for conf in "${SCRIPT_DIR}"/ic/*.conf; do
        [ -f "${conf}" ] || continue

        # Source the IC config to get its variables
        (
            # Subshell to avoid polluting current environment
            source "${conf}" 2>/dev/null
            echo "${IC_DEPLOYED_NAME:-}|${IC_IMAGE_TAG:-}|${IC_GPU_COUNT:-1}|${IC_COPY_COUNT:-1}"
        ) > /tmp/ic_entry_$$ 2>/dev/null

        IC_ENTRY=$(cat /tmp/ic_entry_$$ 2>/dev/null || echo "|||")
        rm -f /tmp/ic_entry_$$

        IC_BASENAME=$(basename "${conf}" .conf)
        IC_ENTRY_IMAGE=$(echo "${IC_ENTRY}" | cut -d'|' -f2)
        IC_ENTRY_GPU=$(echo "${IC_ENTRY}" | cut -d'|' -f3)
        IC_ENTRY_COPY=$(echo "${IC_ENTRY}" | cut -d'|' -f4)

        if [ -n "${IC_ENTRIES}" ]; then
            IC_ENTRIES="${IC_ENTRIES},"
        fi
        IC_ENTRIES="${IC_ENTRIES}{\"name\":\"${IC_BASENAME}\",\"image\":\"${IC_ENTRY_IMAGE}\",\"gpuCount\":${IC_ENTRY_GPU:-1},\"copyCount\":${IC_ENTRY_COPY:-1}}"
        IC_COUNT=$((IC_COUNT + 1))
    done

    if [ "${CI_MODE}" = true ] && [ ${IC_COUNT} -gt 1 ]; then
        # CI mode: only include the first IC (alphabetically) to keep CI costs down
        FIRST_CONF=$(ls "${SCRIPT_DIR}"/ic/*.conf 2>/dev/null | head -1)
        if [ -n "${FIRST_CONF}" ]; then
            (
                source "${FIRST_CONF}" 2>/dev/null
                echo "${IC_DEPLOYED_NAME:-}|${IC_IMAGE_TAG:-}|${IC_GPU_COUNT:-1}|${IC_COPY_COUNT:-1}"
            ) > /tmp/ic_first_$$ 2>/dev/null

            FIRST_ENTRY=$(cat /tmp/ic_first_$$ 2>/dev/null || echo "|||")
            rm -f /tmp/ic_first_$$

            FIRST_BASENAME=$(basename "${FIRST_CONF}" .conf)
            FIRST_IMAGE=$(echo "${FIRST_ENTRY}" | cut -d'|' -f2)
            FIRST_GPU=$(echo "${FIRST_ENTRY}" | cut -d'|' -f3)
            FIRST_COPY=$(echo "${FIRST_ENTRY}" | cut -d'|' -f4)

            IC_LIST_JSON="[{\"name\":\"${FIRST_BASENAME}\",\"image\":\"${FIRST_IMAGE}\",\"gpuCount\":${FIRST_GPU:-1},\"copyCount\":${FIRST_COPY:-1}}]"
        fi
    else
        IC_LIST_JSON="[${IC_ENTRIES}]"
    fi
else
    # Legacy: single IC from do/config
    IC_LIST_JSON="[{\"name\":\"default\",\"image\":\"${PROJECT_NAME}-latest\",\"gpuCount\":${IC_GPU_COUNT:-1},\"copyCount\":${IC_COPY_COUNT:-1}}]"
fi

# Append adapter entries from do/adapters/*.conf
ADAPTER_COUNT=0
if [ -d "${SCRIPT_DIR}/adapters" ]; then
    ADAPTER_ENTRIES=""
    for conf in "${SCRIPT_DIR}"/adapters/*.conf; do
        [ -f "${conf}" ] || continue
        [[ "$(basename "${conf}")" == ".gitkeep" ]] && continue

        ADAPTER_NAME_VAL=""
        ADAPTER_WEIGHTS_VAL=""
        ADAPTER_IC_VAL=""
        eval "$(grep '^export ADAPTER_NAME=' "${conf}" 2>/dev/null)"
        eval "$(grep '^export ADAPTER_WEIGHTS_URI=' "${conf}" 2>/dev/null)"
        eval "$(grep '^export ADAPTER_IC_NAME=' "${conf}" 2>/dev/null)"
        ADAPTER_NAME_VAL="${ADAPTER_NAME:-$(basename "${conf}" .conf)}"
        ADAPTER_WEIGHTS_VAL="${ADAPTER_WEIGHTS_URI:-}"
        ADAPTER_IC_VAL="${ADAPTER_IC_NAME:-}"

        if [ -n "${ADAPTER_ENTRIES}" ]; then
            ADAPTER_ENTRIES="${ADAPTER_ENTRIES},"
        fi
        ADAPTER_ENTRIES="${ADAPTER_ENTRIES}{\"name\":\"${ADAPTER_NAME_VAL}\",\"isAdapter\":true,\"baseIcName\":\"${ADAPTER_IC_VAL}\",\"artifactUrl\":\"${ADAPTER_WEIGHTS_VAL}\",\"gpuCount\":0,\"copyCount\":1}"
        ADAPTER_COUNT=$((ADAPTER_COUNT + 1))
        unset ADAPTER_NAME ADAPTER_WEIGHTS_URI ADAPTER_IC_NAME
    done

    if [ -n "${ADAPTER_ENTRIES}" ] && [ "${IC_LIST_JSON}" != "[]" ]; then
        # Append adapters to existing IC list
        IC_LIST_JSON="${IC_LIST_JSON%]},${ADAPTER_ENTRIES}]"
    elif [ -n "${ADAPTER_ENTRIES}" ]; then
        IC_LIST_JSON="[${ADAPTER_ENTRIES}]"
    fi
fi
<% } %>

# ============================================================
# Derive architecture and backend from DEPLOYMENT_CONFIG
# ============================================================

# DEPLOYMENT_CONFIG format: <architecture>-<backend> (e.g., transformers-vllm, http-flask, triton-fil)
# Special case: marketplace has no backend
if [ "${DEPLOYMENT_CONFIG}" = "marketplace" ]; then
    ARCHITECTURE="marketplace"
    BACKEND=""
else
    ARCHITECTURE="${DEPLOYMENT_CONFIG%%-*}"
    BACKEND="${DEPLOYMENT_CONFIG#*-}"
fi

echo "📋 Registering deployment to registry"
echo "   Project: ${PROJECT_NAME}"
echo "   Deployment config: ${DEPLOYMENT_CONFIG}"
echo "   Architecture: ${ARCHITECTURE}"
echo "   Backend: ${BACKEND}"
echo ""

# ============================================================
# Engine prefix mapping for transformer env var filtering
# ============================================================

get_engine_prefix() {
    local backend="$1"
    case "${backend}" in
        vllm)         echo "VLLM_" ;;
        vllm-omni)    echo "VLLM_" ;;
        sglang)       echo "SGLANG_" ;;
        tensorrt-llm) echo "TRTLLM_" ;;
        lmi)          echo "LMI_" ;;
        djl)          echo "DJL_" ;;
        *)            echo "" ;;
    esac
}

# ============================================================
# Build PARAMETERS JSON from shell environment (sourced by do/config)
# ============================================================

PARAMETERS="{}"

<%
// Build a helper to determine sensitive keys for redaction
const sensitivePatterns = ['HF_TOKEN', 'AWS_SECRET_ACCESS_KEY', 'AWS_SESSION_TOKEN'];
const sensitiveSubstrings = ['SECRET', 'TOKEN'];

function isSensitiveKey(key) {
    if (sensitivePatterns.includes(key)) return true;
    const upper = key.toUpperCase();
    return sensitiveSubstrings.some(sub => upper.includes(sub));
}

function redactValue(key, value) {
    return isSensitiveKey(key) ? '***REDACTED***' : value;
}
%>

case "${ARCHITECTURE}" in
    transformers|diffusors)
        # Extract env vars matching engine prefix + HF_TOKEN + HF_MODEL_ID
        ENGINE_PREFIX=$(get_engine_prefix "${BACKEND}")

        PARAMETERS=$(python3 -c "
import os, json
prefix = '${ENGINE_PREFIX}'
result = {}
for key, value in os.environ.items():
    if prefix and key.startswith(prefix):
        result[key] = value
    elif key == 'HF_TOKEN':
        result[key] = '***REDACTED***'
    elif key == 'HF_MODEL_ID':
        result[key] = value
<% if (modelEnvVars && Object.keys(modelEnvVars).length > 0) { %>
# Include model env vars
<% Object.entries(modelEnvVars).forEach(([key, value]) => { %>
result['<%= key %>'] = '<%= redactValue(key, value) %>'
<% }); %>
<% } %>
<% if (serverEnvVars && Object.keys(serverEnvVars).length > 0) { %>
# Include server env vars (with engine prefix applied at generation time)
<% Object.entries(serverEnvVars).forEach(([key, value]) => { %>
result['<%= key %>'] = '<%= redactValue(key, value) %>'
<% }); %>
<% } %>
print(json.dumps(result))
" 2>/dev/null || echo "{}")
        ;;

    http)
        # Extract all env vars from do/config except system vars, redact secrets
        PARAMETERS=$(python3 -c "
import os, json
exclude = {
    'PATH', 'PYTHONPATH', 'SAGEMAKER_BIND_TO_PORT', 'LANG', 'GPG_KEY',
    'PYTHON_VERSION', 'PYTHON_PIP_VERSION', 'PYTHON_SETUPTOOLS_VERSION',
    'PYTHON_GET_PIP_URL', 'PYTHON_GET_PIP_SHA256'
}
redact = {'HF_TOKEN', 'AWS_SECRET_ACCESS_KEY', 'AWS_SESSION_TOKEN'}
sensitive_substrings = ['SECRET', 'TOKEN']
result = {}
for key, value in os.environ.items():
    if key not in exclude:
        is_sensitive = key in redact or any(sub in key.upper() for sub in sensitive_substrings)
        result[key] = '***REDACTED***' if is_sensitive else value
<% if (modelEnvVars && Object.keys(modelEnvVars).length > 0) { %>
# Include model env vars
<% Object.entries(modelEnvVars).forEach(([key, value]) => { %>
result['<%= key %>'] = '<%= redactValue(key, value) %>'
<% }); %>
<% } %>
<% if (serverEnvVars && Object.keys(serverEnvVars).length > 0) { %>
# Include server env vars
<% Object.entries(serverEnvVars).forEach(([key, value]) => { %>
result['<%= key %>'] = '<%= redactValue(key, value) %>'
<% }); %>
<% } %>
print(json.dumps(result))
" 2>/dev/null || echo "{}")
        ;;

    triton)
        # Read config.pbtxt from model repository
        TRITON_MODEL_REPO="${TRITON_MODEL_REPOSITORY:-/opt/ml/model/model_repository}"
        CONFIG_PBTXT=""

        # Try to find config.pbtxt in the project directory
        PROJECT_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)"
        if [ -f "${PROJECT_DIR}/model_repository/model/config.pbtxt" ]; then
            CONFIG_PBTXT=$(cat "${PROJECT_DIR}/model_repository/model/config.pbtxt" 2>/dev/null || echo "")
        elif [ -f "${PROJECT_DIR}/model_repository/config.pbtxt" ]; then
            CONFIG_PBTXT=$(cat "${PROJECT_DIR}/model_repository/config.pbtxt" 2>/dev/null || echo "")
        fi

        # Read TRITON_MODEL_REPOSITORY from shell env (exported by do/config)
        TRITON_REPO_VAR="${TRITON_MODEL_REPOSITORY:-}"

        # Build parameters JSON (pipe config.pbtxt via stdin to avoid quoting issues)
        PARAMETERS=$(printf '%s' "${CONFIG_PBTXT}" | python3 -c "
import sys, json
config_pbtxt = sys.stdin.read()
result = {}
if config_pbtxt.strip():
    result['config.pbtxt'] = config_pbtxt.strip()
triton_repo = '${TRITON_REPO_VAR}'
if triton_repo:
    result['TRITON_MODEL_REPOSITORY'] = triton_repo
<% if (modelEnvVars && Object.keys(modelEnvVars).length > 0) { %>
# Include model env vars
<% Object.entries(modelEnvVars).forEach(([key, value]) => { %>
result['<%= key %>'] = '<%= redactValue(key, value) %>'
<% }); %>
<% } %>
<% if (serverEnvVars && Object.keys(serverEnvVars).length > 0) { %>
# Include server env vars
<% Object.entries(serverEnvVars).forEach(([key, value]) => { %>
result['<%= key %>'] = '<%= redactValue(key, value) %>'
<% }); %>
<% } %>
print(json.dumps(result))
" 2>/dev/null || echo "{}")
        ;;
esac

# Count captured parameters
PARAM_COUNT=$(echo "${PARAMETERS}" | python3 -c "import sys,json; print(len(json.load(sys.stdin)))" 2>/dev/null || echo "0")

# ============================================================
# Display summary before writing
# ============================================================

echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo "📋 Registration Summary"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo "   Deployment config: ${DEPLOYMENT_CONFIG}"
echo "   Architecture:      ${ARCHITECTURE}"
echo "   Backend:           ${BACKEND}"
<% if (framework === 'transformers') { %>
echo "   Model name:        ${MODEL_NAME:-N/A}"
<% } %>
<% if (deploymentTarget === 'realtime-inference') { %>
echo "   Instance type:     ${INSTANCE_TYPE}"
<% } else if (deploymentTarget === 'batch-transform') { %>
echo "   Instance:          ${INSTANCE_TYPE} x ${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 "   Region:            ${AWS_REGION}"
echo "   Status:            ${STATUS}"
echo "   Env var source:    do/config"
echo "   Parameters:        ${PARAM_COUNT} captured"
if [ -n "${BASE_IMAGE:-}" ]; then
echo "   Base image:        ${BASE_IMAGE}"
fi
if [ -n "${NOTES}" ]; then
echo "   Notes:             ${NOTES}"
fi
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo ""

# ============================================================
# Compute configId (deterministic hash of canonical fields)
# ============================================================

compute_config_id() {
    local input="${DEPLOYMENT_CONFIG}:${MODEL_NAME:-none}:${INSTANCE_TYPE}:${AWS_REGION}:${DEPLOYMENT_TARGET}:ic${IC_COUNT:-1}:adapt${ADAPTER_COUNT:-0}"
    # Use sha256sum (Linux) with fallback to shasum (macOS)
    if command -v sha256sum &> /dev/null; then
        echo -n "$input" | sha256sum | cut -c1-16
    else
        echo -n "$input" | shasum -a 256 | cut -c1-16
    fi
}

# ============================================================
# CI Table write logic (--ci flag)
# ============================================================

write_ci_record() {
    local config_id="$1"

    # Extract promoted attributes
    local promoted_deployment_config="${DEPLOYMENT_CONFIG}"
    local promoted_base_image="${BASE_IMAGE:-}"
    local promoted_project_name="${PROJECT_NAME}"

    # Extract baseImageVersion from BASE_IMAGE tag (e.g., "vllm/vllm-openai:v0.8.5" -> "v0.8.5")
    local promoted_base_image_version=""
    if [ -n "${promoted_base_image}" ]; then
        case "${promoted_base_image}" in
            *:*)
                promoted_base_image_version="${promoted_base_image##*:}"
                ;;
        esac
    fi

    # Build the configJson object from current config variables
    local config_json
    config_json=$(cat <<CJEOF
{
  "projectName": "${PROJECT_NAME}",
  "deploymentConfig": "${DEPLOYMENT_CONFIG}",
  "architecture": "${ARCHITECTURE}",
  "backend": "${BACKEND}",
  "modelName": "${MODEL_NAME:-}",
  "instanceType": "${INSTANCE_TYPE}",
  "awsRegion": "${AWS_REGION}",
  "deploymentTarget": "${DEPLOYMENT_TARGET}",
  "buildTarget": "${BUILD_TARGET}",
  "baseImage": "${BASE_IMAGE:-}",
  "buildStrategy": "${CI_BUILD_STRATEGY}",
<% if (endpointInitialInstanceCount != null || endpointDataCapturePercent != null || endpointVariantName != null || endpointVolumeSize != null) { %>
  "endpointConfig": {
<% if (endpointInitialInstanceCount != null) { %>
    "initialInstanceCount": ${ENDPOINT_INITIAL_INSTANCE_COUNT}<%= (endpointDataCapturePercent != null || endpointVariantName != null || endpointVolumeSize != null) ? ',' : '' %>
<% } %>
<% if (endpointDataCapturePercent != null) { %>
    "dataCapturePercent": ${ENDPOINT_DATA_CAPTURE_PERCENT}<%= (endpointVariantName != null || endpointVolumeSize != null) ? ',' : '' %>
<% } %>
<% if (endpointVariantName != null) { %>
    "variantName": "${ENDPOINT_VARIANT_NAME}"<%= (endpointVolumeSize != null) ? ',' : '' %>
<% } %>
<% if (endpointVolumeSize != null) { %>
    "volumeSize": ${ENDPOINT_VOLUME_SIZE}
<% } %>
  },
<% } %>
<% if (icCpuCount != null || icMemorySize != null || icGpuCount != null || icCopyCount != null || icModelWeight != null) { %>
  "icConfig": {
<% if (icCpuCount != null) { %>
    "cpuCount": ${IC_CPU_COUNT}<%= (icMemorySize != null || icGpuCount != null || icCopyCount != null || icModelWeight != null) ? ',' : '' %>
<% } %>
<% if (icMemorySize != null) { %>
    "memorySize": ${IC_MEMORY_SIZE}<%= (icGpuCount != null || icCopyCount != null || icModelWeight != null) ? ',' : '' %>
<% } %>
<% if (icGpuCount != null) { %>
    "gpuCount": ${IC_GPU_COUNT}<%= (icCopyCount != null || icModelWeight != null) ? ',' : '' %>
<% } %>
<% if (icCopyCount != null) { %>
    "copyCount": ${IC_COPY_COUNT}<%= (icModelWeight != null) ? ',' : '' %>
<% } %>
<% if (icModelWeight != null) { %>
    "modelWeight": ${IC_MODEL_WEIGHT}
<% } %>
  },
<% } %>
<% if (deploymentTarget === 'realtime-inference') { %>
  "icList": ${IC_LIST_JSON},
<% } %>
  "parameters": ${PARAMETERS}
}
CJEOF
)

    # Compact the JSON (remove whitespace)
    config_json=$(echo "$config_json" | python3 -c "import sys,json; print(json.dumps(json.load(sys.stdin), separators=(',',':')))" 2>/dev/null || echo "$config_json")

    local created_at
    created_at=$(date -u +%Y-%m-%dT%H:%M:%SZ)

    # Escape configJson for DynamoDB JSON attribute (double-escape quotes)
    local escaped_config_json
    escaped_config_json=$(echo "$config_json" | python3 -c "import sys,json; print(json.dumps(sys.stdin.read().strip()))" 2>/dev/null)
    # Remove the outer quotes added by json.dumps for string
    escaped_config_json="${escaped_config_json:1:${#escaped_config_json}-2}"

    # Try put-item with condition (new record)
    if aws dynamodb put-item \
        --region "${AWS_REGION}" \
        --table-name "${CI_TABLE_NAME}" \
        --item "{
            \"configId\": {\"S\": \"${config_id}\"},
            \"schemaVersion\": {\"N\": \"1\"},
            \"configJson\": {\"S\": \"${escaped_config_json}\"},
            \"testStatus\": {\"S\": \"untested\"},
            \"lastTestTimestamp\": {\"S\": \"1970-01-01T00:00:00Z\"},
            \"deploymentConfig\": {\"S\": \"${promoted_deployment_config}\"},
            \"baseImage\": {\"S\": \"${promoted_base_image}\"},
            \"baseImageVersion\": {\"S\": \"${promoted_base_image_version}\"},
            \"buildStrategy\": {\"S\": \"${CI_BUILD_STRATEGY}\"},
            \"projectName\": {\"S\": \"${promoted_project_name}\"},
            \"createdAt\": {\"S\": \"${created_at}\"}
        }" \
        --condition-expression "attribute_not_exists(configId)" 2>/dev/null; then
        echo "✅ New CI record created: ${config_id}"
    else
        # Record already exists — update it (reset testStatus, update configJson, preserve createdAt)
        if aws dynamodb update-item \
            --region "${AWS_REGION}" \
            --table-name "${CI_TABLE_NAME}" \
            --key "{\"configId\": {\"S\": \"${config_id}\"}}" \
            --update-expression "SET configJson = :cj, testStatus = :ts, deploymentConfig = :dc, baseImage = :bi, baseImageVersion = :bv, buildStrategy = :bs, projectName = :pn, schemaVersion = :sv" \
            --expression-attribute-values "{
                \":cj\": {\"S\": \"${escaped_config_json}\"},
                \":ts\": {\"S\": \"untested\"},
                \":dc\": {\"S\": \"${promoted_deployment_config}\"},
                \":bi\": {\"S\": \"${promoted_base_image}\"},
                \":bv\": {\"S\": \"${promoted_base_image_version}\"},
                \":bs\": {\"S\": \"${CI_BUILD_STRATEGY}\"},
                \":pn\": {\"S\": \"${promoted_project_name}\"},
                \":sv\": {\"N\": \"1\"}
            }" 2>/dev/null; then
            echo "✅ Existing CI record updated (testStatus reset to untested): ${config_id}"
        else
            echo "❌ Failed to update CI record: ${config_id}"
            return 1
        fi
    fi
}

# ============================================================
# Handle --json and --ci output modes
# ============================================================

if [ "${JSON_OUTPUT}" = true ] || [ "${CI_MODE}" = true ]; then
    # Build JSON deployment entry
    DEPLOYMENT_JSON=$(cat <<DJEOF
{
  "projectName": "${PROJECT_NAME}",
  "deploymentConfig": "${DEPLOYMENT_CONFIG}",
  "architecture": "${ARCHITECTURE}",
  "backend": "${BACKEND}",
  "modelName": "${MODEL_NAME:-}",
  "instanceType": "${INSTANCE_TYPE}",
  "awsRegion": "${AWS_REGION}",
  "deploymentTarget": "${DEPLOYMENT_TARGET}",
  "buildTarget": "${BUILD_TARGET}",
  "baseImage": "${BASE_IMAGE:-}",
  "status": "${STATUS}",
<% if (endpointInitialInstanceCount != null || endpointDataCapturePercent != null || endpointVariantName != null || endpointVolumeSize != null) { %>
  "endpointConfig": {
<% if (endpointInitialInstanceCount != null) { %>
    "initialInstanceCount": ${ENDPOINT_INITIAL_INSTANCE_COUNT}<%= (endpointDataCapturePercent != null || endpointVariantName != null || endpointVolumeSize != null) ? ',' : '' %>
<% } %>
<% if (endpointDataCapturePercent != null) { %>
    "dataCapturePercent": ${ENDPOINT_DATA_CAPTURE_PERCENT}<%= (endpointVariantName != null || endpointVolumeSize != null) ? ',' : '' %>
<% } %>
<% if (endpointVariantName != null) { %>
    "variantName": "${ENDPOINT_VARIANT_NAME}"<%= (endpointVolumeSize != null) ? ',' : '' %>
<% } %>
<% if (endpointVolumeSize != null) { %>
    "volumeSize": ${ENDPOINT_VOLUME_SIZE}
<% } %>
  },
<% } %>
<% if (icCpuCount != null || icMemorySize != null || icGpuCount != null || icCopyCount != null || icModelWeight != null) { %>
  "icConfig": {
<% if (icCpuCount != null) { %>
    "cpuCount": ${IC_CPU_COUNT}<%= (icMemorySize != null || icGpuCount != null || icCopyCount != null || icModelWeight != null) ? ',' : '' %>
<% } %>
<% if (icMemorySize != null) { %>
    "memorySize": ${IC_MEMORY_SIZE}<%= (icGpuCount != null || icCopyCount != null || icModelWeight != null) ? ',' : '' %>
<% } %>
<% if (icGpuCount != null) { %>
    "gpuCount": ${IC_GPU_COUNT}<%= (icCopyCount != null || icModelWeight != null) ? ',' : '' %>
<% } %>
<% if (icCopyCount != null) { %>
    "copyCount": ${IC_COPY_COUNT}<%= (icModelWeight != null) ? ',' : '' %>
<% } %>
<% if (icModelWeight != null) { %>
    "modelWeight": ${IC_MODEL_WEIGHT}
<% } %>
  },
<% } %>
<% if (deploymentTarget === 'realtime-inference') { %>
  "icList": ${IC_LIST_JSON},
<% } %>
  "parameters": ${PARAMETERS}
}
DJEOF
)

    # Output JSON to stdout
    echo "${DEPLOYMENT_JSON}" | python3 -c "import sys,json; print(json.dumps(json.load(sys.stdin), indent=2))" 2>/dev/null || echo "${DEPLOYMENT_JSON}"

    if [ "${CI_MODE}" = true ]; then
        # Strip capacity reservation ARN for CI — force on-demand deployment
        # CI projects must never use reserved capacity (reservations are account-specific
        # and time-bound; CI replay should always target on-demand instances)
        unset CAPACITY_RESERVATION_ARN 2>/dev/null || true

        echo ""
        echo "⚠️  CI Integration is experimental and currently only tested for"
        echo "   SageMaker Real-Time Inference endpoints."
        echo ""

        # Compute configId
        CONFIG_ID=$(compute_config_id)
        echo ""
        echo "🔑 configId: ${CONFIG_ID}"

        # Check if CI_Table exists before writing
        if ! aws dynamodb describe-table --table-name "${CI_TABLE_NAME}" --region "${AWS_REGION}" &>/dev/null; then
            echo ""
            echo "⚠️  CI infrastructure not provisioned. Run 'ml-container-creator bootstrap' with CI enabled."
            echo "   Skipping CI table write."
        else
            echo "📝 Writing to CI table (${CI_TABLE_NAME})..."
            write_ci_record "${CONFIG_ID}"
        fi
    fi

    exit 0
fi

# ============================================================
# Call ml-container-creator registry log
# ============================================================

# Build the registry log command
CMD_ARGS=(
    "registry" "log"
    "--deployment-config" "${DEPLOYMENT_CONFIG}"
    "--architecture" "${ARCHITECTURE}"
    "--backend" "${BACKEND}"
    "--region" "${AWS_REGION}"
    "--status" "${STATUS}"
    "--deployment-target" "${DEPLOYMENT_TARGET}"
    "--build-target" "${BUILD_TARGET}"
)

<% if (framework === 'transformers' || framework === 'diffusors') { %>
if [ -n "${MODEL_NAME:-}" ]; then
    CMD_ARGS+=("--model-name" "${MODEL_NAME}")
fi
<% } %>

<% if (deploymentTarget === 'realtime-inference' || deploymentTarget === 'async-inference' || deploymentTarget === 'batch-transform') { %>
CMD_ARGS+=("--instance-type" "${INSTANCE_TYPE}")
<% } %>

<% if (modelFormat) { %>
if [ -n "${MODEL_FORMAT:-}" ]; then
    CMD_ARGS+=("--model-format" "${MODEL_FORMAT}")
fi
<% } %>

if [ -n "${BASE_IMAGE:-}" ]; then
    CMD_ARGS+=("--base-image" "${BASE_IMAGE}")
fi

if [ -n "${NOTES}" ]; then
    CMD_ARGS+=("--notes" "${NOTES}")
fi

if [ -n "${PROJECT_FLAG}" ]; then
    CMD_ARGS+=("${PROJECT_FLAG}")
fi

# Pass parameters as JSON string
CMD_ARGS+=("--parameters" "${PARAMETERS}")

# Pass IC list as JSON string
<% if (deploymentTarget === 'realtime-inference') { %>
if [ "${IC_LIST_JSON}" != "[]" ]; then
    CMD_ARGS+=("--ic-list" "${IC_LIST_JSON}")
fi
<% } %>

# Pass generator version from package.json if available
GENERATOR_VERSION=""
if command -v node &> /dev/null; then
    GENERATOR_VERSION=$(node -e "try{console.log(require('$(npm root -g)/@aws/ml-container-creator/package.json').version)}catch(e){console.log('')}" 2>/dev/null || echo "")
fi
if [ -n "${GENERATOR_VERSION}" ]; then
    CMD_ARGS+=("--generator-version" "${GENERATOR_VERSION}")
fi

echo "📝 Writing to registry..."
ml-container-creator "${CMD_ARGS[@]}"
