"""Headless check of the server's generation path: TTS -> ASR -> voice chat, with streaming."""

import asyncio
import sys
import threading
import time
from pathlib import Path

import librosa
import soundfile as sf
import torch

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import server  # noqa: E402

OUT = Path(__file__).resolve().parent.parent / "smoke_out"
OUT.mkdir(exist_ok=True)


def run(engine, conv, clips, prefill, mode):
    loop = asyncio.new_event_loop()
    queue = asyncio.Queue()
    stats, cancel = {}, threading.Event()
    emit = server.Emitter(loop, queue, time.perf_counter())
    emit.send = lambda item: None
    _, inputs = engine.build_inputs(conv, clips, prefill)
    text = server.TextStream(engine, emit, cancel, stats)
    audio = server.AudioStream(engine, emit, cancel, stats) if mode != "asr" else None
    t = time.perf_counter()
    with torch.inference_mode():
        engine.model.generate(**inputs, max_new_tokens=512, streamer=text, audio_streamer=audio,
                              return_dict_in_generate=True, **({"generation_mode": "text"} if mode == "asr" else {}))
    gen = time.perf_counter() - t
    pcm = torch.cat([torch.from_numpy(c) for c in audio.collected]).numpy() if audio and audio.collected else None
    dur = pcm.size / server.OUT_RATE if pcm is not None else 0
    print(f"[{mode}] prompt={inputs['input_ids'].shape[1]} gen={gen:.2f}s audio={dur:.2f}s stats={stats}\n  text={text.text!r}")
    return text.text, pcm


engine = server.Engine()
print(engine.info())
_, pcm = run(engine, [{"role": "system", "content": "Perform TTS in English using a female voice."},
                      {"role": "user", "content": "What is the tallest mountain in Europe?"}], [], "<|audio_start|>", "tts")
sf.write(OUT / "tts.wav", pcm, server.OUT_RATE)
q16 = librosa.resample(pcm, orig_sr=server.OUT_RATE, target_sr=server.IN_RATE)
run(engine, [{"role": "system", "content": "Perform ASR."}, {"role": "user", "content": [{"type": "audio"}]}], [q16], None, "asr")
reply, pcm = run(engine, [{"role": "system", "content": "Respond with interleaved text and audio."},
                          {"role": "user", "content": [{"type": "audio"}]}], [q16], "<|mixed_start|>", "chat")
sf.write(OUT / "chat.wav", pcm, server.OUT_RATE)
# sanity-check the chat reply's speech by transcribing it
r16 = librosa.resample(pcm, orig_sr=server.OUT_RATE, target_sr=server.IN_RATE)
run(engine, [{"role": "system", "content": "Perform ASR."}, {"role": "user", "content": [{"type": "audio"}]}], [r16], None, "asr")