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

APP_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
VENV_DIR="${AUDEX_VENV_DIR:-$APP_DIR/.venv}"
PYTHON="$VENV_DIR/bin/python"

if [[ ! -x "$PYTHON" ]]; then
    echo "Missing local environment. Run: bash $APP_DIR/setup_local.sh" >&2
    exit 1
fi

export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}"

gpu_info="$("$PYTHON" - <<'PY'
import re
import sys

import torch

if not torch.cuda.is_available():
    raise SystemExit("CUDA is unavailable in this PyTorch environment.")

major, minor = torch.cuda.get_device_capability()
cuda = (torch.version.cuda or "unknown").replace(".", "")
torch_version = re.sub(r"[^0-9A-Za-z]+", "", torch.__version__.split("+", 1)[0])
python_version = f"{sys.version_info.major}{sys.version_info.minor}"
memory_gib = torch.cuda.get_device_properties(0).total_memory / 1024**3
print(f"{major}.{minor}\tsm{major}{minor}-cu{cuda}-torch{torch_version}-py{python_version}\t{memory_gib:.0f}")
PY
)"
IFS=$'\t' read -r detected_arch runtime_tag memory_gib <<< "$gpu_info"

export AUDEX_TORCH_ARCH="${AUDEX_TORCH_ARCH:-$detected_arch}"
export AUDEX_WHEELS_DIR="${AUDEX_WHEELS_DIR:-$APP_DIR/.local/wheels/$runtime_tag}"
export AUDEX_BUILD_DIR="${AUDEX_BUILD_DIR:-$APP_DIR/.local/build/$runtime_tag}"
export AUDEX_MAX_AUDIO_SECONDS="${AUDEX_MAX_AUDIO_SECONDS:-900}"
export AUDEX_MAX_NEW_TOKENS="${AUDEX_MAX_NEW_TOKENS:-4096}"
export AUDEX_DEFAULT_MAX_NEW_TOKENS="${AUDEX_DEFAULT_MAX_NEW_TOKENS:-1024}"
export AUDEX_TTS_MAX_NEW_TOKENS="${AUDEX_TTS_MAX_NEW_TOKENS:-512}"
export AUDEX_TTS_DEFAULT_NEW_TOKENS="${AUDEX_TTS_DEFAULT_NEW_TOKENS:-256}"
export AUDEX_S2S_TEXT_MAX_NEW_TOKENS="${AUDEX_S2S_TEXT_MAX_NEW_TOKENS:-4096}"
export AUDEX_S2S_TTS_MAX_NEW_TOKENS="${AUDEX_S2S_TTS_MAX_NEW_TOKENS:-2400}"
export GRADIO_SERVER_NAME="${GRADIO_SERVER_NAME:-0.0.0.0}"
export GRADIO_SERVER_PORT="${GRADIO_SERVER_PORT:-7860}"
export GRADIO_SHARE="${GRADIO_SHARE:-false}"

echo "GPU architecture: $AUDEX_TORCH_ARCH (${memory_gib} GiB)"
echo "Models: Audex-30B-A3B and Audex-2B (downloaded from the Hub if uncached)"

exec "$PYTHON" "$APP_DIR/app.py"