#!/bin/bash
# Apply Sortformer diarization patches to a running parakeet-asr container.
#
# Run from the spark-control repo root on the laptop:
#   bash image/parakeet_patches/apply.sh <spark2-host> <ssh-user>
#
# What it does:
#   1. Backs up the current /opt/parakeet/app/main.py inside the container
#      (writable layer; survives docker restart but NOT docker rm).
#   2. Copies the patched main.py + new diarizer.py into the container.
#   3. Restarts the container so the new code + Sortformer model load.
#
# Reversibility:
#   - The backup of main.py is at /opt/parakeet/app/main.py.pre-sortformer
#     inside the container. Restore with:
#       docker exec parakeet-asr cp /opt/parakeet/app/main.py.pre-sortformer /opt/parakeet/app/main.py
#       docker exec parakeet-asr rm -f /opt/parakeet/app/diarizer.py
#       docker restart parakeet-asr
#   - If the container is ever `docker rm`'d (volume rebuild), re-run this
#     script. We will eventually fold this into spark-control as an action.

set -e

HOST="${1:?usage: apply.sh <spark2-host> <ssh-user>}"
USER="${2:?usage: apply.sh <spark2-host> <ssh-user>}"
CONTAINER="${CONTAINER:-parakeet-asr}"

REPO_DIR="$(cd "$(dirname "$0")" && pwd)"

echo "→ Backing up current main.py inside ${CONTAINER}..."
ssh "${USER}@${HOST}" "docker exec ${CONTAINER} sh -c \
  'test -f /opt/parakeet/app/main.py.pre-sortformer || cp /opt/parakeet/app/main.py /opt/parakeet/app/main.py.pre-sortformer'"

echo "→ Copying diarizer.py into container..."
ssh "${USER}@${HOST}" "docker exec -i ${CONTAINER} sh -c \
  'cat > /opt/parakeet/app/diarizer.py'" < "${REPO_DIR}/diarizer.py"

echo "→ Copying patched main.py into container..."
ssh "${USER}@${HOST}" "docker exec -i ${CONTAINER} sh -c \
  'cat > /opt/parakeet/app/main.py'" < "${REPO_DIR}/main.py"

echo "→ Verifying syntax inside container..."
ssh "${USER}@${HOST}" "docker exec ${CONTAINER} python3 -c \
  'import ast; ast.parse(open(\"/opt/parakeet/app/diarizer.py\").read()); ast.parse(open(\"/opt/parakeet/app/main.py\").read()); print(\"py OK\")'"

echo "→ Restarting ${CONTAINER}..."
ssh "${USER}@${HOST}" "docker restart ${CONTAINER}"

echo
echo "✔ Patches applied. Sortformer model (~150 MB) will download on first load — wait ~30s before testing."
echo
echo "Test once it's ready:"
echo "  curl -sS http://${HOST}:8000/health"
echo "  curl -sS -X POST http://${HOST}:8000/v1/audio/diarize -F file=@some-audio.mp3 | head -c 500"