"""Streaming web voice demo for LiquidAI/LFM3-Audio-2B.

One WebSocket per browser tab. The browser streams mic PCM (float32, native rate) while the
talk button is held, and the server encodes and prefills it into the model's KV cache as it
arrives (see session.py). On release only the last audio token and the assistant prefix remain,
then generation streams text tokens (JSON) and decoded 24 kHz speech (binary float32 PCM) back
frame by frame, together with debug events and timings. The conversation stays in the cache
between turns.
"""

import asyncio
import json
import os
import queue
import re
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import numpy as np
import torch
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import FileResponse
from transformers import AutoModelForMultimodalLM, AutoProcessor
from transformers.generation.streamers import BaseStreamer

import tools
from session import StreamingSession

MODEL_ID = os.environ.get("LFM_MODEL", "LiquidAI/LFM3-Audio-2B")
IN_RATE = 16_000
OUT_RATE = 24_000
FRAME_SAMPLES = 1920  # one codec frame = 80 ms at 24 kHz
CHUNK_FRAMES = int(os.environ.get("LFM_CHUNK_FRAMES", "1"))  # frames per decoded audio chunk sent to the browser
MAX_TOOL_ROUNDS = 4
MIXED_START = "<|mixed_start|>"

MODES = {
    "chat": {"system": "Respond with interleaved text and audio.", "prefill": "<|mixed_start|>"},
    "asr": {"system": "Perform ASR.", "prefill": None},
    "tts": {"system": "Perform TTS in English using a female voice.", "prefill": "<|audio_start|>"},
    # no prefill: the model chooses between <tool_call> and <|mixed_start|>
    "tools": {"system": tools.SYSTEM_PROMPT, "prefill": None},
}


def pick_device() -> tuple[str, torch.dtype]:
    device = os.environ.get("LFM_DEVICE")
    if not device:
        device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
    default_dtype = "float32" if device == "cpu" else "bfloat16"
    return device, getattr(torch, os.environ.get("LFM_DTYPE", default_dtype))


class Engine:
    def __init__(self):
        self.device, self.dtype = pick_device()
        t0 = time.perf_counter()
        self.model = AutoModelForMultimodalLM.from_pretrained(
            MODEL_ID, dtype=self.dtype, device_map=self.device, trust_remote_code=True
        ).eval()
        self.processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
        # the codec is ~87x realtime on one CPU thread; keeping it on CPU avoids contending with the LM
        self.codec_device = os.environ.get("LFM_CODEC_DEVICE", "cpu")
        self.codec = self.processor.audio_tokenizer.to(device=self.codec_device, dtype=torch.float32).eval()
        self.load_s = time.perf_counter() - t0
        self.tokenizer = self.processor.tokenizer
        self.lock = threading.Lock()  # one generation at a time across all connections

    def info(self) -> dict:
        return {
            "model": MODEL_ID,
            "device": self.device,
            "dtype": str(self.dtype).removeprefix("torch."),
            "codec_device": self.codec_device,
            "load_s": round(self.load_s, 2),
            "chunk_frames": CHUNK_FRAMES,
            "interleaved_n_text": self.model.config.interleaved_n_text,
            "interleaved_n_audio": self.model.config.interleaved_n_audio,
            "torch": torch.__version__,
        }

    def build_inputs(self, conversation: list[dict], clips: list[np.ndarray], prefill: str | None, tool_defs=None):
        if prefill:
            conversation = conversation + [{"role": "assistant", "content": prefill}]
        prompt = self.processor.apply_chat_template(
            conversation,
            tokenize=False,
            continue_final_message=bool(prefill),
            add_generation_prompt=not prefill,
            tools=tool_defs,
        )
        enc = self.processor(
            text=prompt,
            audio=clips or None,
            sampling_rate=IN_RATE,
            add_special_tokens=False,
            return_tensors="pt",
        )
        keys = ("input_ids", "attention_mask", "input_features", "feature_attention_mask")
        inputs = {k: enc[k].to(self.model.device) for k in keys if k in enc}
        if "input_features" in inputs:
            inputs["input_features"] = inputs["input_features"].to(self.dtype)
        return prompt, inputs


