ai-video-fullstack/frontend/src/hooks/use-voice-preview.ts

"use client";

/**
 * 语音预览。
 *
 * WebRTC 生命周期交给 Pipecat 官方 SmallWebRTCTransport：它负责固定的
 * audio/video transceiver、trackStatus、keepalive、重协商和 ICE 恢复。
 * 本文件只把平台已有的业务消息映射到 React 状态。
 */

import { useCallback, useEffect, useRef, useState } from "react";
import type { PipecatClientOptions, RTVIMessage } from "@pipecat-ai/client-js";
import { SmallWebRTCTransport } from "@pipecat-ai/small-webrtc-transport";

import { API_BASE, webrtcApi } from "@/lib/api";

export type VoicePreviewStatus = "idle" | "connecting" | "connected" | "failed";
export type NetworkQuality = "unknown" | "good" | "fair" | "poor";

const NETWORK_SAMPLE_INTERVAL_MS = 2_000;

type NetworkStatsSample = {
  packetsSent: number;
  packetsLost: number;
};

type NetworkMetrics = NetworkStatsSample & {
  roundTripTime: number | null;
  availableOutgoingBitrate: number | null;
};

type ConnectOptions = {
  visionEnabled?: boolean;
  videoStream?: MediaStream | null;
  dynamicVariables?: Record<string, string | number | boolean>;
};

export type ChatMessage = {
  id: string;
  role: "user" | "assistant";
  content: string;
  timestamp: string;
  sequence: number;
  turnId?: string;
  streaming?: boolean;
};

type AppMessage = Record<string, unknown> & { type?: string };

class AppSmallWebRTCTransport extends SmallWebRTCTransport {
  onAppMessage?: (message: AppMessage) => void;

  override handleMessage(raw: string): void {
    try {
      const message = JSON.parse(raw) as AppMessage;
      if (message.type === "signalling" || message.label === "rtvi-ai") {
        super.handleMessage(raw);
      } else {
        this.onAppMessage?.(message);
      }
    } catch {
      super.handleMessage(raw);
    }
  }

  sendAppMessage(message: AppMessage): void {
    super.sendMessage(message as unknown as RTVIMessage);
  }

  async connectionStats(): Promise<RTCStatsReport | null> {
    const peerConnection = (
      this as unknown as { pc: RTCPeerConnection | null }
    ).pc;
    return peerConnection ? peerConnection.getStats() : null;
  }
}

function errorMessage(error: unknown, fallback: string): string {
  if (error instanceof Error && error.message) return error.message;
  return fallback;
}

function messageOrder(message: ChatMessage): number {
  const timestamp = Date.parse(message.timestamp);
  return Number.isNaN(timestamp) ? Number.MAX_SAFE_INTEGER : timestamp;
}

function sortMessages(messages: ChatMessage[]): ChatMessage[] {
  return messages.sort(
    (a, b) => messageOrder(a) - messageOrder(b) || a.sequence - b.sequence,
  );
}

