#!/bin/bash
# inferon setup: MeloTTS venv + checkpoint + whisper.cpp sanity check.
# Idempotent — safe to re-run.
set -euo pipefail

MELO_ROOT="$HOME/Projects/MeloTTS"
WHISPER_ROOT="$HOME/Projects/whisper.cpp"
CKPT_DIR="$MELO_ROOT/ckpt/MeloTTS-English-v3"
HF_BASE="https://huggingface.co/myshell-ai/MeloTTS-English-v3/resolve/main"

say() { printf '\033[1;36m[setup]\033[0m %s\n' "$*"; }

# ---------------------------------------------------------------- MeloTTS ---

if [ ! -d "$MELO_ROOT" ]; then
    say "cloning MeloTTS -> $MELO_ROOT"
    git clone --depth 1 https://github.com/myshell-ai/MeloTTS "$MELO_ROOT"
else
    say "MeloTTS repo present"
fi

if [ ! -x "$MELO_ROOT/.venv/bin/python" ]; then
    say "creating Python 3.11 venv (uv)"
    command -v uv >/dev/null || { say "ERROR: uv not found — install from https://docs.astral.sh/uv/"; exit 1; }
    (cd "$MELO_ROOT" && uv venv --python 3.11 .venv)
fi

if [ ! -e "$MELO_ROOT/.venv/.deps-done" ]; then
    say "installing Python deps (English set; Japanese stubs needed by import chain)"
    grep -v -E "^gradio$|^tensorboard" "$MELO_ROOT/requirements.txt" \
        | uv pip install -r /dev/stdin --python "$MELO_ROOT/.venv/bin/python"
    uv pip install "setuptools<81" mecab-python3==1.0.9 unidic-lite==1.0.8 pykakasi==2.2.1 fugashi==1.3.0 \
        --python "$MELO_ROOT/.venv/bin/python"
    "$MELO_ROOT/.venv/bin/python" -c "import nltk; nltk.download('averaged_perceptron_tagger_eng')"
    touch "$MELO_ROOT/.venv/.deps-done"
else
    say "Python deps already installed"
fi

if [ ! -f "$CKPT_DIR/checkpoint.pth" ]; then
    say "downloading MeloTTS-English-v3 checkpoint (~200MB)"
    mkdir -p "$CKPT_DIR"
    curl -L --fail -o "$CKPT_DIR/checkpoint.pth" "$HF_BASE/checkpoint.pth"
    curl -L --fail -o "$CKPT_DIR/config.json" "$HF_BASE/config.json"
else
    say "checkpoint present"
fi

# --------------------------------------------------------------- whisper ---

if [ ! -x "$WHISPER_ROOT/build/bin/whisper-server" ]; then
    say "WARNING: whisper-server not found at $WHISPER_ROOT/build/bin/whisper-server"
    say "         build whisper.cpp (cmake -B build && cmake --build build) and"
    say "         put the large-v3-turbo q5_0 model in $WHISPER_ROOT/models/"
else
    say "whisper-server OK"
fi

# ----------------------------------------------------------------- misc ----

command -v letta >/dev/null || say "WARNING: letta CLI not on PATH"
command -v pw-record >/dev/null || say "WARNING: pw-record not found (pipewire-tools)"

say "done. build with: zig build"