class Cancelled(Exception):
    pass


class PlainTextStart(Exception):
    """The reply opened with plain text where a tool call or speech marker was required."""


class Emitter:
    """Thread-safe bridge from the generation thread to the connection's asyncio send queue."""

    def __init__(self, loop: asyncio.AbstractEventLoop, queue: asyncio.Queue, t0: float):
        self.loop, self.queue, self.t0 = loop, queue, t0

    def ms(self) -> float:
        return round((time.perf_counter() - self.t0) * 1000, 1)

    def send(self, item):
        self.loop.call_soon_threadsafe(self.queue.put_nowait, item)


class TextStream(BaseStreamer):
    def __init__(self, engine: Engine, emit: Emitter, cancel: threading.Event, stats: dict, first_allowed=None):
        self.tok, self.emit, self.cancel, self.stats = engine.tokenizer, emit, cancel, stats
        self.first_allowed = first_allowed  # token ids the reply may open with; None allows anything
        self.special = set(self.tok.all_special_ids)
        self.ids: list[int] = []
        self.text = ""  # everything generated, tool-call markup included
        self.visible = ""  # what the chat bubble shows
        self.prompt_seen = False

    def put(self, value):
        if self.cancel.is_set():
            raise Cancelled()
        if not self.prompt_seen:  # generate() first hands the streamer the prompt
            self.prompt_seen = True
            return
        tid = int(value.reshape(-1)[-1])
        if self.first_allowed is not None and not self.ids and tid not in self.first_allowed:
            raise PlainTextStart()
        self.first_allowed = None
        t = self.emit.ms()
        self.stats.setdefault("first_text_ms", t)
        self.stats["text_tokens"] = self.stats.get("text_tokens", 0) + 1
        raw = self.tok.decode([tid], skip_special_tokens=False)
        delta = ""
        if tid not in self.special:
            self.ids.append(tid)
            full = self.tok.decode(self.ids, skip_special_tokens=True)
            if not full.endswith("�"):  # wait for the rest of a split multi-byte char
                self.text = full
                visible = tools.visible_text(full)
                if visible.startswith(self.visible):
                    delta, self.visible = visible[len(self.visible):], visible
        self.emit.send({"type": "token", "id": tid, "raw": raw, "delta": delta, "special": tid in self.special, "t_ms": t})

    def end(self):
        pass


