#!/usr/bin/env python3
"""最小 aiortc 客户端（实验用）。

**为什么需要它**：Pipecat 的官方客户端 SDK 覆盖浏览器 / 移动端 / C++ / ESP32，
唯独没有适配"跑 Linux 的 x86 盒子接 USB 麦克风阵列"——也就是 N5095 的形态。
所以客户端必须自己写，这个文件是雏形。

它同时是验证工具：不依赖浏览器就能测服务端。用法：

    # 发一段 wav，收回应答音频
    python client/test_client.py --server http://<服务器>:7860 \
        --wav p0/stt_tts_bench/wavs/stt_1.wav --out /tmp/reply.wav

    # 只做连接测试，不发音频
    python client/test_client.py --server http://<服务器>:7860 --connect-only

**它还不是什么**：不做 AEC、不做播客参考、不做打断（那些要真实声卡和硬件），
只用文件音频验证信令与媒体通路。
"""

from __future__ import annotations

import argparse
import asyncio
import fractions
import json
import logging
import sys
import time
import urllib.request
import wave
from pathlib import Path

import numpy as np
from aiortc import (
    RTCConfiguration,
    RTCPeerConnection,
    RTCSessionDescription,
)
from aiortc.mediastreams import AudioStreamTrack, MediaStreamError

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("client")


