# Deploy OmniVoice (https://github.com/k2-fsa/OmniVoice) to irv-ml1, GPU 0
# (RTX 3090). Apache-2.0 zero-shot multilingual voice-cloning TTS, served
# behind our OWN FastAPI wrapper (app.py): batch /v1/audio/speech plus a
# streaming /tts driven by the vendored buffer-ratchet scheduler.
#
# Builds the image locally from stacks/omnivoice/Dockerfile (CUDA 12.8 +
# torch 2.8.0 + omnivoice from PyPI + vendored scheduler.py/sanitize.py),
# stages the build context under /opt/docker/compose/omnivoice/, brings it
# up, and waits for /healthz on :8199.
#
# First run is slow: ~5-10 min docker build + a one-time HF weight pre-warm
# (k2-fsa/OmniVoice) on first container start (entrypoint.sh). The wait loop
# below allows up to ~20 min for build-then-up + pre-warm.
#
# Usage:
#   scripts/elway irv-ml1 --playbook playbooks/deploy-omnivoice.yaml
#
# Idempotent — every step is creates-/when-gated; rerun is safe.

vars:
  compose_dir: /opt/docker/compose/omnivoice
  cache_dir: /worktank/omnivoice/hf_cache
  voices_dir: /worktank/omnivoice/voices
  host_port: "8199"

steps:
  # ── host-side dirs ──────────────────────────────────────────────────

  - name: Ensure /worktank/omnivoice root exists (one-time, sudo)
    shell: mkdir -p /worktank/omnivoice
    sudo: true
    creates: /worktank/omnivoice

  - name: Chown /worktank/omnivoice to lkraven
    shell: chown lkraven:lkraven /worktank/omnivoice
    sudo: true
    when: '[ "$(stat -c %U /worktank/omnivoice)" != lkraven ]'

  - name: Ensure cache dir exists
    shell: mkdir -p {{ cache_dir }}
    creates: "{{ cache_dir }}"

  - name: Ensure voices dir exists
    shell: mkdir -p {{ voices_dir }}
    creates: "{{ voices_dir }}"

  - name: Ensure compose dir exists
    shell: mkdir -p {{ compose_dir }}
    creates: "{{ compose_dir }}"

  # ── deploy build context (compose, dockerfile, entrypoint, env) ───────

  - name: Upload compose.yaml
    upload:
      src: stacks/omnivoice/compose.yaml
      dest: "{{ compose_dir }}/compose.yaml"
      mode: "0644"

  - name: Upload Dockerfile
    upload:
      src: stacks/omnivoice/Dockerfile
      dest: "{{ compose_dir }}/Dockerfile"
      mode: "0644"

  - name: Upload app.py (batch + streaming FastAPI wrapper)
    upload:
      src: stacks/omnivoice/app.py
      dest: "{{ compose_dir }}/app.py"
      mode: "0644"

  - name: Upload scheduler.py (vendored buffer-ratchet streaming scheduler)
    upload:
      src: stacks/omnivoice/scheduler.py
      dest: "{{ compose_dir }}/scheduler.py"
      mode: "0644"

  - name: Upload sanitize.py (language-safe TTS text sanitizer)
    upload:
      src: stacks/omnivoice/sanitize.py
      dest: "{{ compose_dir }}/sanitize.py"
      mode: "0644"

  - name: Stage chatterbox reference voices for cloning (skip _*.wav artifacts)
    shell: |
      set -e
      mkdir -p {{ voices_dir }}
      docker exec chatterbox-fast sh -c 'ls /refs/*.wav' | while read -r f; do
        b=$(basename "$f")
        case "$b" in _*) continue;; esac
        docker cp "chatterbox-fast:$f" "{{ voices_dir }}/$b"
      done
      echo "staged:"; ls {{ voices_dir }}
    # Skip if already staged (Emily.wav is a proxy for "voices present").
    when: "[ ! -f {{ voices_dir }}/Emily.wav ]"

  - name: Upload entrypoint.sh
    upload:
      src: stacks/omnivoice/entrypoint.sh
      dest: "{{ compose_dir }}/entrypoint.sh"
      mode: "0755"

  - name: Seed .env from template (only if absent)
    upload:
      src: stacks/omnivoice/.env.example
      dest: "{{ compose_dir }}/.env"
      mode: "0644"
    when: "[ ! -f {{ compose_dir }}/.env ]"

  # ── build + bring up ────────────────────────────────────────────────

  - name: docker compose build (~5-10 min first time; cached after)
    shell: |
      set -o pipefail
      cd {{ compose_dir }} && docker compose build --progress=plain 2>&1 \
        | grep -vE '^#[0-9]+ [0-9.]+ (Downloading|Collecting|Requirement|Using cached|Installing collected|Successfully (installed|built)|━|Resolved|Prepared|Built)'

  - name: docker compose up -d
    shell: cd {{ compose_dir }} && docker compose up -d

  - name: Wait for /healthz (allow ~25 min for weight + Whisper pre-warm + voice cloning)
    shell: |
      for i in $(seq 1 300); do
        curl -sf -o /dev/null --max-time 3 http://localhost:{{ host_port }}/healthz && exit 0
        sleep 5
      done
      exit 1
    changed_when: "false"

verify:
  - name: /healthz returns 200
    shell: curl -sf -o /dev/null http://localhost:{{ host_port }}/healthz
    changed_when: "false"

  - name: /v1/audio/voices lists the reused chatterbox voices
    shell: curl -sf http://localhost:{{ host_port }}/v1/audio/voices | grep -q '"voices"'
    changed_when: "false"

  - name: Container is running
    shell: docker inspect omnivoice --format '{{.State.Status}}' | grep -q running
    changed_when: "false"