class AudioStream(BaseStreamer):
    """Collects speech frames and decodes them incrementally with the codec's streaming cache.

    Decoding runs on its own thread so the codec never stalls the language model's sampling loop.
    """

    def __init__(self, engine: Engine, emit: Emitter, cancel: threading.Event, stats: dict):
        self.engine, self.emit, self.cancel, self.stats = engine, emit, cancel, stats
        # the cache class lives in the codec's remote-code module
        codec_module = sys.modules[type(engine.codec).__module__]
        self.cache = codec_module.LiquidCodecStreamingCache(engine.codec.config)
        self.frames = 0
        self.decode_s = 0.0
        self.max_backlog = 0
        self.collected: list[np.ndarray] = []
        self.inbox: queue.SimpleQueue = queue.SimpleQueue()
        self.worker = threading.Thread(target=self._decode_loop, daemon=True)
        self.worker.start()

    def put(self, codes):
        if self.cancel.is_set():
            raise Cancelled()
        self.stats.setdefault("first_frame_ms", self.emit.ms())
        self.inbox.put(codes.detach().to("cpu"))
        self.frames += 1

    def end(self):
        self.inbox.put(None)
        self.worker.join()
        self.stats["audio_frames"] = self.frames
        self.stats["codec_decode_ms"] = round(self.decode_s * 1000, 1)
        self.stats["max_decode_backlog"] = self.max_backlog

    def _decode_loop(self):
        pending: list[torch.Tensor] = []
        done = False
        while not done:
            item = self.inbox.get()
            # drain whatever else is ready so a slow decode catches up in one bigger chunk
            items = [item]
            while not self.inbox.empty():
                items.append(self.inbox.get())
            done = items[-1] is None
            pending.extend(i for i in items if i is not None)
            self.max_backlog = max(self.max_backlog, len(pending))
            if done:
                if self.cache.started or pending:
                    self._flush(pending, finalize=True)
            elif len(pending) >= CHUNK_FRAMES:
                self._flush(pending, finalize=False)
                pending = []

    @torch.no_grad()
    def _flush(self, pending: list[torch.Tensor], finalize: bool):
        n = len(pending)
        codes = (
            torch.stack(pending, dim=-1)[None].to(self.engine.codec_device)
            if n
            else torch.zeros((1, self.engine.codec.config.n_q, 0), dtype=torch.long)
        )
        first_codes = pending[0].tolist() if n else None
        t = time.perf_counter()
        pcm = self.engine.codec.decode(codes, streaming_cache=self.cache, finalize=finalize).audio_values[0]
        self.decode_s += time.perf_counter() - t
        pcm = pcm.cpu().numpy().astype(np.float32)
        if not pcm.size:
            return
        self.collected.append(pcm)
        t_ms = self.emit.ms()
        self.stats.setdefault("first_audio_ms", t_ms)
        self.emit.send(
            {"type": "audio_chunk", "frames": n, "samples": int(pcm.size), "total_frames": self.frames,
             "codes": first_codes, "final": finalize, "t_ms": t_ms}
        )
        self.emit.send(pcm.tobytes())


app = FastAPI()
engine: Engine | None = None
STATIC = Path(__file__).parent / "static"


@app.on_event("startup")
def load():
    global engine
    engine = Engine()
    print("model ready:", json.dumps(engine.info()))


@app.get("/")
def index():
    return FileResponse(STATIC / "index.html")


