#!/usr/bin/env python3
"""Pipecat SmallWebRTC 服务端（实验）。

**目的**：验证 Pipecat 传输层与自有引擎能否协同工作，并拿到真实传输下的延迟数字。
不是产品代码——见文件末尾"这还不是什么"。

结构：

    ┌──────────────── Pipecat ────────────────┐
    │ transport.input() → EngineProcessor     │
    │                      ↓                  │
    │                   transport.output()    │
    └─────────────────────────────────────────┘
                          ↑
                  我们的 ConversationEngine
                  （STT→记忆→LLM→TTS 全在里面）

EngineProcessor 把整条管线"吞"进一个处理器——这是 A 方案的核心假设：
**Pipecat 提供传输与帧模型，会话语义归我们的引擎。** 这个实验就是在检验它成不成立。

用法：
    # 服务器上
    .venv/bin/python services/bot_server.py --host 0.0.0.0 --port 7860
    # 浏览器打开 http://<服务器IP>:7860/  点连接后说话
"""

from __future__ import annotations

import argparse
import asyncio
import logging
import os
import sys
import time
from collections import deque
from pathlib import Path

import numpy as np

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

from fastapi import FastAPI
from fastapi.responses import RedirectResponse
from pipecat.frames.frames import (
    CancelFrame,
    EndFrame,
    InputAudioRawFrame,
    InterruptionFrame,
    OutputAudioRawFrame,
    StartFrame,
    TextFrame,
)
from pipecat.pipeline.pipeline import Pipeline
from pipecat.pipeline.worker import PipelineParams, PipelineWorker
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
from pipecat.transports.base_transport import TransportParams
from pipecat.transports.smallwebrtc.request_handler import (
    SmallWebRTCRequest,
    SmallWebRTCRequestHandler,
)
from pipecat.transports.smallwebrtc.transport import SmallWebRTCTransport
from pipecat.utils.asyncio.task_manager import TaskManager
from pipecat.workers.base_worker import WorkerParams

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

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

# Pipecat 会在每个 worker 启动时预热 NLTK 分句器。固定到私有、可预装的目录，
# 避免使用权限过宽的 ~/nltk_data，也避免首个会话临时联网下载。
DEFAULT_NLTK_DATA = Path.home() / ".local/share/ai-companion/nltk_data"
os.environ.setdefault("NLTK_DATA", str(DEFAULT_NLTK_DATA))

#: 输入音频格式（WebRTC 侧协商后的格式）
IN_RATE = 16000
IN_CHANNELS = 1
#: 输出音频格式（Qwen3-TTS 输出 24kHz）
OUT_RATE = 24000

#: 管线参数。给 PipelineWorker 构造用。
PIPELINE_PARAMS = PipelineParams(
    audio_in_sample_rate=IN_RATE,
    audio_out_sample_rate=OUT_RATE,
    enable_metrics=True,
)

#: 注意：PipelineWorker.run() 要的是 WorkerParams（内含 task_manager），
#: **不是** PipelineParams。传错会依次报两个错，且都发生在后台任务里，
#: 主日志里什么都看不到：
#:   1. 不传 → "PipelineWorker.run() missing 1 required positional argument: 'params'"
#:   2. 传错 → "'PipelineParams' object has no attribute 'task_manager'"
#: 而且 TaskManager() 内部调用 asyncio.get_running_loop()，
#: **不能在模块级创建**（会报 "no running event loop"），必须在异步上下文里建。
#: 参考实现：pipecat/workers/runner.py


# ---------------------------------------------------------------- 能量 VAD

class SimpleVAD:
    """极简能量判定。

    **这是实验用的占位实现**，不是最终方案。方案 §5 把"说完判定"列为需要
    实测的部分，正式的 VAD 还没选（候选：Pipecat 的 Silero、
    sherpa-onnx VAD、或客户端 VAD + 服务端二次确认）。

    用它的唯一原因是：先把传输层跑通，VAD 的好坏不影响"Pipecat 能不能
    当传输层用"这个结论。
    """

    def __init__(self, *, threshold: float = 350.0,
                 silence_ms: int = 700, min_speech_ms: int = 300):
        self.threshold = threshold          # PCM16 的 RMS 阈值
        self.silence_ms = silence_ms
        self.min_speech_ms = min_speech_ms
        self._silence = 0.0
        self._speech = 0.0

    def reset(self) -> None:
        self._silence = 0.0
        self._speech = 0.0

    @property
    def active(self) -> bool:
        """是否已经进入一段候选语音。"""
        return self._speech > 0.0

    def feed(self, pcm16: np.ndarray, chunk_ms: float) -> str | None:
        """喂一块音频，返回 start / end / discard / None。"""
        if pcm16.size == 0:
            return None
        rms = float(np.sqrt(np.mean(pcm16.astype(np.float32) ** 2)))
        event = None
        if rms >= self.threshold:
            if self._speech == 0.0:
                event = "start"
            self._speech += chunk_ms
            self._silence = 0.0
        else:
            if self._speech > 0:
                self._silence += chunk_ms
                if self._silence >= self.silence_ms:
                    if self._speech >= self.min_speech_ms:
                        event = "end"
                    else:
                        event = "discard"
                    self.reset()
        return event


