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

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
VENV_DIR="$ROOT_DIR/.venv"
VENV_PYTHON="$VENV_DIR/bin/python"
VENV_PIP="$VENV_DIR/bin/pip"
VIBEVOICE_DIR="${VIBEVOICE_DIR:-/home/azurian/VibeVoice}"

cd "$ROOT_DIR"

if [[ -f ".env" ]]; then
  set -a
  # shellcheck disable=SC1091
  source ".env"
  set +a
fi

if [[ ! -x "$VENV_PYTHON" ]]; then
  echo "Creating virtual environment in $VENV_DIR"
  python3 -m venv "$VENV_DIR"
fi

create_clean_venv() {
  echo "Recreating clean virtual environment in $VENV_DIR"
  rm -rf "$VENV_DIR"
  python3 -m venv "$VENV_DIR"
}

install_project() {
  if ! "$VENV_PIP" install -e "$ROOT_DIR"; then
    echo "Initial install failed, retrying from a clean virtual environment"
    create_clean_venv
    "$VENV_PIP" install -e "$ROOT_DIR"
  fi
}

ensure_vibevoice_runtime() {
  local provider="${TTS_PROVIDER:-vibevoice-1.5b}"
  case "${provider,,}" in
    vibevoice|vibevoice-1.5b|vibevoice15b)
      if ! "$VENV_PYTHON" - <<'PY'
import importlib.util
import sys
sys.exit(0 if importlib.util.find_spec("vibevoice") else 1)
PY
      then
        if [[ -d "$VIBEVOICE_DIR" ]]; then
          echo "Installing local VibeVoice runtime from $VIBEVOICE_DIR"
          "$VENV_PIP" install -e "$VIBEVOICE_DIR"
        else
          echo "VibeVoice runtime is required but $VIBEVOICE_DIR was not found" >&2
          exit 1
        fi
      fi
      ;;
  esac
}

ensure_requirements() {
  local missing=0
  if ! "$VENV_PYTHON" - <<'PY'
import importlib.util
import sys

required_modules = [
    "fastapi",
    "uvicorn",
    "numpy",
    "soundfile",
    "tts",
]

missing = [name for name in required_modules if importlib.util.find_spec(name) is None]
if missing:
    print("Missing modules:", ", ".join(missing))
    sys.exit(1)
PY
  then
    missing=1
  fi

  if [[ "$missing" -eq 1 ]]; then
    echo "Installing project requirements into .venv"
    install_project
  fi
}

ensure_requirements
ensure_vibevoice_runtime

exec "$VENV_PYTHON" -m tts.fastapi_app "$@"