class SilentTrack(AudioStreamTrack):
    """静音音轨。只为触发 ICE 候选收集，不发真实音频。"""

    kind = "audio"

    def __init__(self, sample_rate: int = 16000):
        super().__init__()
        self.sample_rate = sample_rate
        self._n = 0

    async def recv(self):
        import fractions as _fr
        from av import AudioFrame as _AF
        await asyncio.sleep(0.02)
        pcm = np.zeros(self.sample_rate // 50, dtype=np.int16)
        f = _AF(format="s16", layout="mono", samples=len(pcm))
        f.planes[0].update(pcm.tobytes())
        f.sample_rate = self.sample_rate
        f.pts = self._n
        f.time_base = _fr.Fraction(1, self.sample_rate)
        self._n += len(pcm)
        return f


class FileAudioTrack(AudioStreamTrack):
    """把 wav 文件当作麦克风音源，按实时速率吐出来。

    实时速率很重要——如果一次性灌完，服务端的 VAD 会把它当成
    一段超长的连续语音，测出来的延迟也不真实。
    """

    kind = "audio"

    def __init__(self, wav_path: Path, *, sample_rate: int = 16000,
                 tail_silence_s: float = 1.2):
        super().__init__()
        with wave.open(str(wav_path)) as w:
            assert w.getnchannels() == 1, "需要单声道"
            self.src_rate = w.getframerate()
            pcm = np.frombuffer(w.readframes(w.getnframes()), dtype=np.int16)
        self.samples = pcm.astype(np.float32) / 32768.0
        self.sample_rate = sample_rate
        if self.src_rate != sample_rate:
            n = int(len(self.samples) * sample_rate / self.src_rate)
            self.samples = np.interp(np.linspace(0, len(self.samples) - 1, n),
                                     np.arange(len(self.samples)), self.samples)
        self._source_pos = 0
        self._timestamp_samples = 0
        self._tail_remaining = int(tail_silence_s * sample_rate)
        self.source_finished_at: float | None = None
        self._t0: float | None = None
        self._frame_samples = sample_rate // 100     # 10ms

    async def recv(self):
        """按 10ms 一块、以真实时间节奏产出。"""
        if self._t0 is None:
            self._t0 = time.perf_counter()
            self._source_pos = 0

        frame_pts = self._timestamp_samples
        if self._source_pos >= len(self.samples):
            if self.source_finished_at is None:
                self.source_finished_at = time.perf_counter()
            audio = np.zeros(self._frame_samples, dtype=np.float32)
            self._tail_remaining = max(0, self._tail_remaining - self._frame_samples)
        else:
            audio = self.samples[self._source_pos:self._source_pos + self._frame_samples]
            if len(audio) < self._frame_samples:
                audio = np.pad(audio, (0, self._frame_samples - len(audio)))
            self._source_pos += self._frame_samples

        # 即使源文件已经结束，静音帧也必须持续推进时钟和 PTS。
        self._timestamp_samples += self._frame_samples

        # 实时节奏
        target = self._t0 + self._timestamp_samples / self.sample_rate
        delay = target - time.perf_counter()
        if delay > 0:
            await asyncio.sleep(delay)

        pcm = (np.clip(audio, -1, 1) * 32767).astype(np.int16)
        from av import AudioFrame
        f = AudioFrame(format="s16", layout="mono", samples=len(pcm))
        f.planes[0].update(pcm.tobytes())
        f.sample_rate = self.sample_rate
        f.pts = int(frame_pts * 90000 / self.sample_rate)
        f.time_base = fractions.Fraction(1, 90000)
        return f

    @property
    def finished(self) -> bool:
        return self._source_pos >= len(self.samples) and self._tail_remaining <= 0


class ReplyRecorder:
    """收服务端音频，写 wav，并记录首音延迟。"""

    def __init__(self, out_path: Path, sample_rate: int = 24000,
                 speech_rms_threshold: float = 80.0):
        self.out_path = out_path
        self.sample_rate = sample_rate
        self.chunks: list[np.ndarray] = []
        self.first_non_silent_at: float | None = None
        self.speech_rms_threshold = speech_rms_threshold

    def feed(self, frame) -> None:
        arr = frame.to_ndarray()
        channels = len(frame.layout.channels)
        if frame.format.is_planar:
            raw = arr.astype(np.float32).mean(axis=0)
        else:
            raw = arr.astype(np.float32).reshape(-1, channels).mean(axis=1)
        # WebRTC/Opus 通常以 48kHz 交付，即使服务端原始输出是 24kHz。
        # WAV 必须记录接收帧的实际采样率，否则时长会被算错。
        if not self.chunks:
            self.sample_rate = frame.sample_rate
        elif frame.sample_rate != self.sample_rate:
            raise ValueError(
                f"回复音轨采样率发生变化: {self.sample_rate} -> {frame.sample_rate}")
        raw = raw.reshape(-1)
        rms = float(np.sqrt(np.mean(raw ** 2))) if raw.size else 0.0
        if self.first_non_silent_at is None and rms >= self.speech_rms_threshold:
            self.first_non_silent_at = time.perf_counter()
        self.chunks.append(raw / 32768.0)

    def save(self) -> float:
        if not self.chunks:
            return 0.0
        audio = np.concatenate(self.chunks)
        pcm = (np.clip(audio, -1, 1) * 32767).astype(np.int16)
        self.out_path.parent.mkdir(parents=True, exist_ok=True)
        with wave.open(str(self.out_path), "wb") as w:
            w.setnchannels(1)
            w.setsampwidth(2)
            w.setframerate(self.sample_rate)
            w.writeframes(pcm.tobytes())
        return len(audio) / self.sample_rate

    @property
    def has_speech(self) -> bool:
        return self.first_non_silent_at is not None

    def first_audio_ms_from(self, started_at: float | None) -> float | None:
        if started_at is None or self.first_non_silent_at is None:
            return None
        return (self.first_non_silent_at - started_at) * 1000


def rtvi_send_text(channel, text: str, *, audio_response: bool = True) -> None:
    """通过 data channel 发一条 RTVI send-text。

    信封格式来自 pipecat 的 RTVI 模型（实测校正过两次）：
      label='rtvi-ai'，**字段是 id 不是 msg_id**，type='send-text'，
      data.content（**是 content 不是 text**）。
      写错任何一个，服务端都会以 "Invalid RTVI transport message: Field required" 拒掉，
      而客户端看不到任何反馈——表现为"发了消息但助手毫无反应"。
    """
    import uuid
    msg = {
        "label": "rtvi-ai",
        "id": uuid.uuid4().hex,
        "type": "send-text",
        "data": {"content": text,
                 "options": {"audio_response": audio_response}},
    }
    channel.send(json.dumps(msg))


def post_offer(server: str, sdp: str, type_: str) -> dict:
    """把 offer 发给服务端，取回 answer。"""
    payload = json.dumps({"sdp": sdp, "type": type_, "pc_id": None,
                          "restart_pc": False, "request_data": None}).encode()
    req = urllib.request.Request(f"{server.rstrip('/')}/api/offer", data=payload,
                                 headers={"Content-Type": "application/json"},
                                 method="POST")
    with urllib.request.urlopen(req, timeout=30) as r:
        return json.loads(r.read().decode())


async def run(server: str, wav: Path | None, out: Path, *, connect_only: bool,
              timeout_s: float, texts: list[str] | None = None) -> int:
    pc = RTCPeerConnection(RTCConfiguration(iceServers=[]))
    recorder = ReplyRecorder(out)
    track = None
    got_audio = asyncio.Event()

    @pc.on("connectionstatechange")
    async def _state():
        log.info("连接状态: %s", pc.connectionState)
        if pc.connectionState in ("failed", "closed"):
            got_audio.set()

    @pc.on("track")
    def _track(t):
        log.info("收到服务端音轨: %s", t.kind)

        async def _pump():
            try:
                while True:
                    frame = await t.recv()
                    recorder.feed(frame)
            except (MediaStreamError, asyncio.CancelledError):
                pass
            finally:
                got_audio.set()

        asyncio.ensure_future(_pump())

    if wav and not connect_only:
        track = FileAudioTrack(wav)
        pc.addTrack(track)
        log.info("音源: %s（%d 采样点）", wav.name, len(track.samples))
    elif texts:
        # 文字模式：仍要一条音轨，否则 aiortc 不收集 ICE 候选
        pc.addTrack(SilentTrack())
        log.info("文字模式：将发送 %d 条文本", len(texts))
    else:
        # --connect-only 也要加一条（静音）音轨。
        # **不能省**：aiortc 的 offer 里没有媒体时完全不收集 ICE 候选，
        # 现象是"信令成功、SDP 交换正常、但 ICE 永远停在 new"。
        pc.addTrack(SilentTrack())
        log.info("连接测试模式：已加静音音轨（仅为触发 ICE 收集）")

    # **必须建 data channel 并定期发 "ping"。**
    # Pipecat 的 SmallWebRTC 用 data channel 上的消息来判断"客户端还活着"：
    #   if isinstance(message, str) and message.startswith("ping"):
    #       self._last_received_time = time.time()
    # 而 connection 的 connected 事件、进而 on_client_connected 回调、
    # 进而 StartFrame（音频推流的开关）都依赖它。
    # 不发 ping 的现象是：ICE 连通、音轨协商成功、但**服务端一帧音频都不推**、
    # 助手永远不说话，而且日志里只有一句 "Data channel not established within 10s"。
    channel = pc.createDataChannel("pipecat")
    dc_state = {"ok": False}

    @channel.on("open")
    def _dc_open():
        dc_state["ok"] = True
        log.info("data channel 已打开，开始发心跳")

    async def _keepalive():
        while pc.connectionState not in ("closed", "failed"):
            if dc_state["ok"]:
                try:
                    channel.send("ping")
                except Exception:  # noqa: BLE001
                    pass
            await asyncio.sleep(1.0)

    asyncio.ensure_future(_keepalive())

    offer = await pc.createOffer()
    await pc.setLocalDescription(offer)

    # 等 ICE 收集完成再发 offer——**这步不能省**。
    # 不等的话 SDP 里没有任何候选地址，服务端拿到的是一份"空地址"的 offer，
    # 信令看着成功，但 ICE 永远连不上（状态停在 new）。
    # 这叫非 trickle ICE：把候选都塞进 SDP 一次发完。
    # 真实产品会用 trickle（边收集边发候选），但那样需要额外的信令端点。
    for _ in range(100):
        if pc.iceGatheringState == "complete":
            break
        await asyncio.sleep(0.1)
    n_cand = pc.localDescription.sdp.count("candidate:")
    log.info("ICE 收集: %s，SDP 含 %d 个候选", pc.iceGatheringState, n_cand)
    if n_cand == 0:
        log.warning("SDP 里没有候选地址，ICE 很可能连不上")

    log.info("本地 SDP 已生成，发给服务端…")

    t_send = time.perf_counter()
    answer = post_offer(server, pc.localDescription.sdp, pc.localDescription.type)
    if answer.get("type") not in ("answer", None):
        log.error("服务端返回异常: %s", answer)
        return 1
    if answer.get("sdp"):
        await pc.setRemoteDescription(RTCSessionDescription(
            sdp=answer["sdp"], type=answer["type"]))
        log.info("远端 SDP 已设置（往返 %.0f ms）", (time.perf_counter() - t_send) * 1000)
    else:
        log.info("服务端未返回 SDP（可能是 pc_id 复用）: %s", answer)

    # 等 ICE 连通
    for _ in range(100):
        if pc.connectionState in ("connected", "completed"):
            break
        if pc.connectionState in ("failed", "closed"):
            log.error("连接失败")
            await pc.close()
            return 1
        await asyncio.sleep(0.1)
    log.info("✅ ICE 已连通（%s）", pc.connectionState)
    print(f"  ✅ 连接建立，状态 {pc.connectionState}")

    if texts:
        # 等 data channel 就绪（服务端靠它判断客户端存活）
        for _ in range(60):
            if dc_state["ok"]:
                break
            await asyncio.sleep(0.1)
        log.info("data channel: %s", "已就绪" if dc_state["ok"] else "未就绪")
        for t in texts:
            await asyncio.sleep(0.5)
            rtvi_send_text(channel, t)
            log.info("已发送文本: %r", t)
            # 等**非静音**音频出现。
            # 不能用"收到多少块"来判断——WebRTC 输出轨空闲时也在推静音，
            # 块数会立刻涨上去，导致客户端在服务端还在生成时就断开
            # （服务端日志表现为"本轮被取消"，客户端表现为"没反应"）。
            recorder.first_non_silent_at = None
            for _ in range(int(timeout_s * 10)):
                if recorder.has_speech:
                    break
                await asyncio.sleep(0.1)
            log.info("本轮 %s", "收到语音回复" if recorder.has_speech else "未收到语音")
            # 让它把话说完
            await asyncio.sleep(4)
        await asyncio.sleep(2)
    elif track:
        # 等音频放完 + 留出响应时间
        t0 = time.perf_counter()
        while not track.finished and time.perf_counter() - t0 < timeout_s:
            await asyncio.sleep(0.2)
        log.info("音频已发完，等待回复…")
        # 等首个非静音回复；SmallWebRTC 空闲时也会持续发送静音。
        t_wait = time.perf_counter()
        while time.perf_counter() - t_wait < timeout_s:
            if recorder.has_speech:
                break
            await asyncio.sleep(0.1)
        # 再收一会儿，让回复说完
        await asyncio.sleep(6)

    dur = recorder.save()
    await pc.close()
    got_audio.set()

    print()
    print("=== 结果 ===")
    print(f"  回复: {'✅ 收到非静音语音' if recorder.has_speech else '❌ 只有静音'}")
    if dur:
        print(f"  录制音轨: {dur:.2f} 秒 → {out}")
    fa = recorder.first_audio_ms_from(track.source_finished_at if track else None)
    if fa is not None:
        print(f"  首音延迟（源音频结束→首个非静音输出）: {fa:.0f} ms")
    return 0 if (connect_only or recorder.has_speech) else 1


def main() -> int:
    p = argparse.ArgumentParser()
    p.add_argument("--server", default="http://cadadmin-server.taild0f91d.ts.net:7860",
                   help="服务端地址。默认用 tailscale 域名——LAN 地址会被重新编址"
                        "（2026-09-15 实测从 192.168.0.x 变成 172.16.1.x）")
    p.add_argument("--wav", type=Path)
    p.add_argument("--out", type=Path, default=Path("/tmp/client_reply.wav"))
    p.add_argument("--connect-only", action="store_true")
    p.add_argument("--text", action="append", default=None,
                   help="通过 RTVI 发送文本（可重复），用于测 flows 文字路径")
    p.add_argument("--timeout", type=float, default=30.0)
    a = p.parse_args()
    if not a.connect_only and not a.wav and not a.text:
        print("需要 --wav / --text / --connect-only", file=sys.stderr)
        return 2
    return asyncio.run(run(a.server, a.wav, a.out,
                           connect_only=a.connect_only, timeout_s=a.timeout,
                           texts=a.text))


if __name__ == "__main__":
    sys.exit(main())