"""Check StreamingSession: streamed ASR must match offline ASR; two-turn chat from the cache; latency after release."""

import sys
import time
from pathlib import Path

import librosa
import numpy as np
import soundfile as sf
import torch

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

engine = server.Engine()
tok = engine.tokenizer
wav, rate = sf.read(Path(__file__).resolve().parent.parent / "smoke_out/tts.wav", dtype="float32")
CHUNK = 2048  # what the browser worklet sends


def system_block(text):
    return engine.processor.apply_chat_template([{"role": "system", "content": text}], tokenize=False)


def stream_in(s, audio, rate):
    for i in range(0, len(audio), CHUNK):
        s.feed_audio(audio[i:i + CHUNK], rate)


def decode_text(out, suffix_len):
    seq = out.sequences[0, suffix_len:].tolist()
    return tok.decode([t for t in seq if t != engine.model.config.audio_out_token_id], skip_special_tokens=True)


# offline reference
_, inputs = engine.build_inputs([{"role": "system", "content": "Perform ASR."}, {"role": "user", "content": [{"type": "audio"}]}],
                                [librosa.resample(wav, orig_sr=rate, target_sr=16000)], None)
with torch.inference_mode():
    ref = engine.model.generate(**inputs, max_new_tokens=64, generation_mode="text", return_dict_in_generate=True)
print("offline ASR :", tok.decode(ref.sequences[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True))

s = StreamingSession(engine)
s.begin_turn(("asr", ""), system_block("Perform ASR."), keep_history=False)
stream_in(s, wav, rate)
t = time.perf_counter()
s.feed_audio(np.zeros(0, np.float32), rate, final=True)
out = s.generate("<|im_end|>\n<|im_start|>assistant\n", max_new_tokens=64, generation_mode="text")
print(f"streamed ASR: {decode_text(out, 5)}  ({(time.perf_counter() - t) * 1000:.0f} ms after release, {s.audio_tokens} audio tokens)")


class FirstFrame:
    def __init__(self): self.t = None
    def put(self, v): self.t = self.t or time.perf_counter()
    def end(self): pass


sysb = system_block("Respond with interleaved text and audio.")
conv = [{"role": "system", "content": "Respond with interleaved text and audio."}, {"role": "user", "content": [{"type": "audio"}]}]
w16 = librosa.resample(wav, orig_sr=rate, target_sr=16000)
for label in ("warmup", "offline"):
    ff = None

    class FF:
        t = None
        def put(self, v): self.t = self.t or time.perf_counter()
        def end(self): pass
    ff = FF()
    t = time.perf_counter()
    _, inputs = engine.build_inputs(conv, [w16], "<|mixed_start|>")
    with torch.inference_mode():
        engine.model.generate(**inputs, max_new_tokens=512, audio_streamer=ff, return_dict_in_generate=True)
    print(f"{label} chat: first_frame={(ff.t - t) * 1000:.0f}ms after release (prep + full prefill)")

s = StreamingSession(engine)
for turn in range(3):
    info = s.begin_turn(("chat", ""), sysb, keep_history=True)
    stream_in(s, wav, rate)
    t = time.perf_counter()
    s.feed_audio(np.zeros(0, np.float32), rate, final=True)
    t_flush = time.perf_counter()
    ff = FirstFrame()
    out = s.generate("<|im_end|>\n<|im_start|>assistant\n<|mixed_start|>", max_new_tokens=512, audio_streamer=ff)
    print(f"chat turn {turn}: {info} flush={(t_flush - t) * 1000:.0f}ms first_frame={(ff.t - t) * 1000:.0f}ms "
          f"total={(time.perf_counter() - t):.1f}s ctx={s.length} close={s.needs_close}\n  {decode_text(out, 6)!r}")

# snapshot / restore round trip
snap = s.snapshot()
n = s.length
s.begin_turn(("chat", ""), sysb, keep_history=True)
s.restore(snap)
print("restore ok:", s.length == n)