# -*- coding: utf-8 -*-
"""bench_b2.py — verifiziert B2 (Worker-Pool) end-to-end mit den ECHTEN pocket_server-Funktionen.
Misst die Wall-Clock einer langen, mehrsätzigen Antwort SERIELL (1 Modell) vs. 3 WORKER parallel,
prüft die Reihenfolge (geordnete Ausgabe) und legt beide Ergebnis-WAVs auf den Desktop.
"""
import os, time
import numpy as np, soundfile as sf, librosa
from concurrent.futures import ProcessPoolExecutor
from pocket_tts import TTSModel
import pocket_server as ps

DESK = os.path.join(os.path.expanduser("~"), "Desktop", "lucy_samples")
os.makedirs(DESK, exist_ok=True)

LONG = ("Guten Morgen, Commander. Das nächtliche Backup ist sauber durchgelaufen. "
        "Es gab keine Fehler in den Protokollen. Der Dienst läuft stabil weiter. "
        "Ich habe die Modelle vorgewärmt und die Engine antwortet zügig. "
        "Wenn du möchtest, starte ich jetzt den Tagesbericht und fasse die offenen Punkte zusammen.")


def _serial():
    m = TTSModel.load_model(language=ps.LANG, lsd_decode_steps=ps.LSD, temp=ps.TEMP,
                            noise_clamp=ps.NOISE_CLAMP, quantize=ps.QUANTIZE)
    ref = ps._prep_ref()
    try:
        vs = m.get_state_for_audio_prompt(ps.VOICE_ST)
    except Exception:
        vs = m.get_state_for_audio_prompt(ref)
    ref_audio, _ = librosa.load(ref, sr=m.sample_rate, mono=True)
    ps.STATE.clear()
    ps.STATE.update(m=m, vs=vs, sr=m.sample_rate, ref_fp=ps._fingerprint(ref_audio, m.sample_rate))
    sents = ps._split_sentences(LONG)
    t0 = time.time()
    parts = list(ps._gen_sentences_ordered(sents))   # pool=None -> serieller Zweig (LOCK)
    dt = time.time() - t0
    a = np.concatenate(parts)
    sf.write(os.path.join(DESK, "b2_serial.wav"), a, m.sample_rate)
    ps.STATE.clear(); del m
    return dt, a.size / 24000, len(sents)


def _parallel(nworkers):
    ps.STATE.clear()
    pool = ProcessPoolExecutor(max_workers=nworkers, initializer=ps._worker_init)
    list(pool.map(ps._warmup, range(nworkers)))      # alle Worker vorab hochfahren
    ref = ps._prep_ref()
    ref_audio, _ = librosa.load(ref, sr=24000, mono=True)
    ps.STATE.update(sr=24000, ref_fp=ps._fingerprint(ref_audio, 24000), pool=pool)
    sents = ps._split_sentences(LONG)
    t0 = time.time()
    parts = list(ps._gen_sentences_ordered(sents))   # pool -> parallel, geordnet
    dt = time.time() - t0
    a = np.concatenate(parts)
    sf.write(os.path.join(DESK, f"b2_parallel_{nworkers}w.wav"), a, 24000)
    pool.shutdown(wait=True); ps.STATE.clear()
    return dt, a.size / 24000, len(sents)


if __name__ == "__main__":
    print(f"Text: {len(LONG)} Zeichen")
    ds, audio_s, n = _serial()
    print(f"SERIELL : {ds:5.2f}s wall  ({n} Saetze, {audio_s:4.1f}s Audio)  -> b2_serial.wav")
    dp, _, _ = _parallel(3)
    print(f"3 WORKER: {dp:5.2f}s wall                              -> b2_parallel_3w.wav")
    if dp > 0:
        print(f"Speedup : {ds/dp:.2f}x  (Audio gleich lang -> nur Generierzeit zaehlt)")
    print("B2_DONE")