#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""CosyVoice3 streaming e2e benchmark: concurrency x utterance-length matrix.

Launches the server once per (concurrency, step-size) combo, streams TTS
requests cycling through several utterance lengths, and reports
TTFA/E2E/RTF/per-chunk latency (mean, stdev, range) per cell. "Step size" is
``codec_chunk_frames``, the token-hop length the first emitted chunk waits
for; later chunks wait for a doubling hop capped at ``4 * codec_chunk_frames``
(see ``_token_hop_lengths``). Per-chunk latency is reconstructed from that
same ramp rather than from raw transport read boundaries, which don't align
to model chunks (see ``one_request``'s docstring).

Running with no flags reproduces the default matrix. Per-request correctness
is checked: non-silent, long enough, and (with a fixed ``--seed``)
byte-identical PCM across requests for the same text.
"""

import argparse
import json
import os
import signal
import statistics
import subprocess
import sys
import threading
import time
from pathlib import Path

import httpx
import numpy as np

REPO = Path(__file__).resolve().parents[2]
DEFAULT_MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512"

DEFAULT_REF_AUDIO = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
DEFAULT_REF_TEXT = "希望你以后能够做的比我还好呦。"

DEFAULT_TEXTS = [
    "Thank you so much for the birthday gift you sent me.",
    "The weather is lovely today, with warm sunshine and a gentle breeze. "
    "We decided to take a walk in the park to enjoy the spring flowers and the fresh air.",
    "The weather is lovely today, with warm sunshine and a gentle breeze. We decided to take a walk in the "
    "park to enjoy the spring flowers and the fresh air. Many people were out walking, children were "
    "running and playing on the grass, and older folks sat on benches chatting with one another. We "
    "strolled slowly along the path and came across a cherry tree in full bloom, its pink petals drifting "
    "on the wind. It was so beautiful that we stopped to take some photos to remember the moment. Then we "
    "continued on and reached a small lake, its water so clear you could see straight to the bottom, with "
    "the blue sky and white clouds reflected on its surface. A few ducks paddled back and forth across the "
    "water, occasionally letting out a soft quack. We sat on a bench by the lake for a while, enjoying the "
    "peaceful afternoon. Time passed quickly, and before we knew it evening had arrived, with the setting "
    "sun painting the sky in shades of gold. Reluctant to leave, we finally headed home, already planning "
    "our next visit. On the way back we talked about everything we had seen and heard that day, feeling "
    "wonderfully happy. It was truly an unforgettable weekend.",
]

# PCM16 mono @ 24 kHz (CosyVoice3 output rate)
SAMPLE_RATE = 24000

# Well below any real speech's RMS; catches all-zero/near-silent PCM that byte
# count and duration checks alone pass through.
MIN_RMS = 50.0

# Overlay: only codec_chunk_frames is overridden; must live next to cosyvoice3.yaml (relative base_config)
DEPLOY_TEMPLATE = """# Auto-generated by bench_cosyvoice3_concurrency_matrix.py -- do not edit.
base_config: cosyvoice3.yaml

connectors:
  connector_of_shared_memory:
    name: SharedMemoryConnector
    extra:
      codec_streaming: true
      connector_get_sleep_s: 0.01
      connector_get_max_wait_first_chunk: 3000
      connector_get_max_wait: 300
      codec_chunk_frames: {chunk_frames}
      codec_max_chunk_frames: {max_chunk_frames}
      codec_stream_scale_factor: {stream_scale_factor}
      codec_left_context_frames: 25
      codec_vocab_size: 6561
"""

# CosyVoice3Config constants (transformers_utils/configs/cosyvoice3.py): tokens
# are consumed at token_mel_ratio mel frames each, and HiFT emits
# samples_per_mel PCM16 samples per mel frame. Pinned here (not left to the
# server's defaults) since one_request's byte-boundary reconstruction below
# depends on knowing them exactly.
TOKEN_MEL_RATIO = 2
SAMPLES_PER_MEL = 480
STREAM_SCALE_FACTOR = 2


def _stats(values: list[float]) -> dict[str, float]:
    """Mean/median/stdev/p95/min/max/n for a list of per-request measurements."""
    n = len(values)
    if n == 0:
        return {
            "n": 0,
            "mean": float("nan"),
            "median": float("nan"),
            "stdev": float("nan"),
            "p95": float("nan"),
            "min": float("nan"),
            "max": float("nan"),
        }
    return {
        "n": n,
        "mean": statistics.mean(values),
        "median": statistics.median(values),
        "stdev": statistics.stdev(values) if n >= 2 else 0.0,
        "p95": statistics.quantiles(values, n=20)[18] if n >= 2 else values[0],
        "min": min(values),
        "max": max(values),
    }


def vllm_bin() -> str:
    for c in (REPO / ".venv/bin/vllm", Path(sys.executable).parent / "vllm"):
        if c.is_file():
            return str(c)
    return "vllm"


def wait_ready(url: str, proc: subprocess.Popen, log: Path, timeout: float) -> None:
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if proc.poll() is not None:
            tail = log.read_text(errors="replace").splitlines()[-40:] if log.exists() else []
            raise RuntimeError("server exited during startup:\n  " + "\n  ".join(tail))
        try:
            if httpx.get(f"{url}/health", timeout=2.0).status_code == 200:
                return
        except Exception:
            pass
        time.sleep(2.0)
    raise TimeoutError(f"server not ready within {timeout}s (log: {log})")


def _token_hop_lengths(chunk_frames: int, max_chunk_frames: int):
    """The token-hop-length ramp talker2code2wav_async_chunk uses to decide
    when to fire the next flow-solve/emission: chunk_frames, then doubling
    (STREAM_SCALE_FACTOR) each step, capped at max_chunk_frames. This -- not
    codec_chunk_frames alone -- is what actually paces emitted audio chunks:
    with the defaults here, hops are chunk_frames, 2x, 4x, 4x, 4x, ...

    The very first emission is consistently one token short of chunk_frames
    (confirmed against real speech.audio.delta event sizes) even though the
    ramp's own progression is computed from the nominal chunk_frames, not the
    actual first hop -- so only the first yield gets the -1 adjustment.
    """
    yield chunk_frames - 1
    hop = chunk_frames
    while True:
        hop = min(max_chunk_frames, max(chunk_frames, hop * STREAM_SCALE_FACTOR))
        yield hop


def one_request(client, url, model, text, ref_audio, ref_text, chunk_frames, seed=None):
    """Run one streaming request, returning TTFA, E2E, per-chunk gaps, and audio bytes.

    ``response_format=pcm`` + ``stream_format=audio`` streams raw bytes with no
    framing, so a client read boundary alone is a transport artifact, not a
    model chunk boundary. Reconstructed instead by replicating the server's
    own token-hop-length ramp (see ``_token_hop_lengths``) to compute each
    real chunk's expected byte size ahead of time, and slicing the
    accumulated byte stream at those offsets -- independent of how the
    transport happens to fragment the underlying reads. The final chunk
    (finalize, which releases extra held-back frames) is whatever remains
    when the stream closes, not a predicted hop.
    """
    payload = {
        "model": model,
        "input": text,
        "response_format": "pcm",
        "stream": True,
        "stream_format": "audio",
        "ref_audio": ref_audio,
        "ref_text": ref_text,
    }
    if seed is not None:
        payload["seed"] = seed
    t0 = time.perf_counter()
    ttfa = None
    nbytes = 0
    audio_bytes = bytearray()
    chunk_gaps_ms = []  # inter-chunk arrival gaps (per model chunk), by step index
    prev_t = None
    step_idx = 0

    bytes_per_token = TOKEN_MEL_RATIO * SAMPLES_PER_MEL * 2
    hop_gen = _token_hop_lengths(chunk_frames, 4 * chunk_frames)
    next_boundary = next(hop_gen) * bytes_per_token
    last_crossed_bytes = 0
    received = 0
    last_now = t0

    with client.stream("POST", f"{url}/v1/audio/speech", json=payload, headers={"Authorization": "Bearer EMPTY"}) as r:
        if r.status_code != 200:
            raise RuntimeError(f"{r.status_code}: {r.read()[:300]!r}")
        for chunk in r.iter_bytes():
            if not chunk:
                continue
            now = time.perf_counter()
            if ttfa is None:
                ttfa = now - t0
            last_now = now
            received += len(chunk)
            nbytes += len(chunk)
            audio_bytes += chunk
            while received >= next_boundary:
                if prev_t is not None:
                    chunk_gaps_ms.append((step_idx, (now - prev_t) * 1e3))
                    step_idx += 1
                prev_t = now
                last_crossed_bytes = next_boundary
                next_boundary += next(hop_gen) * bytes_per_token

    if received > last_crossed_bytes:
        # Finalize chunk: releases extra held-back frames, so it doesn't
        # match a predicted hop size -- its arrival is the stream closing.
        if prev_t is not None:
            chunk_gaps_ms.append((step_idx, (last_now - prev_t) * 1e3))
            step_idx += 1

    e2e = time.perf_counter() - t0
    audio_s = nbytes / 2 / SAMPLE_RATE
    return {
        "ttfa_s": ttfa,
        "e2e_s": e2e,
        "audio_s": audio_s,
        "rtf": e2e / audio_s if audio_s > 0 else float("nan"),
        "chunk_gaps_ms": chunk_gaps_ms,
        "nbytes": nbytes,
        "audio_bytes": audio_bytes,
    }


def check_correctness(r: dict, min_audio_s: float) -> list[str]:
    """Return a list of correctness violations (empty = pass)."""
    problems = []
    if r["nbytes"] <= 0:
        problems.append("empty audio (0 bytes)")
    if r["audio_s"] < min_audio_s:
        problems.append(f"audio too short ({r['audio_s']:.2f}s < {min_audio_s:.2f}s)")
    if r["nbytes"] > 0:
        samples = np.frombuffer(bytes(r["audio_bytes"]), dtype=np.int16).astype(np.float64)
        rms = float(np.sqrt(np.mean(samples**2)))
        if rms < MIN_RMS:
            problems.append(f"near-silent audio (rms={rms:.1f} < {MIN_RMS})")
    return problems


def drive(url, args, n, conc, label, results, lock, idx, chunk_frames):
    """Run ``n`` requests at concurrency ``conc``, cycling through ``args.texts``."""
    texts = args.texts

    def worker():
        with httpx.Client(timeout=600.0) as c:
            while True:
                with lock:
                    i = idx[0]
                    idx[0] += 1
                if i >= n:
                    return
                text = texts[i % len(texts)]
                try:
                    r = one_request(
                        c, url, args.model, text, args.ref_audio, args.ref_text, chunk_frames, seed=args.seed
                    )
                    r["text"] = text
                    with lock:
                        results.append(r)
                except Exception as e:
                    print(f"    [{label}] request {i} failed: {e}", flush=True)

    ts = [threading.Thread(target=worker) for _ in range(conc)]
    t0 = time.perf_counter()
    for t in ts:
        t.start()
    for t in ts:
        t.join()
    return time.perf_counter() - t0


def run_combo(args, conc: int, chunk_frames: int, out: Path) -> dict:
    tag = f"c{conc}_chunk{chunk_frames}"
    run_dir = out / tag
    run_dir.mkdir(parents=True, exist_ok=True)
    log = run_dir / "server.log"

    deploy = REPO / "vllm_omni/deploy" / f"_autogen_steps_{tag}.yaml"
    deploy.write_text(
        DEPLOY_TEMPLATE.format(
            chunk_frames=chunk_frames,
            max_chunk_frames=4 * chunk_frames,
            stream_scale_factor=STREAM_SCALE_FACTOR,
        )
    )

    cmd = [
        vllm_bin(),
        "serve",
        args.model,
        "--omni",
        "--deploy-config",
        str(deploy),
        "--host",
        "127.0.0.1",
        "--port",
        str(args.port),
        "--allowed-local-media-path",
        str(REPO),
    ]
    env = dict(os.environ)
    if args.cuda_visible_devices:
        env["CUDA_VISIBLE_DEVICES"] = args.cuda_visible_devices
    if args.trt is not None:
        env["COSYVOICE3_TRT"] = args.trt

    print(f"\n=== concurrency={conc} step_size={chunk_frames} ===", flush=True)
    print(f"  launching server (log: {log})", flush=True)
    with open(log, "w") as lf:
        proc = subprocess.Popen(cmd, env=env, stdout=lf, stderr=subprocess.STDOUT, start_new_session=True)
    url = f"http://127.0.0.1:{args.port}"
    try:
        wait_ready(url, proc, log, args.server_timeout)
        print("  server ready", flush=True)

        results: list[dict] = []
        lock = threading.Lock()
        idx = [0]
        if args.warmup:
            print(f"  warmup ({args.warmup} req)", flush=True)
            drive(url, args, args.warmup, min(conc, args.warmup), "warmup", results, lock, idx, chunk_frames)
            results.clear()
            idx[0] = 0

        print(f"  running {args.num_prompts} req @ concurrency {conc}", flush=True)
        wall = drive(url, args, args.num_prompts, conc, "bench", results, lock, idx, chunk_frames)

        ok = [r for r in results if r.get("ttfa_s") is not None]
        problems = []
        for r in ok:
            problems += check_correctness(r, args.min_audio_s)
        # With a fixed seed, same-text requests must produce byte-identical PCM (per text group)
        if args.seed is not None:
            by_text: dict[str, bytes] = {}
            for r in ok:
                t = r.get("text")
                ref = by_text.get(t)
                if ref is None:
                    by_text[t] = r.get("audio_bytes")
                elif r.get("audio_bytes") != ref:
                    problems.append(f"byte-identity violation for text {t!r}: seeded requests produced different PCM")
        n_ok = len(ok)
        total_audio_s = sum(r["audio_s"] for r in ok)
        all_gaps = [g for r in ok for g in r["chunk_gaps_ms"]]  # (step_index, gap_ms) per request
        gap_vals = [g for _, g in all_gaps]

        # bucket gaps by chunk position within a request, to see any O(n) slope
        by_idx: dict[int, list[float]] = {}
        for step_index, g in all_gaps:
            by_idx.setdefault(step_index, []).append(g)
        step_latency_by_idx = {i: statistics.median(vals) for i, vals in sorted(by_idx.items())}

        agg = {
            "concurrency": conc,
            "step_size": chunk_frames,
            "ok": n_ok,
            "errors": args.num_prompts - n_ok,
            "correctness_problems": len(problems),
            "correctness_sample": problems[:5],
            "total_audio_s": total_audio_s,
            "throughput_audio_s_per_s": total_audio_s / wall if wall > 0 else float("nan"),
            "ttfa_s": _stats([r["ttfa_s"] for r in ok]),
            "e2e_s": _stats([r["e2e_s"] for r in ok]),
            "rtf": _stats([r["rtf"] for r in ok]),
            "step_latency_ms": _stats(gap_vals),
            "step_latency_ms_by_idx": step_latency_by_idx,
            "num_steps": len(all_gaps),
            # per-utterance-length breakdown; n per text = num_prompts / len(texts)
            "by_text": {
                t: {
                    "audio_s": _stats([r["audio_s"] for r in ok if r.get("text") == t]),
                    "e2e_s": _stats([r["e2e_s"] for r in ok if r.get("text") == t]),
                    "rtf": _stats([r["rtf"] for r in ok if r.get("text") == t]),
                }
                for t in sorted({r.get("text") for r in ok})
            },
        }
        print(json.dumps(agg, indent=2, ensure_ascii=False), flush=True)
        return agg
    finally:
        print("  shutting down server", flush=True)
        try:
            os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
        except Exception:
            proc.terminate()
        try:
            proc.wait(timeout=60)
        except subprocess.TimeoutExpired:
            try:
                os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
            except Exception:
                proc.kill()
            proc.wait(timeout=20)
        deploy.unlink(missing_ok=True)


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--model", default=DEFAULT_MODEL)
    p.add_argument(
        "--text",
        dest="texts",
        nargs="+",
        default=DEFAULT_TEXTS,
        help="One or more utterance texts. With several, requests cycle through "
        "them so a single server launch covers all utterance lengths; results "
        "are bucketed per text.",
    )
    p.add_argument("--ref-audio", default=DEFAULT_REF_AUDIO)
    p.add_argument("--ref-text", default=DEFAULT_REF_TEXT)
    p.add_argument("--concurrency", type=int, nargs="+", default=[1, 8])
    p.add_argument(
        "--step-size",
        type=int,
        nargs="+",
        default=[15],
        help="codec_chunk_frames values to sweep (the streaming step size).",
    )
    p.add_argument("--num-prompts", type=int, default=30)
    p.add_argument("--warmup", type=int, default=2)
    p.add_argument("--port", type=int, default=8091)
    p.add_argument(
        "--server-timeout",
        type=float,
        default=1800.0,
        help="Startup budget; first run also downloads weights and builds the TRT engine.",
    )
    p.add_argument("--output-dir", default="./perf_steps")
    p.add_argument("--cuda-visible-devices", default="0")
    p.add_argument(
        "--trt",
        choices=["0", "1"],
        default=None,
        help="Pin COSYVOICE3_TRT. 0=force torch estimator, 1=allow TensorRT. Default: unset.",
    )
    p.add_argument(
        "--min-audio-s",
        type=float,
        default=1.0,
        help="Minimum expected audio seconds per request for the correctness check.",
    )
    p.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Fixed seed for reproducible (byte-identical) audio. Pass --seed '' "
        "equivalent (any non-int) to disable -- there is no unset sentinel here "
        "by design, since the default sweep is meant to be exactly reproducible.",
    )
    return p.parse_args()


def main() -> None:
    args = parse_args()
    out = Path(args.output_dir).resolve()
    out.mkdir(parents=True, exist_ok=True)

    all_results = []
    failed = False
    for conc in args.concurrency:
        for chunk_frames in args.step_size:
            try:
                result = run_combo(args, conc, chunk_frames, out)
                all_results.append(result)
                if result["errors"] > 0 or result["correctness_problems"] > 0:
                    failed = True
            except Exception as e:
                print(f"  combo c={conc} chunk={chunk_frames} FAILED: {e}", flush=True)
                failed = True

    out_path = out / "steps_results.json"
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump({"model": args.model, "texts": args.texts, "runs": all_results}, f, indent=2, ensure_ascii=False)
    print(f"\nResults written to {out_path}", flush=True)

    print("\n=== SUMMARY ===")
    header = (
        f"{'conc':>4} {'step':>4} {'ok':>4} {'corr':>4} "
        f"{'TTFA med(ms)':>12} {'step med(ms)':>12} {'step p95(ms)':>12} {'RTF mean+-sd':>14}"
    )
    print(header)
    print("-" * len(header))
    for r in all_results:
        print(
            f"{r['concurrency']:>4} {r['step_size']:>4} {r['ok']:>4} {r['correctness_problems']:>4} "
            f"{r['ttfa_s']['median'] * 1e3:>12.1f} "
            f"{r['step_latency_ms']['median']:>12.1f} {r['step_latency_ms']['p95']:>12.1f} "
            f"{r['rtf']['mean']:>6.3f}+-{r['rtf']['stdev']:>5.3f}"
        )
        for t, stats in r.get("by_text", {}).items():
            print(
                f"      text {t[:24]!r}: e2e {stats['e2e_s']['mean'] * 1e3:8.1f}"
                f"+-{stats['e2e_s']['stdev'] * 1e3:6.1f} ms  (n={stats['e2e_s']['n']}, "
                f"audio {stats['audio_s']['mean']:.1f}s)"
            )

    if failed:
        sys.exit(1)


if __name__ == "__main__":
    main()