#!/usr/bin/env python3
"""端到端测试：音频/文本 → STT → 身份 → 记忆 → LLM → TTS → 音频。

**全程不需要麦克风、音箱、摄像头、ESP32。**
这是"服务器即助手"架构的直接好处：链路可以在任何有 CPU 的地方验证，
硬件到货后再接最后一段（真实音频进出）即可。

跑完会把合成音频写到 out/ 目录，可以人工试听整条链路的实际效果。

运行：
    cd <项目根>
    ~/venvs/ai_assistant/bin/python app/tests/test_engine_e2e.py
"""

from __future__ import annotations

import asyncio
import shutil
import sys
import tempfile
import time
import wave
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parent.parent.parent
sys.path.insert(0, str(ROOT))

from app.config import AppConfig  # noqa: E402
from app.engine import ConversationEngine, EventKind  # noqa: E402
from app.identity import GUEST_USER_ID, ManualIdentityProvider  # noqa: E402
from app.llm import FallbackProvider, create_provider  # noqa: E402
from app.memory import SQLiteMemoryStore  # noqa: E402
from app.voice import StreamingSTT, create_tts  # noqa: E402

OK, FAIL = [], []
OUT = ROOT / "out"


def check(name: str, cond: bool, detail: str = "") -> None:
    (OK if cond else FAIL).append(name)
    print(f"  [{'PASS' if cond else 'FAIL'}] {name}" + (f" — {detail}" if detail else ""))


