#!/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) ──────────────────
ECR_REPOSITORY_NAME="${ECR_REPOSITORY_NAME:-${_PROFILE[ecrRepositoryName]:-ml-container-creator}}"
export AWS_REGION="${AWS_REGION:-${_PROFILE[awsRegion]:-us-east-1}}"

echo "🚀 Building Docker image for ${PROJECT_NAME}"
echo "   Deployment config: ${DEPLOYMENT_CONFIG}"
echo "   Framework: ${FRAMEWORK}"
echo "   Model server: ${MODEL_SERVER}"

# Validate prerequisites
if ! command -v docker &> /dev/null; then
    echo "❌ Docker is not installed"
    echo "   Install from: https://docs.docker.com/get-docker/"
    exit 2
fi

# Always build for linux/amd64 — SageMaker runs x86_64 instances.
# Without this, Apple Silicon Macs produce arm64 images that silently
# fail on SageMaker with CannotStartContainerError.
PLATFORM_FLAG="--platform linux/amd64"

# --- Secrets Manager resolution (build-time) ---
if [ -n "${HF_TOKEN_ARN:-}" ]; then
    echo "🔐 Resolving HuggingFace token from Secrets Manager..."
    HF_TOKEN=$(aws secretsmanager get-secret-value --secret-id "${HF_TOKEN_ARN}" --query SecretString --output text) || {
        echo "❌ Failed to resolve HuggingFace token from Secrets Manager"
        exit 3
    }
    export HF_TOKEN
fi

if [ -n "${NGC_API_KEY_ARN:-}" ]; then
    echo "🔐 Resolving NGC API key from Secrets Manager..."
    NGC_API_KEY=$(aws secretsmanager get-secret-value --secret-id "${NGC_API_KEY_ARN}" --query SecretString --output text) || {
        echo "❌ Failed to resolve NGC API key from Secrets Manager"
        exit 3
    }
    export NGC_API_KEY
fi

# NOTE: Build-time secrets are passed as --build-arg. The secret value may persist
# in the image layer. A future improvement will use BuildKit --secret mounts.

# Framework-specific build logic
case "${DEPLOYMENT_CONFIG}" in
    transformers-tensorrt-llm)
        echo "🔐 TensorRT-LLM requires NGC authentication"
        if [ -z "${NGC_API_KEY:-}" ]; then
            echo "❌ NGC_API_KEY environment variable not set"
            echo ""
            echo "To build TensorRT-LLM containers, you need an NVIDIA NGC API key:"
            echo "1. Create account at: https://ngc.nvidia.com/"
            echo "2. Generate API key in account settings"
            echo "3. Export key: export NGC_API_KEY=your_key_here"
            echo "4. Run this script again"
            exit 3
        fi
        
        echo "🔑 Authenticating with NVIDIA NGC..."
        echo "${NGC_API_KEY}" | docker login nvcr.io --username '$oauthtoken' --password-stdin
        
        echo "🏗️  Building GPU-enabled image with TensorRT-LLM..."
        docker build ${PLATFORM_FLAG} -t "${PROJECT_NAME}:latest" .
        ;;
        
    transformers-vllm|transformers-sglang)
        echo "🏗️  Building GPU-enabled image..."
        docker build ${PLATFORM_FLAG} -t "${PROJECT_NAME}:latest" .
        ;;
        
    transformers-lmi|transformers-djl)
        echo "🔐 LMI/DJL base images require AWS ECR authentication"
        echo "🔑 Authenticating with AWS Deep Learning Container ECR..."
        aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 763104351884.dkr.ecr.us-east-1.amazonaws.com
        
        echo "🏗️  Building GPU-enabled image..."
        docker build ${PLATFORM_FLAG} -t "${PROJECT_NAME}:latest" .
        ;;
        
    sklearn-*|xgboost-*|tensorflow-*|http-flask|http-fastapi)
        echo "🏗️  Building CPU-optimized image..."
        docker build ${PLATFORM_FLAG} -t "${PROJECT_NAME}:latest" .
        ;;

    triton-*)
        echo "🏗️  Building Triton Inference Server image..."
        docker build ${PLATFORM_FLAG} -t "${PROJECT_NAME}:latest" .
        ;;
        
    *)
        echo "❌ Unknown deployment configuration: ${DEPLOYMENT_CONFIG}"
        exit 3
        ;;
esac

# Tag with timestamp
TIMESTAMP=$(date +%Y%m%d-%H%M%S)
docker tag "${PROJECT_NAME}:latest" "${PROJECT_NAME}:${TIMESTAMP}"

echo "✅ Build complete!"
echo "   Image: ${PROJECT_NAME}:latest"
echo "   Tagged: ${PROJECT_NAME}:${TIMESTAMP}"
echo ""
echo "Next steps:"
echo "  • Test locally: ./do/run"
echo "  • Push to ECR: ./do/push"
echo "  • Deploy to SageMaker: ./do/deploy"
