#!/bin/bash
set -euo pipefail

PROGPATH="$(realpath "$0")"
PROGDIR="$(dirname "${PROGPATH}")"

KNOWN_ARGS=( samples batch-size data-dir language tts-mode tts-voice-count )
# shellcheck source=/dev/null
source "${PROGDIR}/shell.functions"
WAKE_WORD="${POSITIONAL_ARGS[0]:-}"

if [ ${#UNKNOWN_ARGS[@]} -gt 0 ] ; then
    echo "Unknown argument(s): ${UNKNOWN_ARGS[*]}" >&2
    HELP=true
fi

if [ "${HELP}" == "true" ] || [ -z "${WAKE_WORD}" ] ; then
    cat <<EOF >&2
Usage: $0 [ --samples=<samples> ] [ --batch-size=<batch_size> ]
          [ --language=<lang> ] [ --tts-mode=<modern|hybrid|piper> ]
          [ --tts-voice-count=<voices> ] <wake_word>

--samples:            Number of samples to generate. Default: ${DEFAULT_SAMPLES}
--batch-size:         Generation batch size. Default: ${DEFAULT_BATCH_SIZE}
--language:           TTS language code. Default: ${DEFAULT_LANGUAGE}
--tts-mode:           modern, hybrid, or piper. Default: ${DEFAULT_TTS_MODE}
--tts-voice-count:    Deprecated compatibility option; direct generation ignores it.
<wake_word>           Required phrase to synthesize.
EOF
    exit 1
fi

case "${TTS_MODE}" in
    modern|hybrid|piper) ;;
    *)
        echo "ERROR: --tts-mode must be modern, hybrid, or piper." >&2
        exit 2
        ;;
esac

WORK_DIR="${DATA_DIR}/work"
SAMPLES_DIR="${WORK_DIR}/wake_word_samples"
mkdir -p "${WORK_DIR}"

START_TS=$EPOCHSECONDS
echo "===== Generating ${SAMPLES} wake-word samples (language=${LANGUAGE}, tts=${TTS_MODE}) ====="

python3 "${PROGDIR}/tts_generate_samples.py" "${WAKE_WORD}" \
    --samples="${SAMPLES}" \
    --batch-size="${BATCH_SIZE}" \
    --language="${LANGUAGE}" \
    --tts-mode="${TTS_MODE}" \
    --voice-count="${TTS_VOICE_COUNT}" \
    --data-dir="${DATA_DIR}" \
    --output-dir="${SAMPLES_DIR}"

generated_files=$(find "${SAMPLES_DIR}" -maxdepth 1 -name '*.wav' | wc -l)
if [ "${generated_files}" -ne "${SAMPLES}" ] ; then
    echo "ERROR: only generated ${generated_files} of ${SAMPLES} files" >&2
    exit 1
fi

END_TS=$EPOCHSECONDS
print_elapsed_time "${START_TS}" "${END_TS}" "Generated ${SAMPLES} wake word samples."