def save_wav(path: Path, audio: np.ndarray, sr: int) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    pcm = (np.clip(audio, -1, 1) * 32767).astype(np.int16)
    with wave.open(str(path), "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(sr)
        w.writeframes(pcm.tobytes())


async def collect(engine: ConversationEngine, text: str) -> tuple[str, np.ndarray, list]:
    """跑一轮，返回 (回复文本, 拼接音频, 事件列表)。"""
    reply, audios, events = [], [], []
    async for ev in engine.respond_to_text(text):
        events.append(ev)
        if ev.kind == EventKind.RESPONSE_DELTA:
            reply.append(ev.text)
        elif ev.kind == EventKind.SPEAK and ev.audio is not None:
            audios.append(ev.audio)
    audio = np.concatenate(audios) if audios else np.zeros(0, dtype=np.float32)
    return "".join(reply).strip(), audio, events


async def main() -> int:
    cfg = AppConfig.load()
    tmp = Path(tempfile.mkdtemp(prefix="e2e_"))
    OUT.mkdir(parents=True, exist_ok=True)

    print("=== 0. 装配 ===")
    llm = FallbackProvider(create_provider(cfg.llm),
                           create_provider(cfg.llm_fallback))
    mem = SQLiteMemoryStore(tmp / "memory.db")
    await mem.open()
    tts = create_tts(cfg.voice)
    stt = StreamingSTT(cfg.voice)
    ident = ManualIdentityProvider()
    engine = ConversationEngine(cfg, llm, mem, ident, tts, stt)
    print(f"      LLM={cfg.llm.model}  备用={cfg.llm_fallback.model}")
    tts_desc = (f"{cfg.voice.tts_backend}/{cfg.voice.tts_speaker}"
                if cfg.voice.tts_backend == "remote" else f"local/{cfg.voice.tts_threads}线程")
    print(f"      STT={cfg.voice.stt_threads}线程  TTS={tts_desc}")
    check("四个组件装配成功", True)

    ok, detail = await llm.health()
    check("LLM 可用", ok, detail)
    if not ok:
        await mem.close()
        shutil.rmtree(tmp, ignore_errors=True)
        return 1

    await engine.start_session("e2e_sess", user_id="u_ylhan")

    # ---------------------------------------------------------- 1
    print("\n=== 1. 纯文本一轮（LLM → 分句 → TTS）===")
    reply, audio, events = await collect(engine, "现在几点了？随便说一句就行。")
    print(f"      回复: {reply[:60]!r}")
    m = [e for e in events if e.kind == EventKind.METRICS][0].metrics
    print(f"      {m.summary()}")
    check("有回复", bool(reply))
    check("有音频输出", len(audio) > 0, f"{len(audio)} 采样点")
    check("回复不含 Markdown",
          not any(c in reply for c in "*#`|[]"), reply[:40])
    # 首音延迟这里只做"没离谱"的兜底检查，不设严格阈值：
    # 本机 N5095 的 TTS RTF 是 1.86（实测），首音必然被 TTS 拖到秒级；
    # 服务器上同一模型是 0.09，首音在 150ms 量级。
    # 真正该卡指标的是部署机（服务器），所以那个断言放在服务器上的验收里。
    check("首音延迟未失控", m.first_audio_ms < 8000, f"{m.first_audio_ms:.0f}ms")
    save_wav(OUT / "e2e_1_text.wav", audio, tts.sample_rate)

    # ---------------------------------------------------------- 2
    print("\n=== 2. 多句回复按句播（每句独立音频）===")
    _, _, events2 = await collect(
        engine, "用三句话介绍一下你自己，每句短一点。")
    speaks = [e for e in events2 if e.kind == EventKind.SPEAK]
    print(f"      分 {len(speaks)} 段: {[e.text[:14] for e in speaks]}")
    check("按句切成了多段", len(speaks) >= 2, f"{len(speaks)} 段")
    audio2 = np.concatenate([e.audio for e in speaks]) if speaks else np.zeros(0, np.float32)
    save_wav(OUT / "e2e_2_multisentence.wav", audio2, tts.sample_rate)

    # ---------------------------------------------------------- 3
    print("\n=== 3. 记忆：显式记住 → 下一轮被检索到 ===")
    await collect(engine, "记住我喜欢喝美式咖啡，不加糖。")
    await asyncio.sleep(0.5)  # 等异步写入落库
    stored = await mem.list_memories("u_ylhan", include_candidates=True)
    print(f"      库中记忆 {len(stored)} 条: {[m.text[:20] for m in stored]}")
    check("记忆已落库", len(stored) >= 1, f"{len(stored)} 条")
    check("自动抽取的是候选状态",
          all(m.confirm_status.value == "candidate" for m in stored))

    # 手动确认一条，看它是否参与检索
    await mem.remember(type(stored[0])(
        owner_user_id="u_ylhan", kind=stored[0].kind, key="",
        value="用户喜欢喝美式咖啡，不加糖",
        confirm_status=__import__("app.memory", fromlist=["ConfirmStatus"]).ConfirmStatus.CONFIRMED))
    _, _, events3 = await collect(engine, "我想喝点东西，推荐一下？")
    used = [e for e in events3 if e.kind == EventKind.MEMORY_USED]
    n_used = len(used[0].memories) if used else 0
    print(f"      本轮用到记忆 {n_used} 条")
    check("确认的记忆参与检索", n_used >= 1, f"{n_used} 条")

    # ---------------------------------------------------------- 4
    print("\n=== 4. 身份隔离：换人后拿不到上一个人的记忆 ===")
    ident.switch("u_other", reason="测试切换")
    _, _, events4 = await collect(engine, "我想喝点什么？")
    used4 = [e for e in events4 if e.kind == EventKind.MEMORY_USED]
    n4 = len(used4[0].memories) if used4 else 0
    check("新身份没有旧记忆", n4 == 0, f"{n4} 条")
    reply4 = "".join(e.text for e in events4 if e.kind == EventKind.RESPONSE_DELTA)
    check("回复未泄露咖啡偏好", "美式" not in reply4 and "咖啡" not in reply4,
          reply4[:50])

    # ---------------------------------------------------------- 5
    print("\n=== 5. 访客身份不做个性化 ===")
    ident.reset()
    print(f"      当前身份: {ident.current().user_id} (confirmed={ident.current().confirmed})")
    _, _, events5 = await collect(engine, "我喜欢什么？")
    used5 = [e for e in events5 if e.kind == EventKind.MEMORY_USED]
    check("访客不检索私人记忆", len(used5) == 0 or not used5[0].memories)
    reply5 = "".join(e.text for e in events5 if e.kind == EventKind.RESPONSE_DELTA)
    check("访客回复不编造偏好", "美式" not in reply5, reply5[:50])

    # ---------------------------------------------------------- 6
    print("\n=== 6. 音频输入（真人录音 → STT → LLM → TTS）===")
    # 测试音频不在代码目录里，和模型放在一起（都是可重新下载的产物）。
    # 用模型目录反推，注意要 resolve()——服务器上模型目录是个软链，
    # 不解析的话 .parent 会指到 ~/ai-companion 而不是真实位置。
    wav_path = None
    for cand in (
        cfg.voice.model_dir.resolve().parent / "wavs" / "stt_1.wav",
        ROOT / "p0" / "stt_tts_bench" / "wavs" / "stt_1.wav",
    ):
        if cand.is_file():
            wav_path = cand
            break
    if wav_path is not None:
        with wave.open(str(wav_path)) as w:
            samples = np.frombuffer(w.readframes(w.getnframes()),
                                    dtype=np.int16).astype(np.float32) / 32768.0
        ident.switch("u_ylhan", reason="回到主人")
        t0 = time.perf_counter()
        transcript, audios, events6 = [], [], []
        async for ev in engine.respond_to_audio(samples):
            events6.append(ev)
            if ev.kind == EventKind.TRANSCRIPT:
                transcript.append(ev.text)
            elif ev.kind == EventKind.SPEAK and ev.audio is not None:
                audios.append(ev.audio)
        got = "".join(transcript)
        m6 = [e for e in events6 if e.kind == EventKind.METRICS]
        print(f"      STT: {got[:50]!r}")
        if m6 and m6[0].metrics:
            print(f"      {m6[0].metrics.summary()}")
        check("音频链路跑通（能出转写）", bool(got.strip()))
        if audios:
            save_wav(OUT / "e2e_6_fromaudio.wav", np.concatenate(audios), tts.sample_rate)
    else:
        check("音频链路跑通", False, f"测试音频不存在: {wav_path}")

    # ---------------------------------------------------------- 7
    print("\n=== 7. 打断：取消正在进行的回复 ===")
    cancelled = "未取消"
    chunks_seen = 0

    async def long_turn():
        # 必须把整个流消费完才谈得上"取消"。
        # 第一版写成 `async for _ in ...: return "已开始"`——
        # 第一个事件就 return，生成器随之关闭，根本没有正在进行的回复可取消，
        # 测试于是永远"失败"，但失败原因跟取消能力毫无关系。
        nonlocal chunks_seen
        async for _ev in engine.respond_to_text("请详细讲一讲中国历史，越长越好。"):
            chunks_seen += 1
        return "已完成"

    task = asyncio.create_task(long_turn())
    await asyncio.sleep(0.7)
    task.cancel()
    try:
        await task
    except asyncio.CancelledError:
        cancelled = "已取消"
    check("长回复可被取消", cancelled == "已取消" and chunks_seen > 0,
          f"{cancelled}，取消前已收到 {chunks_seen} 个事件")

    await engine.end_session()
    await mem.close()
    await llm.aclose()
    shutil.rmtree(tmp, ignore_errors=True)

    print(f"\n合成音频已写入 {OUT}/ ，可人工试听：")
    for f in sorted(OUT.glob("e2e_*.wav")):
        print(f"  {f.name}")

    print("\n" + "=" * 62)
    print(f"通过 {len(OK)} / {len(OK) + len(FAIL)}")
    for n in FAIL:
        print(f"  未通过: {n}")
    return 0 if not FAIL else 1


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