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

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PYTHON_BIN="${PYTHON_BIN:-python3}"
TEXT="你好，今天是美好的一天。"
VOICE="Junhao"
SEED=1234
MAX_NEW_FRAMES=110
GREEDY_PREFIX_FRAMES=4
DECODE_BACKEND="onnx"
DECODE_ONNX_THREADS=0
DECODE_WARMUP_FRAMES=32
DECODE_ONNX="$ROOT_DIR/models/onnxmodels/moss_tts_decode_step.onnx"
OUTPUT_WAV="$ROOT_DIR/outputs/moss_tts_ax650.wav"
OUTPUT_LOG="$ROOT_DIR/outputs/moss_tts_ax650.log"

usage() {
  echo "usage: $0 --text TEXT [--voice NAME] [--decode-backend onnx|axmodel]"
  echo "       [--decode-onnx FILE] [--decode-onnx-threads N] [--decode-warmup-frames N]"
  echo "       [--seed N] [--max-new-frames N] [--output-wav FILE] [--output-log FILE]"
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --text) TEXT="$2"; shift 2 ;;
    --voice) VOICE="$2"; shift 2 ;;
    --seed) SEED="$2"; shift 2 ;;
    --max-new-frames) MAX_NEW_FRAMES="$2"; shift 2 ;;
    --greedy-prefix-frames) GREEDY_PREFIX_FRAMES="$2"; shift 2 ;;
    --decode-backend) DECODE_BACKEND="$2"; shift 2 ;;
    --decode-onnx) DECODE_ONNX="$2"; shift 2 ;;
    --decode-onnx-threads) DECODE_ONNX_THREADS="$2"; shift 2 ;;
    --decode-warmup-frames) DECODE_WARMUP_FRAMES="$2"; shift 2 ;;
    --output-wav) OUTPUT_WAV="$2"; shift 2 ;;
    --output-log) OUTPUT_LOG="$2"; shift 2 ;;
    -h|--help) usage; exit 0 ;;
    *) echo "unknown arg: $1" >&2; usage >&2; exit 2 ;;
  esac
done

mkdir -p "$(dirname "$OUTPUT_WAV")" "$(dirname "$OUTPUT_LOG")"
export LD_LIBRARY_PATH="$ROOT_DIR/bin/moss_tts_ax650:${LD_LIBRARY_PATH:-}"

if [[ "$DECODE_BACKEND" == "onnx" ]]; then
  [[ -f "$DECODE_ONNX" ]] || { echo "missing decode ONNX: $DECODE_ONNX" >&2; exit 1; }
  EXTRA=(--decode-backend onnx --decode-onnx "$DECODE_ONNX"
         --decode-onnx-threads "$DECODE_ONNX_THREADS"
         --decode-warmup-frames "$DECODE_WARMUP_FRAMES")
elif [[ "$DECODE_BACKEND" == "axmodel" ]]; then
  EXTRA=(--decode-backend axmodel --decode-axmodel "$ROOT_DIR/models/axmodels_650/tts_decode_step.axmodel")
else
  echo "--decode-backend must be onnx or axmodel" >&2
  exit 2
fi

REQUEST="$ROOT_DIR/outputs/request.txt"
"$PYTHON_BIN" "$ROOT_DIR/python/prepare_request.py" \
  --config-dir "$ROOT_DIR/config" --text "$TEXT" --voice "$VOICE" \
  --seed "$SEED" --max-new-frames "$MAX_NEW_FRAMES" \
  --greedy-prefix-frames "$GREEDY_PREFIX_FRAMES" --output "$REQUEST"

"$ROOT_DIR/bin/moss_tts_ax650/moss_tts_ax650" \
  --axmodel-dir "$ROOT_DIR/models/axmodels_650" \
  --onnx-dir "$ROOT_DIR/models/onnxmodels" \
  "${EXTRA[@]}" --request "$REQUEST" \
  --output-wav "$OUTPUT_WAV" --max-new-frames "$MAX_NEW_FRAMES" \
  --seed "$SEED" --greedy-prefix-frames "$GREEDY_PREFIX_FRAMES" \
  2>&1 | tee "$OUTPUT_LOG"