#!/usr/bin/env bash
set -euo pipefail

REPO_ROOT="${REPO_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}"
IMAGE_TAG="${IMAGE_TAG:-laoyao/nemo-megatron:26.06-flashattn4}"
BASE_IMAGE="${BASE_IMAGE:-ti-coding-agent/nemo-bridge-flashattn4:26.06}"
INSTALL_FLASH_ATTN4="${INSTALL_FLASH_ATTN4:-0}"

export http_proxy="${http_proxy:-http://100.72.0.101:8888}"
export https_proxy="${https_proxy:-http://100.72.0.101:8888}"
export HTTP_PROXY="${HTTP_PROXY:-$http_proxy}"
export HTTPS_PROXY="${HTTPS_PROXY:-$https_proxy}"

docker build \
  --network=host \
  --build-arg BASE_IMAGE="$BASE_IMAGE" \
  --build-arg INSTALL_FLASH_ATTN4="$INSTALL_FLASH_ATTN4" \
  --build-arg PIP_INDEX_URL="${PIP_INDEX_URL:-https://mirrors.aliyun.com/pypi/simple/}" \
  --build-arg PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL:-https://pypi.org/simple}" \
  -f "$REPO_ROOT/docker/nemo/Dockerfile.nemo-megatron" \
  -t "$IMAGE_TAG" \
  "$REPO_ROOT"

docker run --rm --gpus all --ipc=host --network=host "$IMAGE_TAG" bash -lc '
python3 - <<PY
import importlib.metadata as metadata
for package in ["torch", "transformer-engine", "flash-attn", "flash-attn-4", "nvidia-cutlass-dsl", "megatron-core", "megatron-bridge"]:
    try:
        print(package, metadata.version(package))
    except metadata.PackageNotFoundError:
        print(package, "MISSING")
import cutlass
print("cutlass import OK")
PY
'

echo "Built image: $IMAGE_TAG"