#!/usr/bin/env python3
"""Phase 1 spike：Pipecat Flows 与当前系统的兼容性验证。

**这一步只做验证，不改生产入口。** 交接文档要求：未经 spike 结论，
不进入 Phase 3 的生产集成。

验证 7 项（编号对应交接文档 §九 Phase 1）：

  1. 找到与 pipecat-ai==1.10.0 兼容的 flows 版本并精确 pin   ← 已在包管理中确认
  2. DeepSeek OpenAI-compatible 接口能在 Flows 中流式输出中文
  3. 流式 function-call arguments 能被正确聚合一次，不重复拼接
  4. 工具结果回灌后能生成最终回答
  5. 任务取消能终止 LLM 请求和后续工具执行
  6. context aggregator 不产生第二条 system message
  7. 能获得 token 级文本事件，满足网页显示和分句 TTS

用法：
    ~/venvs/ai_assistant/bin/python p1/flows_spike.py
    ~/venvs/ai_assistant/bin/python p1/flows_spike.py --only 3,4
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import sys
import time
from pathlib import Path

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

# 载入密钥（与本项目的其它脚本一致）
_env_file = Path.home() / ".config" / "ai_assistant" / "env"
if _env_file.is_file():
    for line in _env_file.read_text().splitlines():
        line = line.strip()
        if line and not line.startswith("#") and "=" in line:
            k, _, v = line.partition("=")
            os.environ.setdefault(k.strip(), v.strip())

from pipecat.frames.frames import (
    EndFrame,
    LLMFullResponseEndFrame,
    LLMFullResponseStartFrame,
    LLMMessagesAppendFrame,
    TextFrame,
)
from pipecat.pipeline.pipeline import Pipeline
from pipecat.pipeline.runner import PipelineRunner
from pipecat.pipeline.worker import PipelineParams, PipelineWorker
from pipecat.processors.aggregators.llm_context import LLMContext
from pipecat.processors.aggregators.llm_response_universal import LLMContextAggregatorPair
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
from pipecat.services.deepseek.llm import DeepSeekLLMService
# Flows 是 Pipecat 1.10 内置的（pipecat.flows）。
# 独立的 pipecat-ai-flows 包**不需要**，装了 Pipecat 会报
# '_warn_if_standalone_flows_installed' 错误提示卸载。
from pipecat.flows import FlowManager, NodeConfig

class DebugDeepSeek(DeepSeekLLMService):
    """抓实际发给 DeepSeek 的请求体。

    Flows 把 role_message 通过 LLMUpdateSettingsFrame 下发成
    system_instruction，由 adapter 拼进 messages；节点里的 task_messages
    又是另一条 system 消息。到底发出去几条、顺序如何，只有看请求体才知道。
    """

    last_request: dict | None = None

    def build_chat_completion_params(self, params):
        out = super().build_chat_completion_params(params)
        DebugDeepSeek.last_request = dict(out)
        return out

    @staticmethod
    def dump():
        r = DebugDeepSeek.last_request or {}
        msgs = r.get("messages") or []
        print("      ── 实际请求 ──")
        print(f"      model={r.get('model')}  tools={len(r.get('tools') or [])}")
        for m in msgs:
            role = m.get("role")
            c = m.get("content")
            c = (c if isinstance(c, str) else json.dumps(c, ensure_ascii=False)) or ""
            print(f"        [{role}] {c[:78]!r}")
        for t in (r.get("tools") or []):
            fn = (t.get("function") or {})
            print(f"        tool: {fn.get('name')}  req={fn.get('parameters', {}).get('required')}")


if os.environ.get("SPIKE_DEBUG"):
    import logging
    logging.basicConfig(level=logging.DEBUG,
                        format="%(asctime)s %(levelname)s %(name)s: %(message)s")

MODEL = os.environ.get("DEEPSEEK_MODEL", "deepseek-flash")
BASE_URL = os.environ.get("DEEPSEEK_BASE_URL", "https://api.deepseek.com/v1")

# NLTK punkt_tab：Pipecat 启动预热分句器时会用它，缺了会在 setup 阶段临时联网
# 下载、20 秒超时，然后整个管线卡住——现象是"帧一个都不来"，极难排查。
# 服务器上装在这里，本地开发机也要有。fail-fast 而不是让它去下载。
_NLTK_DIR = Path.home() / ".local" / "share" / "ai-companion" / "nltk_data"
os.environ.setdefault("NLTK_DATA", str(_NLTK_DIR))
if not (_NLTK_DIR / "tokenizers" / "punkt_tab").is_dir():
    sys.exit(
        f"缺少 NLTK punkt_tab：{_NLTK_DIR / 'tokenizers' / 'punkt_tab'}\n"
        f"  从服务器拷贝：\n"
        f"    rsync -az yonglin@cadadmin-server.taild0f91d.ts.net:"
        f"~/.local/share/ai-companion/nltk_data/tokenizers/punkt_tab "
        f"{_NLTK_DIR}/tokenizers/")


# ---------------------------------------------------------------- 观测处理器

class Observer(FrameProcessor):
    """被动观察管线里流过的帧，用于 spike 取证。"""

    def __init__(self, label: str):
        super().__init__()
        self.label = label
        self.texts: list[tuple[float, str]] = []
        self.starts = 0
        self.ends = 0
        self.other: list[str] = []
        self._t0 = time.perf_counter()

    async def process_frame(self, frame, direction: FrameDirection):
        await super().process_frame(frame, direction)
        if isinstance(frame, TextFrame) and frame.text:
            self.texts.append((time.perf_counter() - self._t0, frame.text))
        elif isinstance(frame, LLMFullResponseStartFrame):
            self.starts += 1
        elif isinstance(frame, LLMFullResponseEndFrame):
            self.ends += 1
        else:
            n = type(frame).__name__
            if n not in self.other:
                self.other.append(n)
        await self.push_frame(frame, direction)

    @property
    def full_text(self) -> str:
        return "".join(t for _, t in self.texts)


# ---------------------------------------------------------------- 搭建

def build(task_holder: dict, *, on_tool=None):
    """建一个最小管线：aggregator → llm，外加观察器。"""
    context = LLMContext()
    aggregator = LLMContextAggregatorPair(context)
    llm = DebugDeepSeek(
        api_key=os.environ["DEEPSEEK_API_KEY"],
        base_url=BASE_URL,
        settings=DebugDeepSeek.Settings(
            model=MODEL,
            # 关闭思考：Pipecat 的 DeepSeek service 默认就是 disabled，
            # 这里显式写出来，避免将来默认值变化后悄悄改变行为。
            thinking=DebugDeepSeek.ThinkingConfig(type="disabled"),
            max_tokens=200,
        ),
    )
    obs = Observer("llm-output")
    pipeline = Pipeline([
        aggregator.user(),
        llm,
        obs,
        aggregator.assistant(),
    ])
    # 注意：Pipecat 1.3.0 起 PipelineTask 已废弃，Flows 1.3.0 也要求
    # 传 PipelineWorker。用 PipelineTask 会报 "FlowManager requires a 'worker'"。
    worker = PipelineWorker(pipeline, params=PipelineParams(
        enable_metrics=True, cancel_on_idle_timeout=False))
    task_holder["task"] = worker
    task_holder["context"] = context
    task_holder["llm"] = llm
    task_holder["obs"] = obs
    return FlowManager(worker=worker, llm=llm, context_aggregator=aggregator)


async def run_flow(flow: FlowManager, task_holder: dict, nodes: NodeConfig,
                   user_text: str, *, timeout_s: float = 40.0) -> dict:
    """跑一轮：初始化 flow → 注入用户消息 → 收完回复。"""
    # **必须 await。** FlowManager.initialize 是协程函数，但它的签名显示
    # `-> None`，看不出是 async；漏掉 await 不报错，只是整个 flow 从不初始化——
    # 表现为 role_message/task_messages/工具 全部没生效，而 LLM 照样会回复
    # （因为 run_llm=True 直接触发了一次推理），极容易误判成"框架不支持这些特性"。
    await flow.initialize(nodes)
    task = task_holder["task"]
    obs = task_holder["obs"]

    # **必须用 PipelineRunner 跑，不能直接 worker.run()。**
    # 直接 run() 时 runner 没挂上，Flows 的工具执行会失败：
    #   "Worker 'PipelineWorker#0': runner is not set; call attach() first."
    # 症状是工具被调用但 handler 从不执行，LLM 收到一句通用的
    # "The function `X` failed and returned no result."，很难定位。
    runner_task = asyncio.create_task(PipelineRunner(handle_sigint=False).run(task))
    await asyncio.sleep(0.8)

    # **必须清空再发用户消息。**
    # Flows 的 NodeConfig.respond_immediately 默认 True——initialize() 会立刻
    # 触发一轮 LLM 回复（针对 task_messages，不是用户问题）。不清空的话，
    # 观察到的是那轮自动回复，会把结论归因错。
    # 本 spike 的节点都显式设 respond_immediately=False，让回复只由用户消息触发。
    obs.texts.clear()
    obs.ends = 0
    await task.queue_frames([LLMMessagesAppendFrame(
        messages=[{"role": "user", "content": user_text}], run_llm=True)])

    # 等到"输出稳定"而不是"第一个响应结束"。
    # 带工具的一轮会有两次 LLM 推理（先出 tool_calls，回灌后再出回答），
    # 只看 obs.ends>0 会在第一轮就退出，把工具回灌后的最终回答漏掉——
    # 那样会误判成"回灌没效果"。
    deadline = time.perf_counter() + timeout_s
    last_count = -1
    last_change = time.perf_counter()
    while time.perf_counter() < deadline:
        n = len(obs.texts)
        if n != last_count:
            last_count = n
            last_change = time.perf_counter()
        elif n > 0 and (time.perf_counter() - last_change) > 2.5:
            # 已有输出，且 2.5 秒没有新块 → 认为说完
            break
        await asyncio.sleep(0.05)
    await asyncio.sleep(0.3)

    result = {
        "text": obs.full_text,
        "text_chunks": len(obs.texts),
        "responses": obs.ends,
        "context_messages": list(task_holder["context"].messages),
    }
    await task.queue_frames([EndFrame()])
    try:
        await asyncio.wait_for(runner_task, timeout=5)
    except (asyncio.TimeoutError, asyncio.CancelledError):
        runner_task.cancel()
    return result


# ---------------------------------------------------------------- 各验证项

def check(name: str, ok: bool, detail: str = "") -> dict:
    print(f"  [{'PASS' if ok else 'FAIL'}] {name}" + (f" — {detail}" if detail else ""))
    return {"name": name, "ok": ok, "detail": detail}


async def item2_chinese() -> list[dict]:
    print("\n=== 2. DeepSeek 在 Flows 中流式输出中文 ===")
    holder: dict = {}
    flow = build(holder)
    node: NodeConfig = {
        "name": "conversation",
        "role_message": "你是桌面陪伴助手，回复简短自然，适合朗读，不超过两句话，不要用 Markdown。",
        "task_messages": [{"role": "system", "content": "当前任务：普通陪伴对话。"}],
        "respond_immediately": False,
    }
    r = await run_flow(flow, holder, node, "用一句话介绍一下你自己。")
    text = r["text"]
    DebugDeepSeek.dump()
    print(f"      回复: {text[:60]!r}")
    print(f"      分块数: {r['text_chunks']}")
    out = [
        check("有中文回复", bool(text.strip()) and any("一" <= c <= "鿿" for c in text),
              f"{len(text)} 字"),
        check("回复不含 Markdown", not any(c in text for c in "*#`|"), text[:40]),
        check("7. token 级文本事件（>1 块）", r["text_chunks"] > 1, f"{r['text_chunks']} 块"),
    ]
    return out


async def item3_4_tools() -> list[dict]:
    print("\n=== 3+4. 流式工具参数聚合 / 工具结果回灌 ===")
    calls: list[dict] = []

    handler_errors: list[str] = []

    async def get_time(args, flow_manager):
        # 把异常打出来：Flows/Pipecat 把它变成给 LLM 的通用错误串，
        # 真正的 traceback 只在 push_error 里，默认看不到。
        try:
            calls.append(dict(args))
            return {"time": "下午三点半"}
        except Exception as e:  # noqa: BLE001
            import traceback
            handler_errors.append(traceback.format_exc())
            print(f"      ⚠️ handler 异常: {type(e).__name__}: {e}")
            raise

    from pipecat.flows import FlowsFunctionSchema
    holder: dict = {}
    flow = build(holder)
    node: NodeConfig = {
        "name": "conversation",
        "role_message": "你是助手。需要知道时间时必须调用 get_time 工具，不要猜。",
        "task_messages": [{"role": "system", "content": "当前任务：回答用户问题。"}],
        "respond_immediately": False,
        "functions": [FlowsFunctionSchema(
            name="get_time",
            description="获取当前时间",
            properties={"city": {"type": "string", "description": "城市名"}},
            required=["city"],
            handler=get_time,
        )],
    }
    r = await run_flow(flow, holder, node, "北京现在几点了？")
    text = r["text"]
    DebugDeepSeek.dump()
    print(f"      工具调用参数: {calls}")
    print(f"      最终回复: {text[:60]!r}")
    if handler_errors:
        print("      异常栈:")
        print("      " + handler_errors[0].replace("\n", "\n      ")[:900])
    out = []
    out.append(check("3. 工具被调用", len(calls) >= 1, f"{len(calls)} 次"))
    if calls:
        # 参数必须是完整合法的 JSON，且 city 只出现一次（防重复拼接）
        raw = json.dumps(calls[0], ensure_ascii=False)
        out.append(check("3. 参数聚合完整（无重复拼接）",
                         calls[0].get("city") == "北京", raw))
    out.append(check("4. 工具结果回灌后有最终回答", bool(text.strip()), text[:40]))
    return out


async def item5_cancel() -> list[dict]:
    print("\n=== 5. 取消能终止 LLM 请求 ===")
    holder: dict = {}
    flow = build(holder)
    node: NodeConfig = {
        "name": "conversation",
        "role_message": "你是助手。",
        "task_messages": [{"role": "system", "content": "用户要求长文时必须完整写出。"}],
        "respond_immediately": False,
    }
    await flow.initialize(node)
    task = holder["task"]
    obs = holder["obs"]
    rt = asyncio.create_task(PipelineRunner(handle_sigint=False).run(task))
    await asyncio.sleep(0.5)

    await task.queue_frames([LLMMessagesAppendFrame(
        messages=[{"role": "user", "content": "请详细写一篇 800 字的长文，主题是窗外的雨。"}],
        run_llm=True)])
    # 等它开始输出
    for _ in range(200):
        if obs.texts:
            break
        await asyncio.sleep(0.05)
    got_before = len(obs.texts)
    t_cancel = time.perf_counter()
    # 用 worker.cancel()，不是 EndFrame。
    # EndFrame 是"优雅停止"——会把已排队的帧处理完，正在进行的 LLM 流不受影响。
    # 取消要求的是"立刻终止 LLM 请求和后续工具执行"，必须走 cancel。
    await task.cancel(reason="spike: 用户打断")
    await asyncio.sleep(1.5)
    got_after = len(obs.texts)
    stopped_in = time.perf_counter() - t_cancel
    try:
        await asyncio.wait_for(rt, timeout=5)
    except (asyncio.TimeoutError, asyncio.CancelledError):
        rt.cancel()
    print(f"      取消前 {got_before} 块 → 取消后 1.5s 内 {got_after} 块")
    return [check("5. 取消后停止输出", got_after - got_before <= 3,
                  f"又收到 {got_after - got_before} 块，{stopped_in:.1f}s 内停")]


async def item6_system_message() -> list[dict]:
    print("\n=== 6. 不产生第二条 system message ===")
    holder: dict = {}
    flow = build(holder)
    node: NodeConfig = {
        "name": "conversation",
        "role_message": "你是桌面陪伴助手。",
        "task_messages": [{"role": "system", "content": "当前任务：陪伴对话。"}],
        "respond_immediately": False,
    }
    r = await run_flow(flow, holder, node, "你好。")
    msgs = r["context_messages"]
    sys_msgs = [m for m in msgs if m.get("role") == "system"]
    roles = [m.get("role") for m in msgs]
    print(f"      context 消息角色序列: {roles}")
    for m in sys_msgs:
        c = (m.get("content") or "")
        print(f"        system: {str(c)[:70]!r}")
    return [
        check("6. system 消息不超过 1 条", len(sys_msgs) <= 1,
              f"{len(sys_msgs)} 条"),
        check("6. 有 user 与 assistant 轮次",
              "user" in roles and "assistant" in roles, str(roles)),
    ]


# ---------------------------------------------------------------- 主流程

ITEMS = {
    "2": ("DeepSeek 中文流式 + token 级事件", item2_chinese),
    "3": ("工具调用与回灌", item3_4_tools),
    "5": ("取消", item5_cancel),
    "6": ("system message", item6_system_message),
}


async def main() -> int:
    p = argparse.ArgumentParser()
    p.add_argument("--only", default="2,3,5,6", help="要跑的项，逗号分隔")
    a = p.parse_args()

    print("=" * 64)
    print("Phase 1 spike · Pipecat Flows 兼容性验证")
    import pipecat
    print(f"  pipecat-ai  {pipecat.__version__}（Flows 内置：pipecat.flows）")
    print(f"  LLM: {MODEL} @ {BASE_URL}")
    print("=" * 64)

    all_results: list[dict] = []
    for key in [k.strip() for k in a.only.split(",") if k.strip()]:
        if key not in ITEMS:
            print(f"\n  （跳过未实现的项 {key}）")
            continue
        _, fn = ITEMS[key]
        try:
            all_results.extend(await fn())
        except Exception as e:  # noqa: BLE001
            import traceback
            traceback.print_exc()
            all_results.append(check(f"{key}. 执行", False, f"{type(e).__name__}: {e}"))

    print("\n" + "=" * 64)
    ok = sum(1 for r in all_results if r["ok"])
    print(f"汇总: {ok} / {len(all_results)}")
    for r in all_results:
        if not r["ok"]:
            print(f"  未通过: {r['name']} — {r['detail']}")
    return 0 if ok == len(all_results) else 1


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