# Deploy Fish Audio S2-Pro (richest paralinguistic open-source TTS) to
# irv-ml1.
#
# Builds the image locally from fishaudio/fish-speech via docker buildx
# git URL context. ~10-15 min cold build (CUDA 12.x + torch + flash-attn
# + Fish's training/inference deps). First start downloads s2-pro
# (~9 GB BF16) into the bind-mounted HF cache. Generous /v1/health
# wait deadline accommodates both.
#
# Usage:
#   scripts/elway irv-ml1 --playbook playbooks/deploy-fish-s2.yaml
#
# Idempotent — every step is creates-/when-gated; rerun is safe.

vars:
  compose_dir: /opt/docker/compose/fish-s2
  references_dir: /worktank/fish-s2/references
  checkpoints_dir: /worktank/fish-s2/checkpoints
  cache_dir: /worktank/fish-s2/hf_cache
  host_port: "8195"

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

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

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

  - name: Ensure references dir exists
    shell: mkdir -p {{ references_dir }}
    creates: "{{ references_dir }}"

  - name: Normalize reference layout (flat <name>.wav → <name>/<name>.wav + .lab)
    # fish-speech resolves reference_id=<name> against the DIRECTORY
    # references/<name>/ (audio + a same-basename .lab transcript), NOT a
    # flat references/<name>.wav. A flat-only layout makes every
    # reference_id silently resolve to nothing → Fish falls back to its
    # default speaker, so every voice in the dropdown sounds identical
    # (this was the 2026-06-01 "no accent" root cause). External voice
    # imports drop flat <name>.wav + <name>.txt here; this step mirrors
    # each into references/<name>/<name>.wav + <name>/<name>.lab. `cp -u`
    # makes it idempotent; the `when:` skips it once every flat wav has a
    # matching nested wav (audio-only voices like glados need no .lab).
    shell: |
      cd {{ references_dir }} || exit 1
      for w in *.wav; do
        [ -e "$w" ] || continue
        n="${w%.wav}"; mkdir -p "$n"
        cp -u "$w" "$n/$n.wav"
        [ -e "$n.txt" ] && cp -u "$n.txt" "$n/$n.lab" || true
      done
    when: "[ \"$(cd {{ references_dir }} 2>/dev/null && ls *.wav 2>/dev/null | wc -l)\" -ne \"$(cd {{ references_dir }} 2>/dev/null && ls */*.wav 2>/dev/null | wc -l)\" ]"

  - name: Ensure checkpoints dir exists
    shell: mkdir -p {{ checkpoints_dir }}
    creates: "{{ checkpoints_dir }}"

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

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

  # ── deploy compose + env ────────────────────────────────────────────

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

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

  # ── pre-pull model checkpoint ───────────────────────────────────────
  # Fish doesn't auto-download on first run — start_server.sh validates
  # checkpoints/s2-pro/ exists and exits cleanly (rc=0) if missing. So
  # we pre-pull fishaudio/s2-pro into the bind-mount via a one-shot
  # python container with hf_transfer (~9 GB at ~100 MB/s).
  # Idempotent — `creates:` skips if the codec file is already there.

  - name: Pre-pull fishaudio/s2-pro into checkpoints (~9 GB, ~90s)
    shell: |
      docker run --rm --user 1000:1000 \
        -e HOME=/tmp/h -e HF_HUB_ENABLE_HF_TRANSFER=1 \
        -v {{ checkpoints_dir }}:/dest \
        python:3.12-slim sh -c 'set -e; mkdir -p /tmp/h /tmp/pip /tmp/site; PIP_CACHE_DIR=/tmp/pip pip install --quiet --target /tmp/site huggingface_hub hf_transfer; PYTHONPATH=/tmp/site python -c "from huggingface_hub import snapshot_download; snapshot_download(repo_id=\"fishaudio/s2-pro\", local_dir=\"/dest/s2-pro\", local_dir_use_symlinks=False, ignore_patterns=[\"*.md\",\"*.png\",\".gitattributes\"])"'
    creates: "{{ checkpoints_dir }}/s2-pro/codec.pth"

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

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

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

  - name: Wait for /v1/health to respond (allow ~15 min for first model download + warmup)
    shell: |
      for i in $(seq 1 180); do
        curl -sf -o /dev/null --max-time 3 http://localhost:{{ host_port }}/v1/health && exit 0
        sleep 5
      done
      exit 1
    changed_when: "false"

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

  - name: /v1/tts returns a real WAV (POST with text body)
    # Fish's API is NOT OpenAI-compatible — there's no /v1/audio/speech
    # and no /v1/audio/voices. The single TTS endpoint is POST /v1/tts
    # with at minimum {"text":"..."} returning audio/wav. Named-voice
    # cloning is via the reference_id field, which resolves against the
    # references/<name>/ directory (audio + .lab) — see the A/B gate
    # below. Verify by POST + asserting the response is a real RIFF WAV.
    shell: |
      out=$(mktemp --suffix=.wav)
      curl -sf -X POST http://localhost:{{ host_port }}/v1/tts \
        -H 'Content-Type: application/json' \
        -d '{"text":"Verify."}' \
        -o "$out" --max-time 30
      file -b "$out" | grep -q '^RIFF.*WAVE'
      rm -f "$out"
    changed_when: "false"

  - name: reference_id actually changes the voice (A/B — guards the empty-dir regression)
    # Two different reference_ids with identical text+seed MUST yield
    # different audio. If byte-identical, reference_id is being ignored
    # (empty references/<id>/ dirs) and every voice has silently
    # collapsed to Fish's default speaker — the 2026-06-01 root cause.
    # Hard-fail the deploy so the regression can never ship silently.
    shell: |
      a=$(mktemp); b=$(mktemp)
      curl -sf -X POST http://localhost:{{ host_port }}/v1/tts -H 'Content-Type: application/json' \
        -d '{"text":"The quick brown fox.","seed":42,"reference_id":"Abigail"}' -o "$a" --max-time 60
      curl -sf -X POST http://localhost:{{ host_port }}/v1/tts -H 'Content-Type: application/json' \
        -d '{"text":"The quick brown fox.","seed":42,"reference_id":"Imogen"}' -o "$b" --max-time 60
      rc=0; [ "$(md5sum < "$a")" = "$(md5sum < "$b")" ] && rc=1
      rm -f "$a" "$b"
      exit $rc
    changed_when: "false"

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