#!/usr/bin/env bash
# tts-engine-warm.sh — the boot warmer of the arm A TTS engine, shipped inside
# the frank-tts-engine image at /app/warm-voices.sh
# (dankclaw/voiceservice/Containerfile.engine).
#
# WHO RUNS IT. Two callers, same binary:
#
#   1. The Quadlet unit, as
#      `ExecStartPost=-/usr/bin/podman exec frank-tts-engine /app/warm-voices.sh`
#      — after every start of the container, including every boot.
#      `podman exec` rather than a host script because the engine publishes no
#      port at all, so there is no address the host could dial; and because an
#      exec inherits the unit's EnvironmentFile, which is what puts every knob of
#      this script in scripts/frank-tts-engine.env next to the engine's own.
#   2. scripts/deploy-remote.sh's deploy_frank_tts_engine, once per deploy, so a
#      voice that did not warm becomes a deploy result instead of only a journal
#      line (the unit's `-` prefix deliberately cannot fail the unit).
#
# WHY IT EXISTS AT ALL (the cold-voice law). The FIRST speech request for a
# registered voice decodes the whole reference clip to snapshot its codec state:
# 6.2-8.0 s, first audio byte at 5,530 ms. Every later request for that voice
# restores the snapshot in 0.5 ms. Two facts make that a boot problem rather than
# a curiosity: the voice registry is MEMORY ONLY (a restart forgets every voice),
# and re-registering a voice does not warm it — only synthesizing does. So a boot
# that registers the voices and stops leaves the first caller of the day paying
# 5.5 s of silence. This script speaks one short sentence per voice instead, then
# speaks a second one so the journal carries the warm number too: the line
# `voice <name> cold=…ms warm=…ms` is the repo's standing proof that the second
# request of a voice answers well under a second.
#
# THE VOICES. /voices — the host directory /home/dank/voices mounted read-only by
# the unit, the same packs frank-tts mounts for its own name list. A pack is a
# directory that contains reference.wav; everything else in it (manifest.json, the
# engine's reference.spk / reference.rvq extraction, a 16 kHz variant) is not read
# here. The registration payload is `{"name": <dir>, "wav_b64": <base64 of
# reference.wav>}` with NO ref_text, which selects the embedding-only (x-vector)
# voice — the shipped default, warm RTF ~0.57-0.63 against ICL's 0.67, and the one
# this fleet can produce honestly because it has no trustworthy transcript for the
# audition crops. Adding ref_text (with the codes) would flip that voice to
# in-context-learning mode: better fidelity on a good transcript, slower, and only
# worth it for batch work like a podcast render. That is a per-voice opt-in this
# script deliberately does NOT exercise, and the door stays open: a pack that
# wants ICL needs its text in the pack plus one key in the payload below.
#
# A lone `spk_b64` is a 400 (src/tts-server.h:480 — it is `wav_b64` or BOTH
# `spk_b64` and `rvq_b64`), which is why the WAV travels and the .spk stays home.
#
# NOTHING HERE IS A SECRET, and nothing is logged that could be: the audio bytes
# are discarded, the synthesized text comes from the env file, and the reference
# clip is base64'd straight into the request body without ever being echoed.
set -euo pipefail

# The engine's own address, from inside its namespace. Loopback is valid HERE:
# this script runs in the container, next to the server. There is deliberately no
# host-visible address to use instead — the unit publishes nothing, because a
# host-loopback publish would be unreachable from the façade anyway (netavark
# DNATs published ports only) while adding a surface on an unauthenticated
# server.
: "${ENGINE_URL:=http://127.0.0.1:8880}"

# Container path; the host directory behind it is named in the Quadlet unit.
: "${VOICES_DIR:=/voices}"

# How long to wait for the engine to answer /health. Model load dominates: the
# two GGUFs (1.2 GiB + 291 MB) are read and the context built before the HTTP
# thread binds, and a boot races every other unit on this box.
: "${HEALTH_WAIT:=300}"

# One short sentence, long enough to exercise the codec snapshot and short enough
# that warming a library of voices is seconds, not minutes. Never varied per run:
# the cold/warm numbers in the journal should mean the same thing every boot. The
# synthesis asks for `wav`, the buffered path — the streamed path (`pcm`) is
# different audio at the same seed, and `wav` is also what frank-tts asks for.
: "${WARM_TEXT:=Warm up, one two three.}"

