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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REGISTRY="${WM_IMAGE_REGISTRY:-algo-image-cn-shanghai.cr.volces.com}"
NAMESPACE="${WM_IMAGE_NAMESPACE:-niuma}"
BASE_IMAGE_NAME="${WM_BASE_IMAGE_NAME:-wm-daemon-runtime-base}"
RUNTIME_IMAGE_PREFIX="${WM_RUNTIME_IMAGE_PREFIX:-wm-daemon-runtime}"
PLATFORMS="${WM_IMAGE_PLATFORMS:-}"
PUSH="${WM_IMAGE_PUSH:-1}"
BUILD_BASE="${WM_BUILD_BASE:-1}"
RUNTIMES="${WM_RUNTIMES-claude codex kimi hermes opencode cursor antigravity}"
FULL_PREFIX="$REGISTRY/$NAMESPACE"
BASE_IMAGE="${WM_BASE_IMAGE_OVERRIDE:-$FULL_PREFIX/$BASE_IMAGE_NAME}"
BASE_OUTPUT_IMAGE="$FULL_PREFIX/$BASE_IMAGE_NAME"

if [[ -n "${WM_IMAGE_REGISTRY_USERNAME:-}" || -n "${WM_IMAGE_REGISTRY_PASSWORD:-}" ]]; then
  if [[ -z "${WM_IMAGE_REGISTRY_USERNAME:-}" || -z "${WM_IMAGE_REGISTRY_PASSWORD:-}" ]]; then
    echo "WM_IMAGE_REGISTRY_USERNAME and WM_IMAGE_REGISTRY_PASSWORD must be set together" >&2
    exit 2
  fi
  printf '%s' "$WM_IMAGE_REGISTRY_PASSWORD" | docker login "$REGISTRY" --username "$WM_IMAGE_REGISTRY_USERNAME" --password-stdin
fi

build_image() {
  local dockerfile="$1"
  local image="$2"
  shift 2
  local args=("$@")
  if [[ -n "$PLATFORMS" ]]; then
    local buildx_cmd=(docker buildx build --pull --platform "$PLATFORMS" -f "$dockerfile" -t "$image:latest")
    if [[ ${#args[@]} -gt 0 ]]; then
      buildx_cmd+=("${args[@]}")
    fi
    buildx_cmd+=(--push "$SCRIPT_DIR")
    "${buildx_cmd[@]}"
  else
    local build_cmd=(docker build --pull -f "$dockerfile" -t "$image:latest")
    if [[ ${#args[@]} -gt 0 ]]; then
      build_cmd+=("${args[@]}")
    fi
    build_cmd+=("$SCRIPT_DIR")
    "${build_cmd[@]}"
    if [[ "$PUSH" != "0" ]]; then
      docker push "$image:latest"
    fi
  fi
}

if [[ "$BUILD_BASE" != "0" ]]; then
  echo "Building base image: $BASE_OUTPUT_IMAGE:latest"
  base_build_args=()
  if [[ -n "${WM_APT_MIRROR:-}" ]]; then
    base_build_args+=(--build-arg "APT_MIRROR=$WM_APT_MIRROR")
  fi
  if [[ -n "${WM_APT_SECURITY_MIRROR:-}" ]]; then
    base_build_args+=(--build-arg "APT_SECURITY_MIRROR=$WM_APT_SECURITY_MIRROR")
  fi
  if [[ -n "${WM_NPM_REGISTRY:-}" ]]; then
    base_build_args+=(--build-arg "NPM_REGISTRY=$WM_NPM_REGISTRY")
  fi
  build_image "$SCRIPT_DIR/Dockerfile.base" "$BASE_OUTPUT_IMAGE" "${base_build_args[@]}"
else
  echo "Skipping base image build: $BASE_IMAGE:latest"
fi

for runtime in $RUNTIMES; do
  image="$FULL_PREFIX/$RUNTIME_IMAGE_PREFIX-$runtime"
  echo "Building runtime image: $image:latest"
  runtime_build_args=(
    --build-arg "BASE_IMAGE=$BASE_IMAGE:latest"
    --build-arg "WM_RUNTIME=$runtime"
  )
  if [[ -n "${WM_PIP_INDEX_URL:-}" ]]; then
    runtime_build_args+=(--build-arg "PIP_INDEX_URL=$WM_PIP_INDEX_URL")
  fi
  build_image "$SCRIPT_DIR/Dockerfile.runtime" "$image" \
    "${runtime_build_args[@]}"
done

echo "$BASE_OUTPUT_IMAGE:latest"
for runtime in $RUNTIMES; do
  echo "$FULL_PREFIX/$RUNTIME_IMAGE_PREFIX-$runtime:latest"
done