# ---------------------------------------------------------------- 引擎处理器

class EngineProcessor(FrameProcessor):
    """把 ConversationEngine 包成一个 Pipecat 处理器。

    A 方案的核心假设在这里被检验：**一个处理器吞掉整条会话逻辑，
    Pipecat 只负责音频进出和帧路由。**
    """

    def __init__(self, cfg: AppConfig, transport: SmallWebRTCTransport):
        super().__init__()
        self.cfg = cfg
        self.transport = transport
        self._vad = SimpleVAD()
        self._pre_roll: deque[bytes] = deque(maxlen=50)  # 约 1 秒
        self._buf = bytearray()          # 当前话语的 PCM16
        self._turn: asyncio.Task | None = None
        self._engine: ConversationEngine | None = None
        self._memory: SQLiteMemoryStore | None = None
        self._llm: FallbackProvider | None = None
        self._session = ""
        self._bot_speaking = False
        # 计时：用于拿真实延迟数字
        self._t_speech_end = 0.0
        self._t_first_audio = 0.0

    # ------------------------------------------------------------ 生命周期

    async def _ensure_engine(self) -> ConversationEngine:
        if self._engine is None:
            self._memory = SQLiteMemoryStore(
                Path.home() / ".local/share/ai_assistant/memory.db")
            await self._memory.open()
            self._llm = FallbackProvider(create_provider(self.cfg.llm),
                                         create_provider(self.cfg.llm_fallback))
            self._engine = ConversationEngine(
                self.cfg, self._llm, self._memory, ManualIdentityProvider(),
                create_tts(self.cfg.voice), StreamingSTT(self.cfg.voice))
            self._session = await self._engine.start_session(user_id="u_ylhan")
        return self._engine

    async def setup(self, setup):
        await super().setup(setup)
        log.info("会话处理器已就绪")

    async def process_frame(self, frame, direction: FrameDirection):
        await super().process_frame(frame, direction)

        if isinstance(frame, StartFrame):
            await self.push_frame(frame, direction)

        elif isinstance(frame, InputAudioRawFrame):
            await self._on_audio(frame)

        elif isinstance(frame, InterruptionFrame):
            # 客户端判定打断后发来的信号（或 Pipecat 内部产生）
            log.info("收到打断信号，取消当前生成")
            await self._cancel_turn()
            await self.push_frame(frame, direction)

        elif isinstance(frame, (EndFrame, CancelFrame)):
            await self._cancel_turn()
            await self.push_frame(frame, direction)

        else:
            await self.push_frame(frame, direction)

    # ------------------------------------------------------------ 音频处理

    async def _on_audio(self, frame: InputAudioRawFrame) -> None:
        pcm16 = np.frombuffer(frame.audio, dtype=np.int16)
        chunk_ms = len(pcm16) / IN_RATE * 1000.0

        # 头几帧打日志，用于确认音频真的到了处理器
        self._n_audio = getattr(self, "_n_audio", 0) + 1
        if self._n_audio <= 3:
            rms = float(np.sqrt(np.mean(pcm16.astype(np.float32) ** 2))) if pcm16.size else 0.0
            log.info("收到音频帧 #%d: %d 采样, %.0fms, RMS=%.0f",
                     self._n_audio, pcm16.size, chunk_ms, rms)

        # 助手正在说话时收到的音频 = 用户插话的候选。
        # 这里只记日志，真正的打断判定在客户端（见方案 §5 的边界决定）。
        if self._bot_speaking and float(np.sqrt(np.mean(pcm16.astype(np.float32) ** 2))) > self._vad.threshold:
            log.info("检测到用户插话（助手正在播放）")

        was_active = self._vad.active
        self._pre_roll.append(frame.audio)
        event = self._vad.feed(pcm16, chunk_ms)

        if event == "start":
            # 带上说话前约 1 秒，避免 VAD 切掉首字。
            self._buf = bytearray().join(self._pre_roll)
        elif was_active:
            self._buf.extend(frame.audio)

        # 防止噪声或异常客户端让常驻进程无限积累音频。
        max_bytes = IN_RATE * 2 * 60
        if len(self._buf) > max_bytes:
            log.warning("单次话语超过 60 秒，强制截断处理")
            event = "end"
            self._vad.reset()

        if event == "discard":
            self._buf.clear()
            return

        if event == "end":
            audio = np.frombuffer(bytes(self._buf), dtype=np.int16).astype(np.float32) / 32768.0
            self._buf.clear()
            if len(audio) > IN_RATE * 0.3:      # 太短的不处理
                self._t_speech_end = time.perf_counter()
                self._t_first_audio = 0.0
                await self._cancel_turn()
                self._turn = self.create_task(self._run_turn(audio), "conversation-turn")

    async def _cancel_turn(self) -> None:
        if self._turn and not self._turn.done():
            await self.cancel_task(self._turn, timeout=5.0)
        self._turn = None

    # ------------------------------------------------------------ 跑一轮

    async def _run_turn(self, audio: np.ndarray) -> None:
        try:
            engine = await self._ensure_engine()
            async for ev in engine.respond_to_audio(audio, session_id=self._session):
                if ev.kind == EventKind.RESPONSE_DELTA and ev.text:
                    await self.push_frame(TextFrame(text=ev.text))
                elif ev.kind == EventKind.SPEAK and ev.audio is not None and len(ev.audio):
                    if self._t_first_audio == 0.0:
                        self._t_first_audio = time.perf_counter()
                        log.info("首音延迟（说完→首段音频）: %.0f ms",
                                 (self._t_first_audio - self._t_speech_end) * 1000)
                    await self._push_audio(ev.audio, ev.sample_rate)
                elif ev.kind == EventKind.METRICS and ev.metrics:
                    log.info("本轮指标: %s", ev.metrics.summary())
        except asyncio.CancelledError:
            log.info("本轮被取消")
            raise
        except Exception as e:  # noqa: BLE001
            log.exception("跑轮失败: %s", e)

    async def _push_audio(self, samples: np.ndarray, sample_rate: int) -> None:
        """把 float32 音频推给传输层。必要时重采样到 OUT_RATE。"""
        if sample_rate != OUT_RATE:
            # 简单线性重采样。实验够用；正式实现应该用 soxr/sox（已装）。
            n = int(len(samples) * OUT_RATE / sample_rate)
            samples = np.interp(np.linspace(0, len(samples) - 1, n),
                                np.arange(len(samples)), samples).astype(np.float32)
        pcm = (np.clip(samples, -1, 1) * 32767).astype(np.int16)
        self._bot_speaking = True
        # 按 20ms 一块推，模拟渐进播放
        step = OUT_RATE // 50
        for i in range(0, len(pcm), step):
            chunk = pcm[i:i + step]
            await self.push_frame(OutputAudioRawFrame(
                audio=chunk.tobytes(), sample_rate=OUT_RATE, num_channels=1))
        self._bot_speaking = False

    async def cleanup(self):
        await self._cancel_turn()
        if self._engine is not None:
            await self._engine.end_session()
        if self._llm is not None:
            await self._llm.aclose()
        if self._memory is not None:
            await self._memory.close()
        self._engine = None
        self._llm = None
        self._memory = None
        await super().cleanup()


