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

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"

VENDOR=""
BASE_IMAGE=""
REGISTRY="unis"
IMAGE_NAME="qwen3-asr"
VERSION="$(date +"%Y%m%d_%H%M%S")"
OUTPUT_DIR="${ROOT_DIR}/build-file"
EXPORT_TAR="true"
NO_CACHE="false"
PYTHON_BIN=""

usage() {
  cat <<EOF
Usage:
  ./scripts/package_vendor_gpu_image.sh --vendor metax|iluvatar|mthreads --base-image IMAGE [options]

Options:
  --vendor NAME        Vendor accelerator: metax, iluvatar, or mthreads
  --base-image IMAGE   Official vendor vLLM image already loaded/pullable by Docker
  -v, --version VER    Output image version tag; default: timestamp
  -r, --registry REG   Image registry/namespace; default: ${REGISTRY}
  -o, --output DIR     Export directory; default: ${OUTPUT_DIR}
  --python-bin BIN     Python inside base image; default: /opt/conda/bin/python for metax, python3 for iluvatar
  --no-export          Build image only, do not docker save tar.gz
  --no-cache           Disable Docker build cache
  -h, --help           Show help

Example:
  ./scripts/package_vendor_gpu_image.sh \\
    --vendor metax \\
    --base-image metax/vllm:3.7.0.38-c500-py310 \\
    -v n260-3.7.0.38

  ./scripts/package_vendor_gpu_image.sh \\
    --vendor iluvatar \\
    --base-image registry.iluvatar.com.cn:10443/customer/sz/vllm0.17.0-4.4.0-x86:v5 \\
    -v vllm0.17.0-4.4.0-v5

  ./scripts/package_vendor_gpu_image.sh \\
    --vendor mthreads \\
    --base-image registry.mthreads.com/presale/devtech/vllm_musa:s4000_4.3.5_d0519 \\
    -v s4000_4.3.5_d0519
EOF
}

die() { echo "[ERROR] $1" >&2; exit 1; }
info() { echo "[INFO] $1" >&2; }
compressor() { command -v pigz >/dev/null 2>&1 && echo "pigz -f" || echo "gzip -f"; }

while [[ $# -gt 0 ]]; do
  case "$1" in
    --vendor) VENDOR="$2"; shift 2 ;;
    --base-image) BASE_IMAGE="$2"; shift 2 ;;
    -v|--version) VERSION="$2"; shift 2 ;;
    -r|--registry) REGISTRY="$2"; shift 2 ;;
    -o|--output) OUTPUT_DIR="$2"; shift 2 ;;
    --python-bin) PYTHON_BIN="$2"; shift 2 ;;
    --no-export) EXPORT_TAR="false"; shift ;;
    --no-cache) NO_CACHE="true"; shift ;;
    -h|--help) usage; exit 0 ;;
    *) die "Unknown argument: $1" ;;
  esac
done

[[ -n "$VENDOR" ]] || die "--vendor is required"
[[ -n "$BASE_IMAGE" ]] || die "--base-image is required"

case "$VENDOR" in
  metax)
    dockerfile="Dockerfile.metax"
    tag="${REGISTRY}/${IMAGE_NAME}:metax-${VERSION}"
    target="metax"
    build_arg_name="METAX_BASE_IMAGE"
    default_python_bin="/opt/conda/bin/python"
    ;;
  iluvatar|ix|tianshu)
    dockerfile="Dockerfile.iluvatar"
    tag="${REGISTRY}/${IMAGE_NAME}:iluvatar-${VERSION}"
    target="iluvatar"
    build_arg_name="ILUVATAR_BASE_IMAGE"
    default_python_bin="python3"
    ;;
  mthreads|musa|moorethreads)
    dockerfile="Dockerfile.mthreads"
    tag="${REGISTRY}/${IMAGE_NAME}:mthreads-${VERSION}"
    target="mthreads"
    build_arg_name="MTHREADS_BASE_IMAGE"
    default_python_bin="python3"
    ;;
  *)
    die "Unsupported vendor: ${VENDOR}"
    ;;
esac

PYTHON_BIN="${PYTHON_BIN:-$default_python_bin}"

docker image inspect "$BASE_IMAGE" >/dev/null 2>&1 || {
  die "Base image not found locally: ${BASE_IMAGE}. Run docker load first or docker pull it."
}

args=(
  build
  -f "$dockerfile"
  -t "$tag"
  --build-arg "${build_arg_name}=${BASE_IMAGE}"
  --build-arg "PYTHON_BIN=${PYTHON_BIN}"
)
[[ "$NO_CACHE" == "true" ]] && args+=(--no-cache)

info "Building fused image: ${tag}"
info "Vendor base image: ${BASE_IMAGE}"
info "Base image Python: ${PYTHON_BIN}"
(
  cd "$ROOT_DIR"
  docker "${args[@]}" .
)

if [[ "$EXPORT_TAR" == "true" ]]; then
  mkdir -p "$OUTPUT_DIR"
  tar_path="${OUTPUT_DIR}/${IMAGE_NAME}-${target}-${VERSION}-amd64.tar"
  info "Exporting image archive: ${tar_path}"
  docker save -o "$tar_path" "$tag"
  info "Compressing archive"
  $(compressor) "$tar_path"
  info "Done: ${tar_path}.gz"
else
  info "Built image: ${tag}"
fi