@app.websocket("/ws")
async def ws_endpoint(ws: WebSocket):
    await ws.accept()
    loop = asyncio.get_running_loop()
    out: asyncio.Queue = asyncio.Queue()
    # every model call for this connection runs on one thread, in arrival order: prefill chunks, then generation
    executor = ThreadPoolExecutor(max_workers=1)

    def run(fn, *args):
        return loop.run_in_executor(executor, fn, *args)

    async def sender():
        while True:
            item = await out.get()
            if isinstance(item, (bytes, bytearray)):
                await ws.send_bytes(item)
            else:
                await ws.send_text(json.dumps(item))

    send_task = asyncio.create_task(sender())
    await out.put({"type": "ready", **engine.info()})

    session = StreamingSession(engine)
    settings = {"mode": "chat", "max_new_tokens": 1024, "audio_temperature": 1.0, "audio_top_k": 128,
                "system": "", "voice": "female", "keep_history": True, "force_speech": True}
    cancel = threading.Event()
    gen_task: asyncio.Task | None = None
    recording = False
    rec: list[np.ndarray] = []
    rec_rate = IN_RATE
    mic_inbox: queue.SimpleQueue = queue.SimpleQueue()  # chunks waiting to be prefilled
    speech = {}  # prefill stats for the utterance being recorded

    def log(msg, **kw):
        out.put_nowait({"type": "log", "msg": msg, **kw})

    def voice_mode():
        return settings["mode"] if settings["mode"] in ("chat", "asr", "tools") else "chat"

    def system_text(mode):
        if settings["system"].strip():
            return settings["system"].strip()
        if mode == "tts":
            return f"Perform TTS in English using a {settings['voice']} voice."
        return MODES[mode]["system"]

    # ---- jobs that run on the executor thread ----------------------------------------------------------------

    def open_turn(mode):
        """Close the previous turn and open a user turn in the cached context."""
        sys_txt = system_text(mode)
        block = engine.processor.apply_chat_template(
            [{"role": "system", "content": sys_txt}], tools=tools.TOOLS if mode == "tools" else None, tokenize=False
        )
        with engine.lock:
            t = time.perf_counter()
            info = session.begin_turn((mode, sys_txt), block, settings["keep_history"] and mode != "asr")
            info["open_ms"] = round((time.perf_counter() - t) * 1000, 1)
        if "context_reset" in info:
            emit_now({"type": "log", "msg": f"context reset ({info['context_reset']})"})
        speech.clear()
        speech.update(info, prefill_calls=0, prefill_busy_ms=0.0)

    def feed_pending(final=False):
        chunks = []
        while not mic_inbox.empty():
            chunks.append(mic_inbox.get())
        if not chunks and not final:
            return
        pcm = np.concatenate(chunks) if chunks else np.zeros(0, np.float32)
        t = time.perf_counter()
        with engine.lock:
            session.feed_audio(pcm, rec_rate, final=final)
        ms = (time.perf_counter() - t) * 1000
        if final:
            speech["flush_ms"] = round(ms, 1)
        else:
            speech["prefill_calls"] += 1
            speech["prefill_busy_ms"] += ms

    def emit_now(item):
        loop.call_soon_threadsafe(out.put_nowait, item)

    def run_round(gen, prompt_text, prompt_tokens, mode, emit, round_no, first_allowed=None):
        stats: dict = {"prompt_tokens": prompt_tokens}
        emit.send({"type": "prompt", "text": prompt_text, "tokens": prompt_tokens, "round": round_no, "t_ms": emit.ms()})
        text_stream = TextStream(engine, emit, cancel, stats, first_allowed)
        audio_stream = AudioStream(engine, emit, cancel, stats) if mode != "asr" else None
        kwargs = dict(
            max_new_tokens=int(settings["max_new_tokens"]),
            audio_temperature=float(settings["audio_temperature"]),
            audio_top_k=int(settings["audio_top_k"]),
            streamer=text_stream,
            audio_streamer=audio_stream,
        )
        if mode == "asr":
            kwargs.update(generation_mode="text")
        stats["cancelled"] = False
        with engine.lock, torch.inference_mode():
            gen_t0 = time.perf_counter()
            try:
                gen(**kwargs)
            except Cancelled:
                stats["cancelled"] = True
            except PlainTextStart:
                stats["plain_start"] = True
            stats["gen_ms"] = round((time.perf_counter() - gen_t0) * 1000, 1)
        audio_s = sum(c.size for c in audio_stream.collected) / OUT_RATE if audio_stream else 0.0
        stats["audio_out_s"] = round(audio_s, 2)
        return text_stream, stats

    def session_round(suffix, mode, emit, round_no, first_allowed=None):
        cached = session.length
        n = len(engine.tokenizer(suffix, add_special_tokens=False).input_ids)
        shown = f"[{cached} cached tokens]" + suffix
        return run_round(lambda **kw: session.generate(suffix, **kw), shown, cached + n, mode, emit, round_no, first_allowed)

    def run_session_turn(user_text, mode, t0, user_meta):
        """Generate the reply to the open user turn; in tools mode, run calls and continue until the model answers."""
        emit = Emitter(loop, out, t0)
        if user_text is None:  # spoken: flush the encoder's last token
            feed_pending(final=True)
        open_suffix = (user_text or "") + "<|im_end|>\n<|im_start|>assistant\n"
        prefill = MODES[mode]["prefill"] or ""
        force_speech = mode == "tools" and settings["force_speech"]
        speech_markers = {engine.tokenizer.convert_tokens_to_ids(t) for t in ("<tool_call>", "<|mixed_start|>", "<|audio_start|>")}
        rounds: list[dict] = []
        suffix = open_suffix + prefill
        try:
            for round_no in range(MAX_TOOL_ROUNDS + 1):
                # the model sometimes answers in plain text only; force speech by prefilling the interleaved marker
                # after a tool result, and by restarting a first round that opens with plain text
                if force_speech and round_no:
                    suffix += MIXED_START
                guard = speech_markers if force_speech and not round_no else None
                snap = session.snapshot()
                text_stream, r = session_round(suffix, mode, emit, round_no, guard)
                if r.get("plain_start"):
                    emit.send({"type": "log", "msg": f"round {round_no} opened with plain text; restarting with speech"})
                    rounds.append(r)
                    session.restore(snap)
                    text_stream, r = session_round(suffix + MIXED_START, mode, emit, round_no)
                rounds.append(r)
                if r["cancelled"]:
                    break
                calls = tools.parse_tool_calls(text_stream.text) if mode == "tools" else []
                if not calls or round_no == MAX_TOOL_ROUNDS:
                    break
                results = []
                for call in calls:
                    fn = call["function"]
                    emit.send({"type": "tool_call", "name": fn["name"], "arguments": fn["arguments"], "round": round_no,
                               "t_ms": emit.ms()})
                    t = time.perf_counter()
                    result = tools.run_tool(fn["name"], fn["arguments"],
                                            lambda action, value: emit.send({"type": "ui", "action": action, "value": value}))
                    results.append(result)
                    emit.send({"type": "tool_result", "name": fn["name"], "result": result,
                               "ms": round((time.perf_counter() - t) * 1000, 1), "t_ms": emit.ms()})
                suffix = (session.close_prefix() + "<|im_start|>tool\n<tool_response>\n" + json.dumps(results)
                          + "\n</tool_response><|im_end|>\n<|im_start|>assistant\n")
        except BaseException:
            session.restore(session.turn_snapshot)
            raise
        if rounds[-1]["cancelled"]:
            session.restore(session.turn_snapshot)  # a cancelled turn leaves no trace in the context
        stats = summarize(rounds, mode, emit, user_meta)
        stats.update({k: (round(v, 1) if isinstance(v, float) else v) for k, v in speech.items()},
                     context_tokens_after=session.length)
        return text_stream.visible, stats

    def run_tts(text, t0):
        emit = Emitter(loop, out, t0)
        conv = [{"role": "system", "content": system_text("tts")}, {"role": "user", "content": text}]
        prompt, inputs = engine.build_inputs(conv, [], MODES["tts"]["prefill"])
        text_stream, r = run_round(lambda **kw: engine.model.generate(**inputs, return_dict_in_generate=True, **kw),
                                   prompt, int(inputs["input_ids"].shape[1]), "tts", emit, 0)
        return text_stream.visible, summarize([r], "tts", emit, {})

    # ---- event-loop side -------------------------------------------------------------------------------------

    async def generate(job, *args):
        cancel.clear()
        try:
            text, stats = await run(job, *args)
            await out.put({"type": "done", "text": text, "stats": stats})
        except Exception as e:  # surface model errors in the UI instead of killing the socket
            import traceback

            traceback.print_exc()
            await out.put({"type": "error", "msg": f"{type(e).__name__}: {e}"})

    try:
        while True:
            msg = await ws.receive()
            if msg["type"] == "websocket.disconnect":
                break
            if msg.get("bytes") is not None:
                if recording:
                    chunk = np.frombuffer(msg["bytes"], dtype=np.float32)
                    rec.append(chunk)
                    mic_inbox.put(chunk)
                    run(feed_pending)  # queued behind open_turn; drains whatever has arrived by then
                continue
            data = json.loads(msg["text"])
            kind = data.get("type")
            busy = gen_task is not None and not gen_task.done()

            if kind == "settings":
                settings.update({k: v for k, v in data.items() if k in settings})
                log("settings updated", settings=settings)
            elif kind == "start":
                if busy:
                    log("busy: still generating, utterance ignored")
                    continue
                recording, rec, rec_rate = True, [], int(data.get("rate", IN_RATE))
                run(open_turn, voice_mode())
                log(f"recording @ {rec_rate} Hz, prefilling while you speak")
            elif kind == "stop":
                if not recording:
                    continue
                recording = False
                t0 = time.perf_counter()
                audio = np.concatenate(rec) if rec else np.zeros(0, np.float32)
                dur = audio.size / rec_rate
                rms = float(np.sqrt(np.mean(audio**2))) if audio.size else 0.0
                peak = float(np.abs(audio).max()) if audio.size else 0.0
                user_meta = {"user_audio_s": round(dur, 2), "user_rms": round(rms, 4), "user_peak": round(peak, 3)}
                await out.put({"type": "user_audio", **user_meta})
                if dur < 0.3:
                    log("utterance too short, ignored")
                    run(lambda: session.restore(session.turn_snapshot))
                    continue
                gen_task = asyncio.create_task(generate(run_session_turn, None, voice_mode(), t0, user_meta))
            elif kind == "text":
                if busy or recording:
                    log("busy")
                    continue
                t0 = time.perf_counter()
                run(open_turn, voice_mode())
                gen_task = asyncio.create_task(generate(run_session_turn, data["text"], voice_mode(), t0, {}))
            elif kind == "tts":
                if busy or recording:
                    log("busy")
                    continue
                gen_task = asyncio.create_task(generate(run_tts, data["text"], time.perf_counter()))
            elif kind == "cancel":
                cancel.set()
                log("cancel requested")
            elif kind == "reset":
                run(session.reset)
                log("context cleared")
    except WebSocketDisconnect:
        pass
    finally:
        cancel.set()
        send_task.cancel()
        executor.shutdown(wait=False, cancel_futures=True)