# ---------------------------------------------------------------- HTTP 服务

app = FastAPI(title="AI Companion · Pipecat experimental server")
handler = SmallWebRTCRequestHandler()
_cfg: AppConfig | None = None
_tasks: list[PipelineWorker] = []
_bg_tasks: set[asyncio.Task] = set()   # 持有后台 task 引用，避免被 GC


@app.post("/api/offer")
async def offer(request: SmallWebRTCRequest):
    """WebRTC 信令：客户端发 offer，我们回 answer。"""

    async def on_connection(connection):
        transport = SmallWebRTCTransport(
            webrtc_connection=connection,
            params=TransportParams(
                audio_in_enabled=True,
                audio_in_sample_rate=IN_RATE,
                audio_in_channels=IN_CHANNELS,
                audio_out_enabled=True,
                audio_out_sample_rate=OUT_RATE,
                audio_out_channels=1,
            ),
        )
        engine_proc = EngineProcessor(_cfg, transport)
        pipeline = Pipeline([transport.input(), engine_proc, transport.output()])
        # 注意：这个版本里 PipelineParams 没有 allow_interruptions——
        # 打断在 1.10 已移到 pipecat/turns/ 框架。实验阶段我们不用它的
        # 轮次框架（方案 §5 的边界决定是"打断由客户端判定"），
        # 所以这里只开指标。
        # 注意：PipelineWorker 的音频参数传构造函数；run() 另收
        # WorkerParams（主要提供 task_manager）。漏传 run 参数会报：
        #   "PipelineWorker.run() missing 1 required positional argument: 'params'"
        # 而这个异常发生在后台任务里，会被静默吞掉——
        # 现象是"连接正常但什么反应都没有"，而且日志里一条错误都看不到。
        # 所以下面的 _run_task 里必须 log.exception，不能让异常无声消失。
        # 保留 RTVI：网页 prebuilt 客户端用它显示状态、转写和指标；
        # N5095 的精简客户端可以忽略不认识的 data-channel 消息。
        task = PipelineWorker(
            pipeline,
            params=PIPELINE_PARAMS,
            enable_rtvi=True,
            idle_timeout_secs=None,
        )
        _tasks.append(task)
        log.info("新连接已建立，管线启动")

        @task.event_handler("on_pipeline_started")
        async def _on_pipeline_started(worker, frame):
            log.info("管线已就绪")

        @task.event_handler("on_pipeline_timeout")
        async def _on_pipeline_timeout(worker, frame):
            log.error("管线等待 %s 超时", type(frame).__name__)

        @task.event_handler("on_pipeline_error")
        async def _on_pipeline_error(worker, frame):
            log.error("管线错误: %s", frame)

        @transport.event_handler("on_client_connected")
        async def _on_conn(t, client):
            # PipelineWorker.run() 会自动发送且只应发送一次 StartFrame。
            log.info("客户端已连接")

        @transport.event_handler("on_client_disconnected")
        async def _on_disc(t, client):
            log.info("客户端断开")
            await task.cancel()

        # **必须作为后台任务跑，不能 await。**
        # 这个回调是被 handle_web_request 调用的，而后者又在 HTTP 请求里被 await。
        # 如果这里 `await task.run()`，就会在信令请求里阻塞整个会话的生命周期，
        # 管线初始化的时序会歪掉——实测表现为 setup()/connect() 不被调用，
        # 于是 on_client_connected 永远不触发、StartFrame 不发、一帧音频都收不到。
        # 官方 runner 用的是 background_tasks.add_task(bot, args)，同样的道理。
        async def _run_task():
            log.info("PipelineWorker 开始运行")
            try:
                # TaskManager 必须在这里建——模块级没有运行中的事件循环
                await task.run(WorkerParams(task_manager=TaskManager()))
            except Exception as e:  # noqa: BLE001
                log.exception("PipelineWorker 异常退出: %s", e)
            log.info("PipelineWorker 已结束")

        bg = asyncio.create_task(_run_task())
        _bg_tasks.add(bg)
        bg.add_done_callback(_bg_tasks.discard)

    answer = await handler.handle_web_request(request, on_connection)
    if answer is None:
        return {"status": "ok"}
    return answer


