#!/bin/bash
set -euo pipefail

IMAGE="cortazar-spark"
CONTAINER_NAME="cortex_infer_spark"
WORKSPACE_DIR="$(pwd)"

if ! docker image inspect "$IMAGE" >/dev/null 2>&1; then
  echo "Image $IMAGE not found. Building from Dockerfile.spark..."
  docker build -f Dockerfile.spark -t "$IMAGE" .
fi

if [ "$(docker ps -a -q -f name=$CONTAINER_NAME)" ]; then
  echo "Container exists. Restarting..."
  docker rm -f "$CONTAINER_NAME"
fi

echo "--- Launching inference container ---"
docker run -d --rm \
  --runtime=nvidia \
  --gpus all \
  --ipc=host \
  --network=host \
  --ulimit memlock=-1 \
  --ulimit stack=67108864 \
  --name "$CONTAINER_NAME" \
  -v "$WORKSPACE_DIR:/workspace" \
  "$IMAGE" \
  sleep infinity

echo "--- Installing inference dependencies (no torch/torchaudio upgrades) ---"
docker exec "$CONTAINER_NAME" bash -c "
  set -e
  BASE_TORCH=\$(python -c 'import torch; print(torch.__version__)')
  BASE_TORCHAUDIO=\$(python -c 'import torchaudio; print(torchaudio.__version__)')
  echo \"Base torch: \$BASE_TORCH\"
  echo \"Base torchaudio: \$BASE_TORCHAUDIO\"

  pip install --no-deps -e F5-TTS/

  pip install pytorch-metric-learning==2.9.0
  pip install google-api-core google-cloud-storage boto3 \
    x_transformers torchdiffeq transformers_stream_generator unidecode pypinyin rjieba
  pip install --no-deps encodec

  pip install --no-cache-dir -r /dev/stdin --no-deps < <(cat scripts/core/requirements.txt | grep -vE '^\s*(torch|torchaudio)([<>=![:space:]]|$)')

  CUR_TORCH=\$(python -c 'import torch; print(torch.__version__)')
  CUR_TORCHAUDIO=\$(python -c 'import torchaudio; print(torchaudio.__version__)')
  if [ \"\$CUR_TORCH\" != \"\$BASE_TORCH\" ] || [ \"\$CUR_TORCHAUDIO\" != \"\$BASE_TORCHAUDIO\" ]; then
    echo \"ERROR: torch/torchaudio version changed (base: \$BASE_TORCH/\$BASE_TORCHAUDIO, current: \$CUR_TORCH/\$CUR_TORCHAUDIO). Aborting.\"
    exit 1
  fi
"

echo "Inference container ready: $CONTAINER_NAME"