#!/usr/bin/env bash
# OMNIVOICE-RUNPOD-INSTALL.sh

set -euo pipefail

INSTALL_DIR="${INSTALL_DIR:-/workspace/OmniVoice}"
VENV_DIR="${VENV_DIR:-/tmp/omnivoice_venv}"
PYTHON_BIN="${PYTHON_BIN:-python3}"
TORCH_INDEX="${TORCH_INDEX:-https://download.pytorch.org/whl/cu128}"

log(){ echo "[OMNIVOICE] $*"; }
die(){ echo "[OMNIVOICE][ERROR] $*" >&2; exit 1; }

log "Checking NVIDIA GPU"
command -v nvidia-smi >/dev/null 2>&1 || die "nvidia-smi not found. Use a RunPod NVIDIA GPU pod."
nvidia-smi

log "Installing system packages"
apt-get update
DEBIAN_FRONTEND=noninteractive apt-get install -y \
  git curl wget ffmpeg \
  python3 python3-venv python3-pip \
  build-essential

log "Preparing folders"
mkdir -p "$INSTALL_DIR/hf_cache" "$INSTALL_DIR/torch_cache"

log "Removing old venv"
rm -rf "$VENV_DIR"

log "Creating venv in $VENV_DIR"
"$PYTHON_BIN" -m venv "$VENV_DIR"
source "$VENV_DIR/bin/activate"

log "Installing pip + uv"
python -m pip install -U pip wheel setuptools uv

export UV_LINK_MODE=copy
export UV_CACHE_DIR="/tmp/uv_cache"
export PYTHONNOUSERSITE=1
unset PYTHONPATH

export HF_HOME="$INSTALL_DIR/hf_cache"
export HF_HUB_DISABLE_SYMLINKS_WARNING=1
export TORCH_HOME="$INSTALL_DIR/torch_cache"
export PIP_DISABLE_PIP_VERSION_CHECK=1
export PYTHONUTF8=1
export PYTHONIOENCODING=utf-8

log "Installing PyTorch + torchaudio from cu128"
uv pip install torch==2.8.0 torchaudio==2.8.0 --index-url "$TORCH_INDEX"

log "Installing OmniVoice"
uv pip install --upgrade omnivoice

log "Re-applying PyTorch + torchaudio from cu128"
uv pip install torch==2.8.0 torchaudio==2.8.0 --index-url "$TORCH_INDEX"

log "Testing CUDA + OmniVoice"
python - <<'PY'
import torch
import torchaudio
import omnivoice

print("[OK] torch:", torch.__version__)
print("[OK] torchaudio:", torchaudio.__version__)
print("[OK] cuda available:", torch.cuda.is_available())
print("[OK] torch cuda:", torch.version.cuda)

if not torch.cuda.is_available():
    raise SystemExit("[FAIL] CUDA is not available in PyTorch")

print("[OK] gpu:", torch.cuda.get_device_name(0))
x = torch.randn((1,), device="cuda")
print("[OK] CUDA test passed")
print("[OK] OmniVoice import passed")
PY

PORT="$(python - <<'PY'
import socket
s = socket.socket()
s.bind(("127.0.0.1", 0))
print(s.getsockname()[1])
s.close()
PY
)"

log "Starting OmniVoice UI with Gradio public share link"
echo
echo "============================================================"
echo "No RunPod HTTP port needed."
echo "Wait for Gradio to print:"
echo
echo "  Running on public URL: https://....gradio.live"
echo
echo "Then Click that public URL. DONT CLICK THE LINK ABOVE IT'S AN EXAMPLE YOU DUMMY I SEE YOU"
echo "============================================================"
echo

omnivoice-demo --ip 0.0.0.0 --port "$PORT" --share