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

"use client";

/**
 * 语音预览:把麦克风接到后端 /ws/voice(WebRTC 信令),听到助手实时回应。
 *
 * 走原生 RTCPeerConnection + 一条 ws 信令通道,与后端 voice_webrtc.py 的约定对齐:
 *   client → {type:"offer",         payload:{pc_id, sdp, type, assistant_id}}
 *   server → {type:"answer",        payload:{pc_id, sdp, type}}
 *   client → {type:"ice-candidate", payload:{pc_id, candidate:{...}}}
 * 音频本身走 WebRTC 媒体流(Opus),不经 ws;后端 TTS 帧从 ontrack 拿到直接播放。
 *
 * 纯本机(localhost)即可跑:localhost 是 secure context,麦克风可用,ws 用明文。
 * 局域网/别的设备要 https+wss,见 deploy/README.md。
 */

import { useCallback, useEffect, useRef, useState } from "react";

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

export type VoicePreviewStatus = "idle" | "connecting" | "connected" | "failed";

// http→ws、https→wss,自动跟随 API 基址(同源反代时也对)
function wsBaseUrl(): string {
  const url = new URL(API_BASE, window.location.origin);
  url.protocol = url.protocol === "https:" ? "wss:" : "ws:";
  return url.toString().replace(/\/$/, "");
}

function generatePcId(): string {
  const bytes = new Uint8Array(16);
  crypto.getRandomValues(bytes);
  return (
    "PC-" +
    Array.from(bytes)
      .map((b) => b.toString(16).padStart(2, "0"))
      .join("")
  );
}

type UseVoicePreviewOptions = {
  /** 取麦克风失败(权限/无设备)时回调,供 UI 提示。 */
  onMicError?: () => void;
};

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