# Per-request ceiling. A cold voice is 6-8 s by measurement; anything approaching
# this is a stalled engine and should be reported rather than waited on forever —
# systemd counts the whole warm pass against the unit's start deadline.
: "${REQUEST_TIMEOUT:=120}"

fail=0
log() { printf 'frank-tts-engine-warm: %s\n' "$*"; }

# Wait for /health. It answers from the HTTP thread, so a 200 means the server is
# listening; it says nothing about voices, which is what the rest of this script
# is for.
ready=0
for ((i = 1; i <= HEALTH_WAIT; i++)); do
    if curl -fsS --max-time 5 "${ENGINE_URL}/health" >/dev/null 2>&1; then
        ready=1
        break
    fi
    sleep 1
done
if (( ready == 0 )); then
    log "engine did not answer ${ENGINE_URL}/health within ${HEALTH_WAIT}s; nothing warmed"
    exit 1
fi

if [[ ! -d "$VOICES_DIR" ]]; then
    log "no voice directory at ${VOICES_DIR}; nothing to warm (frank-tts will answer 400 with an empty voice list until packs are built)"
    exit 0
fi

warmed=0
for dir in "$VOICES_DIR"/*/; do
    [[ -d "$dir" ]] || continue
    name="${dir%/}"
    name="${name##*/}"
    if [[ ! -f "${dir}reference.wav" ]]; then
        continue          # not a pack; frank-tts refuses this name too
    fi
    # The engine lowercases every voice name (tts_voice_name) while frank-tts
    # matches the directory name exactly, so a mixed-case pack would work by
    # accident on one side and not the other. Refuse it loudly instead. Pack
    # names are also JSON string content here, so the character class is what
    # keeps a directory name from writing structure into the request body.
    if [[ ! "$name" =~ ^[a-z0-9][a-z0-9._-]*$ ]]; then
        log "voice '${name}': name must be lowercase [a-z0-9._-] (the engine lowercases names; the shim does not) — skipped"
        fail=1
        continue
    fi

    body="$(printf '{"name":"%s","wav_b64":"%s"}' "$name" "$(base64 -w0 "${dir}reference.wav")")"
    if ! code="$(curl -sS --max-time "$REQUEST_TIMEOUT" -o /dev/null -w '%{http_code}' \
            -H 'Content-Type: application/json' --data-binary "$body" \
            "${ENGINE_URL}/v1/audio/voices")" || [[ "$code" != "200" ]]; then
        log "voice '${name}': registration failed (HTTP ${code:-no response}) — the shim will answer 400 for this name"
        fail=1
        continue
    fi

    # The cold pass. Its output is discarded; the point is the codec snapshot it
    # leaves inside the engine.
    start=$(date +%s%3N)
    if ! curl -fsS --max-time "$REQUEST_TIMEOUT" -o /dev/null \
            -H 'Content-Type: application/json' \
            --data-binary "{\"input\":\"${WARM_TEXT}\",\"voice\":\"${name}\",\"response_format\":\"wav\"}" \
            "${ENGINE_URL}/v1/audio/speech"; then
        log "voice '${name}': registered, but the warm synthesis failed — first live request stays cold"
        fail=1
        continue
    fi
    cold=$(( $(date +%s%3N) - start ))

    # The warm pass, same request. This is the number the fleet actually
    # experiences, and it is what "the second request answers under a second" is
    # measured against.
    start=$(date +%s%3N)
    if ! curl -fsS --max-time "$REQUEST_TIMEOUT" -o /dev/null \
            -H 'Content-Type: application/json' \
            --data-binary "{\"input\":\"${WARM_TEXT}\",\"voice\":\"${name}\",\"response_format\":\"wav\"}" \
            "${ENGINE_URL}/v1/audio/speech"; then
        log "voice '${name}': warm synthesis failed after a cold one that worked"
        fail=1
        continue
    fi
    warm=$(( $(date +%s%3N) - start ))

    log "voice ${name} cold=${cold}ms warm=${warm}ms"
    warmed=$(( warmed + 1 ))
done

# Failure first, emptiness second, and not the other way round: a pass where
# every voice failed also has `warmed == 0`, and reading that as "nothing to do,
# success" is how an endpoint that answers 400 for every voice name gets deployed
# on a Friday.
if (( fail )); then
    log "${warmed} voice(s) warm, at least one failed (see the lines above)"
    exit 1
fi
if (( warmed == 0 )); then
    log "no voice with a reference.wav under ${VOICES_DIR}; engine is up but speaks nothing"
    exit 0
fi
log "${warmed} voice(s) registered and warm"