def summarize(rounds: list[dict], mode: str, emit: Emitter, user_meta: dict) -> dict:
    gen_s = sum(r["gen_ms"] for r in rounds) / 1000
    audio_s = sum(r["audio_out_s"] for r in rounds)
    total = lambda key: sum(r.get(key, 0) for r in rounds)  # noqa: E731
    first = lambda key: next((r[key] for r in rounds if key in r), None)  # noqa: E731
    return {
        "mode": mode, **user_meta,
        "rounds": len(rounds),
        "speech_restarts": sum(1 for r in rounds if r.get("plain_start")),
        "prompt_tokens": rounds[-1]["prompt_tokens"],
        "first_text_ms": first("first_text_ms"),
        "first_frame_ms": first("first_frame_ms"),
        "first_audio_ms": first("first_audio_ms"),
        "text_tokens": total("text_tokens"),
        "audio_frames": total("audio_frames"),
        "codec_decode_ms": round(total("codec_decode_ms"), 1),
        "max_decode_backlog": max((r.get("max_decode_backlog", 0) for r in rounds), default=0),
        "cancelled": rounds[-1]["cancelled"],
        "gen_ms": round(gen_s * 1000, 1),
        "total_ms": emit.ms(),
        "audio_out_s": round(audio_s, 2),
        "rtf": round(gen_s / audio_s, 3) if audio_s else None,
        "text_tok_per_s": round(total("text_tokens") / gen_s, 1) if gen_s else None,
        "frames_per_s": round(total("audio_frames") / gen_s, 1) if gen_s else None,
        "per_round": rounds,
    }


def collapse_placeholders(prompt: str) -> str:
    """Show runs of audio placeholders as <|audio|>×N so the prompt stays readable."""

    return re.sub(r"((<\|audio(?:_out)?\|>))\1+", lambda m: f"{m.group(2)}×{len(m.group(0)) // len(m.group(2))}", prompt)


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host=os.environ.get("HOST", "127.0.0.1"), port=int(os.environ.get("PORT", "8000")))