export function useVoicePreview(
  assistantId: string | null,
  { onMicError }: UseVoicePreviewOptions = {},
) {
  const [status, setStatus] = useState<VoicePreviewStatus>("idle");
  const [error, setError] = useState<string | null>(null);
  const [localStream, setLocalStream] = useState<MediaStream | null>(null);
  const audioRef = useRef<HTMLAudioElement | null>(null);
  const pcRef = useRef<RTCPeerConnection | null>(null);
  const wsRef = useRef<WebSocket | null>(null);
  const localStreamRef = useRef<MediaStream | null>(null);
  const startingRef = useRef(false);

  const releaseResources = useCallback(() => {
    const ws = wsRef.current;
    wsRef.current = null;
    if (ws) {
      ws.onclose = null;
      ws.onerror = null;
      ws.onmessage = null;
      ws.close();
    }

    const pc = pcRef.current;
    pcRef.current = null;
    if (pc) {
      pc.onconnectionstatechange = null;
      pc.onicecandidate = null;
      pc.oniceconnectionstatechange = null;
      pc.ontrack = null;
      pc.close();
    }

    localStreamRef.current?.getTracks().forEach((track) => track.stop());
    localStreamRef.current = null;
    if (audioRef.current) audioRef.current.srcObject = null;
    startingRef.current = false;
  }, []);

  const disconnect = useCallback(() => {
    releaseResources();
    setLocalStream(null);
    setError(null);
    setStatus("idle");
  }, [releaseResources]);

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

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

    startingRef.current = true;
    setError(null);
    setStatus("connecting");

    const pcId = generatePcId();
    const ws = new WebSocket(`${wsBaseUrl()}/ws/voice`);
    wsRef.current = ws;

    ws.onmessage = async (event) => {
      try {
        const msg = JSON.parse(event.data);
        if (msg.type === "answer") {
          await pcRef.current?.setRemoteDescription({
            type: "answer",
            sdp: msg.payload.sdp,
          });
        } else if (msg.type === "ice-candidate" && msg.payload?.candidate) {
          // 后端当前不主动 trickle,留兼容
          try {
            await pcRef.current?.addIceCandidate(msg.payload.candidate);
          } catch {
            /* 忽略迟到/重复 candidate */
          }
        } else if (msg.type === "error") {
          fail(msg.payload?.message || "后端无法启动语音会话。");
        }
      } catch {
        /* 非 JSON / 未知消息,忽略 */
      }
    };

    try {
      // 1) 等 ws 连上
      await new Promise<void>((resolve, reject) => {
        ws.onopen = () => resolve();
        ws.onerror = (e) => reject(e);
        ws.onclose = () => reject(new Error("语音信令连接已关闭。"));
      });
      // 连上后,信令异常或关闭都结束当前会话并保留失败状态。
      ws.onerror = () => {
        if (wsRef.current === ws) fail("语音信令连接失败。");
      };
      ws.onclose = () => {
        if (wsRef.current === ws) fail("语音信令连接已断开。");
      };

      // 2) 建 PeerConnection(纯 STUN,本机/局域网够用)
      const pc = new RTCPeerConnection({
        iceServers: [{ urls: "stun:stun.l.google.com:19302" }],
      });
      pcRef.current = pc;

      pc.onicecandidate = (e) => {
        if (ws.readyState !== WebSocket.OPEN) return;
        ws.send(
          JSON.stringify({
            type: "ice-candidate",
            payload: {
              pc_id: pcId,
              candidate: e.candidate
                ? {
                    candidate: e.candidate.candidate,
                    sdpMid: e.candidate.sdpMid,
                    sdpMLineIndex: e.candidate.sdpMLineIndex,
                  }
                : null,
            },
          }),
        );
      };

      pc.ontrack = (e) => {
        if (e.track.kind === "audio" && audioRef.current) {
          audioRef.current.srcObject =
            e.streams[0] ?? new MediaStream([e.track]);
          void audioRef.current.play().catch(() => {});
        }
      };

      pc.onconnectionstatechange = () => {
        if (pcRef.current !== pc) return;
        if (pc.connectionState === "connected") setStatus("connected");
        else if (pc.connectionState === "failed")
          fail("WebRTC 音频连接失败。");
      };

      pc.oniceconnectionstatechange = () => {
        if (pcRef.current !== pc) return;
        const st = pc.iceConnectionState;
        if (st === "connected" || st === "completed") setStatus("connected");
        else if (st === "failed") fail("WebRTC 音频连接失败。");
        else if (st === "disconnected") fail("WebRTC 音频连接已断开。");
      };

      // 3) 取麦克风 → 加入连接
      let stream: MediaStream;
      try {
        stream = await navigator.mediaDevices.getUserMedia({
          audio: {
            echoCancellation: true,
            noiseSuppression: true,
            autoGainControl: true,
          },
        });
      } catch (mediaError) {
        onMicError?.();
        fail(errorMessage(mediaError, "无法访问麦克风。"));
        return;
      }
      localStreamRef.current = stream;
      setLocalStream(stream);
      stream.getTracks().forEach((track) => pc.addTrack(track, stream));

      // 4) 生成 offer 并发给后端(assistant_id 在 payload 顶层)
      const offer = await pc.createOffer();
      await pc.setLocalDescription(offer);
      const localDescription = pc.localDescription;
      if (!localDescription?.sdp) {
        throw new Error("浏览器无法创建 WebRTC offer。");
      }
      ws.send(
        JSON.stringify({
          type: "offer",
          payload: {
            pc_id: pcId,
            sdp: localDescription.sdp,
            type: localDescription.type,
            assistant_id: assistantId,
          },
        }),
      );
    } catch (connectionError) {
      fail(errorMessage(connectionError, "无法连接语音服务。"));
    } finally {
      startingRef.current = false;
    }
  }, [assistantId, fail, onMicError]);

  // 卸载时收尾
  useEffect(() => releaseResources, [releaseResources]);

  return { status, error, localStream, connect, disconnect, audioRef };
}