@app.post("/start")
async def start_web_client():
    """给 Pipecat Playground 返回本服务的 WebRTC 信令入口。

    prebuilt 网页客户端会先 POST /start，再用响应里的 endpoint 发 SDP offer；
    自写 aiortc/N5095 客户端则可以继续直接调用 /api/offer。
    """
    return {
        "webrtcRequestParams": {
            "endpoint": "/api/offer",
        },
    }


@app.get("/health")
async def health():
    return {"status": "ok", "connections": len([t for t in _tasks if not t.has_finished()])}


@app.get("/")
async def root():
    return RedirectResponse("/client/")


def mount_client() -> bool:
    """挂载 prebuilt 网页客户端（用于测试）。"""
    try:
        from pipecat_ai_small_webrtc_prebuilt.frontend import SmallWebRTCPrebuiltUI
        app.mount("/client", SmallWebRTCPrebuiltUI, name="client")
        return True
    except Exception as e:  # noqa: BLE001
        log.warning("prebuilt 客户端不可用（不影响服务端）: %s", e)
        return False


def check_runtime_data() -> None:
    """启动前验证运行时数据，避免首个会话里偷偷联网下载。"""
    import nltk

    try:
        nltk.data.find("tokenizers/punkt_tab")
    except LookupError as e:
        raise RuntimeError(
            "缺少 NLTK punkt_tab。请在部署阶段运行 "
            "`python -m nltk.downloader punkt_tab`，并设置 NLTK_DATA。"
        ) from e


def main() -> int:
    global _cfg
    p = argparse.ArgumentParser()
    p.add_argument("--host", default="0.0.0.0")
    p.add_argument("--port", type=int, default=7860)
    a = p.parse_args()

    check_runtime_data()
    _cfg = AppConfig.load()
    log.info("LLM=%s  TTS=%s/%s  STT=%d线程",
             _cfg.llm.model, _cfg.voice.tts_backend, _cfg.voice.tts_speaker,
             _cfg.voice.stt_threads)
    mounted = mount_client()
    log.info("网页客户端: %s", "已挂载 /client/" if mounted else "未挂载")

    import uvicorn
    uvicorn.run(app, host=a.host, port=a.port, log_level="warning")
    return 0


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