function encodeHeaderJson(value: unknown): string {
  const bytes = new TextEncoder().encode(JSON.stringify(value));
  let binary = "";
  bytes.forEach((byte) => {
    binary += String.fromCharCode(byte);
  });
  return btoa(binary).replace(/\+/g, "-").replace(/\//g, "_").replace(/=+$/, "");
}

function readNetworkMetrics(report: RTCStatsReport): NetworkMetrics | null {
  const outbound: Partial<Record<"audio" | "video", number>> = {};
  const lost: Partial<Record<"audio" | "video", number>> = {};
  let roundTripTime: number | null = null;
  let availableOutgoingBitrate: number | null = null;

  report.forEach((stat) => {
    const value = stat as RTCStats & Record<string, unknown>;
    const kind = value.kind ?? value.mediaType;
    if (
      value.type === "outbound-rtp" &&
      (kind === "audio" || kind === "video") &&
      typeof value.packetsSent === "number"
    ) {
      outbound[kind] = value.packetsSent;
    } else if (
      value.type === "remote-inbound-rtp" &&
      (kind === "audio" || kind === "video")
    ) {
      if (typeof value.packetsLost === "number") lost[kind] = value.packetsLost;
      if (typeof value.roundTripTime === "number") {
        roundTripTime = value.roundTripTime;
      }
    } else if (
      value.type === "candidate-pair" &&
      value.state === "succeeded" &&
      (value.nominated === true || value.selected === true)
    ) {
      if (typeof value.currentRoundTripTime === "number") {
        roundTripTime = value.currentRoundTripTime;
      }
      if (typeof value.availableOutgoingBitrate === "number") {
        availableOutgoingBitrate = value.availableOutgoingBitrate;
      }
    }
  });

  const kind = outbound.video !== undefined ? "video" : "audio";
  if (outbound[kind] === undefined && roundTripTime === null) return null;
  return {
    packetsSent: outbound[kind] ?? 0,
    packetsLost: lost[kind] ?? 0,
    roundTripTime,
    availableOutgoingBitrate,
  };
}

function classifyNetworkQuality(
  metrics: NetworkMetrics,
  previous: NetworkStatsSample | null,
): NetworkQuality {
  const sentDelta = previous
    ? Math.max(0, metrics.packetsSent - previous.packetsSent)
    : 0;
  const lostDelta = previous
    ? Math.max(0, metrics.packetsLost - previous.packetsLost)
    : 0;
  const lossRate = sentDelta > 0 ? lostDelta / sentDelta : null;
  const { roundTripTime: rtt, availableOutgoingBitrate: bandwidth } = metrics;

  if (lossRate === null && rtt === null && bandwidth === null) return "unknown";
  if (
    (lossRate !== null && lossRate >= 0.08) ||
    (rtt !== null && rtt >= 0.45) ||
    (bandwidth !== null && bandwidth < 350_000)
  ) return "poor";
  if (
    (lossRate !== null && lossRate >= 0.03) ||
    (rtt !== null && rtt >= 0.25) ||
    (bandwidth !== null && bandwidth < 550_000)
  ) return "fair";
  return "good";
}

export function useVoicePreview(
  assistantId: string | null,
  onNodeActive?: (nodeId: string | null) => void,
) {
  const [status, setStatus] = useState<VoicePreviewStatus>("idle");
  const [error, setError] = useState<string | null>(null);
  const [micWarning, setMicWarning] = useState<string | null>(null);
  const [localStream, setLocalStream] = useState<MediaStream | null>(null);
  const [videoStream, setVideoStream] = useState<MediaStream | null>(null);
  const [remoteStream, setRemoteStream] = useState<MediaStream | null>(null);
  const [messages, setMessages] = useState<ChatMessage[]>([]);
  const [callEnded, setCallEnded] = useState(false);
  const [networkQuality, setNetworkQuality] =
    useState<NetworkQuality>("unknown");
  const [audioInputs, setAudioInputs] = useState<MediaDeviceInfo[]>([]);
  const [audioOutputs, setAudioOutputs] = useState<MediaDeviceInfo[]>([]);
  const [selectedDeviceId, setSelectedDeviceId] = useState("");
  const [selectedOutputDeviceId, setSelectedOutputDeviceId] = useState("");

  const transportRef = useRef<AppSmallWebRTCTransport | null>(null);
  const startingRef = useRef(false);
  const messageSeqRef = useRef(0);
  const networkStatsRef = useRef<NetworkStatsSample | null>(null);
  const pendingAssistantTurnsRef = useRef(
    new Map<string, { timestamp: string }>(),
  );
  const endedByServerRef = useRef(false);
  const selectedDeviceIdRef = useRef("");
  const selectedOutputDeviceIdRef = useRef("");
  const onNodeActiveRef = useRef(onNodeActive);
  const audioRef = useRef<HTMLAudioElement | null>(null);

  useEffect(() => {
    onNodeActiveRef.current = onNodeActive;
  }, [onNodeActive]);

  const supportsOutputSelection =
    typeof HTMLAudioElement !== "undefined" &&
    "setSinkId" in HTMLAudioElement.prototype;

  const applyOutputDevice = useCallback(async (deviceId: string) => {
    const audio = audioRef.current;
    if (!audio || !supportsOutputSelection) return;
    try {
      await audio.setSinkId(deviceId);
    } catch {
      // 设备可能已拔出，交给浏览器回退默认输出。
    }
  }, [supportsOutputSelection]);

  const refreshDevices = useCallback(async () => {
    if (!navigator.mediaDevices?.enumerateDevices) return;
    try {
      const devices = await navigator.mediaDevices.enumerateDevices();
      setAudioInputs(
        devices.filter((item) => item.kind === "audioinput" && item.deviceId),
      );
      setAudioOutputs(
        devices.filter((item) => item.kind === "audiooutput" && item.deviceId),
      );
    } catch {
      // 设备枚举失败不影响仅接收模式。
    }
  }, []);

  useEffect(() => {
    // eslint-disable-next-line react-hooks/set-state-in-effect
    void refreshDevices();
    navigator.mediaDevices?.addEventListener("devicechange", refreshDevices);
    return () =>
      navigator.mediaDevices?.removeEventListener("devicechange", refreshDevices);
  }, [refreshDevices]);

  const releaseResources = useCallback(() => {
    const transport = transportRef.current;
    transportRef.current = null;
    transport?.disconnect().catch(() => {});
    if (audioRef.current) audioRef.current.srcObject = null;
    startingRef.current = false;
    pendingAssistantTurnsRef.current.clear();
    networkStatsRef.current = null;
    onNodeActiveRef.current?.(null);
  }, []);

  const disconnect = useCallback(() => {
    releaseResources();
    setLocalStream(null);
    setVideoStream(null);
    setRemoteStream(null);
    setMessages([]);
    messageSeqRef.current = 0;
    pendingAssistantTurnsRef.current.clear();
    setError(null);
    setMicWarning(null);
    setNetworkQuality("unknown");
    setStatus("idle");
  }, [releaseResources]);

  const fail = useCallback((message: string) => {
    releaseResources();
    setLocalStream(null);
    setVideoStream(null);
    setRemoteStream(null);
    setError(message);
    setNetworkQuality("unknown");
    setStatus("failed");
  }, [releaseResources]);

  const handleAppMessage = useCallback((msg: AppMessage) => {
    if (
      msg.type === "assistant-text-start" &&
      typeof msg.turn_id === "string"
    ) {
      pendingAssistantTurnsRef.current.set(msg.turn_id, {
        timestamp:
          typeof msg.timestamp === "string"
            ? msg.timestamp
            : new Date().toISOString(),
      });
    } else if (
      msg.type === "assistant-text-delta" &&
      typeof msg.turn_id === "string" &&
      typeof msg.delta === "string" &&
      msg.delta.length > 0
    ) {
      const pending = pendingAssistantTurnsRef.current.get(msg.turn_id);
      if (pending) {
        pendingAssistantTurnsRef.current.delete(msg.turn_id);
        messageSeqRef.current += 1;
        const sequence = messageSeqRef.current;
        setMessages((previous) =>
          sortMessages([
            ...previous,
            {
              id: `assistant-${msg.turn_id as string}`,
              role: "assistant",
              content: msg.delta as string,
              timestamp: pending.timestamp,
              sequence,
              turnId: msg.turn_id as string,
              streaming: true,
            },
          ]),
        );
      } else {
        setMessages((previous) =>
          previous.map((message) =>
            message.turnId === msg.turn_id
              ? { ...message, content: message.content + msg.delta }
              : message,
          ),
        );
      }
    } else if (
      msg.type === "assistant-text-end" &&
      typeof msg.turn_id === "string"
    ) {
      const pending = pendingAssistantTurnsRef.current.get(msg.turn_id);
      pendingAssistantTurnsRef.current.delete(msg.turn_id);
      const finalContent =
        typeof msg.content === "string" ? msg.content.trim() : "";
      setMessages((previous) => {
        const existing = previous.some((message) => message.turnId === msg.turn_id);
        if (existing) {
          return previous.map((message) =>
            message.turnId === msg.turn_id
              ? {
                  ...message,
                  content: finalContent || message.content,
                  streaming: false,
                }
              : message,
          );
        }
        if (!finalContent) return previous;
        messageSeqRef.current += 1;
        return sortMessages([
          ...previous,
          {
            id: `assistant-${msg.turn_id as string}`,
            role: "assistant",
            content: finalContent,
            timestamp: pending?.timestamp ?? new Date().toISOString(),
            sequence: messageSeqRef.current,
            turnId: msg.turn_id as string,
            streaming: false,
          },
        ]);
      });
    } else if (
      msg.type === "transcript" &&
      (msg.role === "user" || msg.role === "assistant") &&
      typeof msg.content === "string" &&
      msg.content.trim()
    ) {
      messageSeqRef.current += 1;
      setMessages((previous) =>
        sortMessages([
          ...previous,
          {
            id: `msg-${messageSeqRef.current}`,
            role: msg.role as "user" | "assistant",
            content: msg.content as string,
            timestamp:
              typeof msg.timestamp === "string"
                ? msg.timestamp
                : new Date().toISOString(),
            sequence: messageSeqRef.current,
          },
        ]),
      );
    } else if (msg.type === "node-active" && typeof msg.nodeId === "string") {
      onNodeActiveRef.current?.(msg.nodeId);
    } else if (msg.type === "call-ended") {
      endedByServerRef.current = true;
      setCallEnded(true);
      disconnect();
    }
  }, [disconnect]);

  const connect = useCallback(async (options: ConnectOptions = {}) => {
    if (startingRef.current || transportRef.current) return;
    if (!assistantId) {
      setError("请先保存助手，再开始语音预览。");
      setStatus("failed");
      return;
    }

    startingRef.current = true;
    setStatus("connecting");
    setError(null);
    setMicWarning(null);
    setMessages([]);
    pendingAssistantTurnsRef.current.clear();
    setCallEnded(false);
    endedByServerRef.current = false;

    const iceServers = await webrtcApi
      .iceServers()
      .then((response) => response.iceServers)
      .catch(() => [{ urls: "stun:stun.l.google.com:19302" }]);
    const transport = new AppSmallWebRTCTransport({ iceServers });
    transportRef.current = transport;
    transport.onAppMessage = handleAppMessage;

    const callbacks: NonNullable<PipecatClientOptions["callbacks"]> = {
      onConnected: () => setStatus("connected"),
      onDisconnected: () => {
        if (transportRef.current !== transport) return;
        if (endedByServerRef.current) disconnect();
        else fail("WebRTC 连接已断开。");
      },
      onTrackStarted: (track) => {
        if (track.kind !== "audio") return;
        const stream = new MediaStream([track]);
        setRemoteStream(stream);
        if (audioRef.current) {
          audioRef.current.srcObject = stream;
          void applyOutputDevice(selectedOutputDeviceIdRef.current).then(() =>
            audioRef.current?.play().catch(() => {}),
          );
        }
      },
      onDeviceError: (deviceError) => {
        setMicWarning(deviceError.message || "无法访问麦克风。已尝试继续连接。");
      },
      onCamUpdated: () => {
        const track = transport.tracks().local.video;
        setVideoStream(track ? new MediaStream([track]) : null);
      },
      onTransportStateChanged: (nextState) => {
        if (nextState === "connected" || nextState === "ready") {
          setStatus("connected");
        } else if (nextState === "error") {
          fail("WebRTC 连接失败。");
        }
      },
    };

    transport.initialize(
      {
        transport,
        enableMic: true,
        enableCam: Boolean(options.visionEnabled),
        callbacks,
      },
      () => {},
    );

    try {
      await transport.initDevices();
      if (selectedDeviceIdRef.current) {
        await transport.updateMic(selectedDeviceIdRef.current);
      }
      const localAudio = transport.tracks().local.audio;
      const localVideo = transport.tracks().local.video;
      setLocalStream(localAudio ? new MediaStream([localAudio]) : null);
      setVideoStream(localVideo ? new MediaStream([localVideo]) : null);
      void refreshDevices();

      const request = new Request(`${API_BASE}/api/webrtc/offer`, {
        method: "POST",
        headers: {
          "Content-Type": "application/json",
          "X-Pipecat-Assistant-ID": assistantId,
          "X-Pipecat-Dynamic-Variables": encodeHeaderJson(
            options.dynamicVariables ?? {},
          ),
        },
        credentials: "include",
      });
      await transport.connect({
        webrtcRequestParams: {
          endpoint: request,
          requestData: {
            assistant_id: assistantId,
            vision_enabled: Boolean(options.visionEnabled),
            dynamic_variables: options.dynamicVariables ?? {},
          },
        },
      });
      transport.sendAppMessage({ type: "client-ready" });
      setStatus("connected");
    } catch (connectionError) {
      fail(errorMessage(connectionError, "无法连接语音服务。"));
    } finally {
      startingRef.current = false;
    }
  }, [
    assistantId,
    applyOutputDevice,
    disconnect,
    fail,
    handleAppMessage,
    refreshDevices,
  ]);

  const replaceVideoStream = useCallback(async (videoStream: MediaStream | null) => {
    const transport = transportRef.current;
    const deviceId = videoStream?.getVideoTracks()[0]?.getSettings().deviceId;
    if (transport && deviceId) await transport.updateCam(deviceId);
  }, []);

  const selectCamera = useCallback((deviceId: string) => {
    transportRef.current?.updateCam(deviceId);
  }, []);

  const selectDevice = useCallback(async (deviceId: string) => {
    setSelectedDeviceId(deviceId);
    selectedDeviceIdRef.current = deviceId;
    const transport = transportRef.current;
    if (!transport) return;
    try {
      await transport.updateMic(deviceId);
      const audio = transport.tracks().local.audio;
      setLocalStream(audio ? new MediaStream([audio]) : null);
      setMicWarning(null);
    } catch (deviceError) {
      setMicWarning(errorMessage(deviceError, "无法切换麦克风。"));
    }
  }, []);

  const selectOutputDevice = useCallback((deviceId: string) => {
    setSelectedOutputDeviceId(deviceId);
    selectedOutputDeviceIdRef.current = deviceId;
    void applyOutputDevice(deviceId);
  }, [applyOutputDevice]);

  useEffect(() => {
    if (status !== "connected") return;
    let cancelled = false;
    const sample = async () => {
      const transport = transportRef.current;
      if (!transport) return;
      try {
        const report = await transport.connectionStats();
        const metrics = report ? readNetworkMetrics(report) : null;
        if (cancelled || transportRef.current !== transport) return;
        if (!metrics) {
          setNetworkQuality("unknown");
          return;
        }
        setNetworkQuality(
          classifyNetworkQuality(metrics, networkStatsRef.current),
        );
        networkStatsRef.current = {
          packetsSent: metrics.packetsSent,
          packetsLost: metrics.packetsLost,
        };
      } catch {
        if (!cancelled) setNetworkQuality("unknown");
      }
    };
    void sample();
    const interval = window.setInterval(() => void sample(), NETWORK_SAMPLE_INTERVAL_MS);
    return () => {
      cancelled = true;
      window.clearInterval(interval);
    };
  }, [status]);

  const sendText = useCallback((text: string): boolean => {
    const trimmed = text.trim();
    const transport = transportRef.current;
    if (!trimmed || !transport || transport.state !== "connected") return false;
    transport.sendAppMessage({ type: "user-text", text: trimmed });
    return true;
  }, []);

  useEffect(() => releaseResources, [releaseResources]);

  return {
    status,
    error,
    micWarning,
    localStream,
    videoStream,
    remoteStream,
    messages,
    callEnded,
    networkQuality,
    audioInputs,
    audioOutputs,
    selectedDeviceId,
    selectedOutputDeviceId,
    setSelectedDeviceId,
    selectDevice,
    selectOutputDevice,
    supportsOutputSelection,
    sendText,
    connect,
    replaceVideoStream,
    selectCamera,
    disconnect,
    audioRef,
  };
}

export type VoicePreview = ReturnType<typeof useVoicePreview>;