nanobot-voice-interface/frontend/src/hooks/useWebRTC.ts

import { useCallback, useEffect, useRef, useState } from "preact/hooks";
import type {
  AgentState,
  CardItem,
  CardLane,
  CardMessageMetadata,
  CardState,
  ClientMessage,
  LogLine,
  ServerMessage,
} from "../types";

const BACKEND_URL = import.meta.env.VITE_BACKEND_URL ?? "";

let cardIdCounter = 0;
let logIdCounter = 0;

const LANE_RANK: Record<CardLane, number> = {
  attention: 0,
  work: 1,
  context: 2,
  history: 3,
};

const STATE_RANK: Record<CardState, number> = {
  active: 0,
  stale: 1,
  resolved: 2,
  superseded: 3,
  archived: 4,
};

interface WebRTCState {
  connected: boolean;
  connecting: boolean;
  textOnly: boolean;
  agentState: AgentState;
  logLines: LogLine[];
  cards: CardItem[];
  voiceStatus: string;
  statusVisible: boolean;
  remoteAudioEl: HTMLAudioElement | null;
  remoteStream: MediaStream | null;
  sendJson(msg: ClientMessage): void;
  sendTextMessage(text: string, metadata?: CardMessageMetadata): Promise<void>;
  dismissCard(id: number): void;
  setTextOnly(enabled: boolean): void;
  connect(): Promise<void>;
}

type AppendLine = (role: string, text: string, timestamp: string) => void;
type UpsertCard = (item: Omit<CardItem, "id">) => void;
type SetAgentState = (updater: (prev: AgentState) => AgentState) => void;
type RawPersistedCard =
  | Extract<ServerMessage, { type: "card" }>
  | (Omit<Extract<ServerMessage, { type: "card" }>, "type"> & { type?: "card" });
interface IdleFallbackControls {
  clear(): void;
  schedule(delayMs?: number): void;
}

interface RTCRefs {
  pcRef: { current: RTCPeerConnection | null };
  dcRef: { current: RTCDataChannel | null };
  remoteAudioRef: { current: HTMLAudioElement | null };
  micSendersRef: { current: RTCRtpSender[] };
  localTracksRef: { current: MediaStreamTrack[] };
}

interface RTCCallbacks {
  setConnected: (v: boolean) => void;
  setConnecting: (v: boolean) => void;
  setRemoteStream: (s: MediaStream | null) => void;
  showStatus: (text: string, persistMs?: number) => void;
  appendLine: AppendLine;
  onDcMessage: (raw: string) => void;
  onDcOpen: () => void;
  closePC: () => void;
}

function compareCards(a: CardItem, b: CardItem): number {
  const laneDiff = LANE_RANK[a.lane] - LANE_RANK[b.lane];
  if (laneDiff !== 0) return laneDiff;
  const stateDiff = STATE_RANK[a.state] - STATE_RANK[b.state];
  if (stateDiff !== 0) return stateDiff;
  if (a.priority !== b.priority) return b.priority - a.priority;
  const updatedDiff = b.updatedAt.localeCompare(a.updatedAt);
  if (updatedDiff !== 0) return updatedDiff;
  return b.createdAt.localeCompare(a.createdAt);
}

function sortCards(items: CardItem[]): CardItem[] {
  return [...items].sort(compareCards);
}

function stopMediaTracks(tracks: Iterable<MediaStreamTrack | null | undefined>): void {
  const seen = new Set<MediaStreamTrack>();
  for (const track of tracks) {
    if (!track || seen.has(track)) continue;
    seen.add(track);
    track.stop();
  }
}

function toCardItem(msg: Extract<ServerMessage, { type: "card" }>): Omit<CardItem, "id"> {
  return {
    serverId: msg.id,
    kind: msg.kind,
    content: msg.content,
    title: msg.title,
    question: msg.question || undefined,
    choices: msg.choices.length > 0 ? msg.choices : undefined,
    responseValue: msg.response_value || undefined,
    slot: msg.slot || undefined,
    lane: msg.lane,
    priority: msg.priority,
    state: msg.state,
    templateKey: msg.template_key || undefined,
    contextSummary: msg.context_summary || undefined,
    createdAt: msg.created_at || new Date().toISOString(),
    updatedAt: msg.updated_at || new Date().toISOString(),
  };
}

function appendLogLineEntry(
  prev: LogLine[],
  role: string,
  text: string,
  timestamp: string,
): LogLine[] {
  const next = [
    ...prev,
    { id: logIdCounter++, role, text, timestamp: timestamp || new Date().toISOString() },
  ];
  return next.length > 250 ? next.slice(next.length - 250) : next;
}

function mergeCardItem(prev: CardItem[], item: Omit<CardItem, "id">): CardItem[] {
  const existingIndex = item.serverId
    ? prev.findIndex((card) => card.serverId === item.serverId)
    : -1;
  if (existingIndex >= 0) {
    const next = [...prev];
    next[existingIndex] = { ...next[existingIndex], ...item };
    return sortCards(next);
  }
  return sortCards([...prev, { ...item, id: cardIdCounter++ }]);
}

function normalizePersistedCard(raw: RawPersistedCard): Omit<CardItem, "id"> {
  return toCardItem({
    type: "card",
    ...(raw as Omit<Extract<ServerMessage, { type: "card" }>, "type">),
  });
}

function reconcilePersistedCards(prev: CardItem[], rawCards: RawPersistedCard[]): CardItem[] {
  const byServerId = new Map(
    prev.filter((card) => card.serverId).map((card) => [card.serverId as string, card.id]),
  );
  const next = rawCards.map((raw) => {
    const card = normalizePersistedCard(raw);
    return {
      ...card,
      id:
        card.serverId && byServerId.has(card.serverId)
          ? (byServerId.get(card.serverId) as number)
          : cardIdCounter++,
    };
  });
  return sortCards(next);
}

function handleTypedMessage(
  msg: Extract<ServerMessage, { type: string }>,
  setAgentState: SetAgentState,
  appendLine: AppendLine,
  upsertCard: UpsertCard,
  idleFallback: IdleFallbackControls,
): void {
  if (msg.type === "agent_state") {
    idleFallback.clear();
    setAgentState((prev) => (prev === "listening" ? prev : msg.state));
    return;
  }
  if (msg.type === "message") {
    if (msg.is_tool_hint) {
      appendLine("tool", msg.content, msg.timestamp);
      return;
    }
    if (!msg.is_progress) {
      appendLine(msg.role, msg.content, msg.timestamp);
      idleFallback.schedule();
    }
    return;
  }
  if (msg.type === "card") {
    upsertCard(toCardItem(msg));
    idleFallback.schedule();
    return;
  }
  if (msg.type === "error") {
    appendLine("system", msg.error, "");
    idleFallback.schedule();
  }
}

async function acquireMicStream(): Promise<MediaStream> {
  try {
    return await navigator.mediaDevices.getUserMedia({
      audio: {
        channelCount: 1,
        sampleRate: 48000,
        echoCancellation: true,
        noiseSuppression: true,
        autoGainControl: false,
      },
      video: false,
    });
  } catch {
    return navigator.mediaDevices.getUserMedia({ audio: true, video: false });
  }
}

function waitForIceComplete(pc: RTCPeerConnection): Promise<void> {
  return new Promise<void>((resolve) => {
    if (pc.iceGatheringState === "complete") {
      resolve();
      return;
    }
    const check = () => {
      if (pc.iceGatheringState === "complete") {
        pc.removeEventListener("icegatheringstatechange", check);
        resolve();
      }
    };
    pc.addEventListener("icegatheringstatechange", check);
    setTimeout(resolve, 5000);
  });
}

async function exchangeSdp(
  localDesc: RTCSessionDescription,
): Promise<{ sdp: string; rtcType: string }> {
  const rtcUrl = BACKEND_URL ? `${BACKEND_URL}/rtc/offer` : "/rtc/offer";
  const resp = await fetch(rtcUrl, {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({ sdp: localDesc.sdp, rtcType: localDesc.type }),
  });
  if (!resp.ok) throw new Error(`/rtc/offer returned ${resp.status}`);
  return resp.json() as Promise<{ sdp: string; rtcType: string }>;
}

async function runConnect(
  refs: RTCRefs,
  cbs: RTCCallbacks,
  opts: { textOnly: boolean },
): Promise<void> {
  if (refs.pcRef.current) return;
  if (!window.RTCPeerConnection) {
    cbs.showStatus("WebRTC unavailable in this browser.", 4000);
    return;
  }
  cbs.setConnecting(true);
  cbs.showStatus("Connecting...");

  let micStream: MediaStream | null = null;
  try {
    refs.localTracksRef.current = [];
    if (!opts.textOnly) {
      micStream = await acquireMicStream();
      const audioTracks = micStream.getAudioTracks();
      audioTracks.forEach((track) => {
        track.enabled = false;
      });
      refs.localTracksRef.current = audioTracks;
    }

    const pc = new RTCPeerConnection({ iceServers: [{ urls: "stun:stun.l.google.com:19302" }] });
    refs.pcRef.current = pc;

    const newRemoteStream = new MediaStream();
    cbs.setRemoteStream(newRemoteStream);
    if (refs.remoteAudioRef.current) {
      refs.remoteAudioRef.current.srcObject = newRemoteStream;
      refs.remoteAudioRef.current.play().catch(() => {});
    }

    pc.ontrack = (event) => {
      if (event.track.kind !== "audio") return;
      newRemoteStream.addTrack(event.track);
      refs.remoteAudioRef.current?.play().catch(() => {});
    };

    const dc = pc.createDataChannel("app", { ordered: true });
    refs.dcRef.current = dc;
    dc.onopen = () => {
      cbs.setConnected(true);
      cbs.setConnecting(false);
      cbs.showStatus(opts.textOnly ? "Text-only mode enabled" : "Hold anywhere to talk", 2500);
      cbs.appendLine("system", "Connected.", new Date().toISOString());
      cbs.onDcOpen();
    };
    dc.onclose = () => {
      cbs.appendLine("system", "Disconnected.", new Date().toISOString());
      cbs.closePC();
    };
    dc.onmessage = (e) => cbs.onDcMessage(e.data as string);

    refs.micSendersRef.current = [];
    if (micStream) {
      micStream.getAudioTracks().forEach((track) => {
        pc.addTrack(track, micStream as MediaStream);
      });
      refs.micSendersRef.current = pc.getSenders().filter((s) => s.track?.kind === "audio");
    }

    const offer = await pc.createOffer();
    await pc.setLocalDescription(offer);
    await waitForIceComplete(pc);

    const localDesc = pc.localDescription;
    if (!localDesc) throw new Error("No local description after ICE gathering");
    const answer = await exchangeSdp(localDesc);
    await pc.setRemoteDescription({ type: answer.rtcType as RTCSdpType, sdp: answer.sdp });
  } catch (err) {
    cbs.appendLine("system", `Connection failed: ${err}`, new Date().toISOString());
    cbs.showStatus("Connection failed.", 3000);
    cbs.closePC();
  }
}

function useBackendActions() {
  const sendTextMessage = useCallback(async (text: string, metadata?: CardMessageMetadata) => {
    const message = text.trim();
    if (!message) return;
    const url = BACKEND_URL ? `${BACKEND_URL}/message` : "/message";
    const resp = await fetch(url, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({ text: message, metadata: metadata ?? {} }),
    });
    if (!resp.ok) throw new Error(`Send failed (${resp.status})`);
  }, []);

  return { sendTextMessage };
}

function useCardPolling(loadPersistedCards: () => Promise<void>) {
  useEffect(() => {
    loadPersistedCards().catch(() => {});
    const pollId = window.setInterval(() => {
      loadPersistedCards().catch(() => {});
    }, 10000);
    const onVisible = () => {
      if (document.visibilityState === "visible") loadPersistedCards().catch(() => {});
    };
    const onCardsRefresh = () => {
      loadPersistedCards().catch(() => {});
    };
    window.addEventListener("focus", onVisible);
    window.addEventListener("nanobot:cards-refresh", onCardsRefresh);
    document.addEventListener("visibilitychange", onVisible);
    return () => {
      window.clearInterval(pollId);
      window.removeEventListener("focus", onVisible);
      window.removeEventListener("nanobot:cards-refresh", onCardsRefresh);
      document.removeEventListener("visibilitychange", onVisible);
    };
  }, [loadPersistedCards]);
}

function useRemoteAudioBindings({
  textOnly,
  connected,
  showStatus,
  remoteAudioRef,
  micSendersRef,
  dcRef,
  textOnlyRef,
}: {
  textOnly: boolean;
  connected: boolean;
  showStatus: (text: string, persistMs?: number) => void;
  remoteAudioRef: { current: HTMLAudioElement | null };
  micSendersRef: { current: RTCRtpSender[] };
  dcRef: { current: RTCDataChannel | null };
  textOnlyRef: { current: boolean };
}) {
  useEffect(() => {
    textOnlyRef.current = textOnly;
  }, [textOnly, textOnlyRef]);

  useEffect(() => {
    const audio = new Audio();
    audio.autoplay = true;
    (audio as HTMLAudioElement & { playsInline: boolean }).playsInline = true;
    audio.muted = textOnlyRef.current;
    remoteAudioRef.current = audio;
    return () => {
      audio.srcObject = null;
    };
  }, [remoteAudioRef, textOnlyRef]);

  useEffect(() => {
    const handler = (e: Event) => {
      const enabled = (e as CustomEvent<{ enabled: boolean }>).detail?.enabled ?? false;
      micSendersRef.current.forEach((sender) => {
        if (sender.track) sender.track.enabled = enabled && !textOnlyRef.current;
      });
    };
    window.addEventListener("nanobot-mic-enable", handler);
    return () => window.removeEventListener("nanobot-mic-enable", handler);
  }, [micSendersRef, textOnlyRef]);

  useEffect(() => {
    if (remoteAudioRef.current) {
      remoteAudioRef.current.muted = textOnly;
      if (textOnly) remoteAudioRef.current.pause();
      else remoteAudioRef.current.play().catch(() => {});
    }
    micSendersRef.current.forEach((sender) => {
      if (sender.track) sender.track.enabled = false;
    });
    if (textOnly) {
      const dc = dcRef.current;
      if (dc?.readyState === "open") {
        dc.send(JSON.stringify({ type: "voice-ptt", pressed: false } satisfies ClientMessage));
      }
    }
    if (connected) showStatus(textOnly ? "Text-only mode enabled" : "Hold anywhere to talk", 2000);
  }, [connected, dcRef, micSendersRef, remoteAudioRef, showStatus, textOnly]);
}

function useIdleFallback(setAgentState: SetAgentState): IdleFallbackControls {
  const idleTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);

  const clear = useCallback(() => {
    if (!idleTimerRef.current) return;
    clearTimeout(idleTimerRef.current);
    idleTimerRef.current = null;
  }, []);

  const schedule = useCallback(
    (delayMs = 450) => {
      clear();
      idleTimerRef.current = setTimeout(() => {
        idleTimerRef.current = null;
        setAgentState((prev) => {
          if (prev === "listening" || prev === "speaking") return prev;
          return "idle";
        });
      }, delayMs);
    },
    [clear, setAgentState],
  );

  useEffect(() => clear, [clear]);
  return { clear, schedule };
}

function useLogState() {
  const [logLines, setLogLines] = useState<LogLine[]>([]);

  const appendLine = useCallback((role: string, text: string, timestamp: string) => {
    setLogLines((prev) => appendLogLineEntry(prev, role, text, timestamp));
  }, []);

  return { logLines, appendLine };
}

async function fetchPersistedCardsFromBackend(): Promise<RawPersistedCard[] | null> {
  const url = BACKEND_URL ? `${BACKEND_URL}/cards` : "/cards";
  const resp = await fetch(url, { cache: "no-store" });
  if (!resp.ok) {
    console.warn(`[cards] /cards returned ${resp.status}`);
    return null;
  }
  return (await resp.json()) as RawPersistedCard[];
}

function useCardsState() {
  const [cards, setCards] = useState<CardItem[]>([]);

  const upsertCard = useCallback((item: Omit<CardItem, "id">) => {
    setCards((prev) => mergeCardItem(prev, item));
  }, []);

  const dismissCard = useCallback((id: number) => {
    setCards((prev) => {
      const card = prev.find((entry) => entry.id === id);
      if (card?.serverId) {
        const url = BACKEND_URL
          ? `${BACKEND_URL}/cards/${card.serverId}`
          : `/cards/${card.serverId}`;
        fetch(url, { method: "DELETE" }).catch(() => {});
      }
      return prev.filter((entry) => entry.id !== id);
    });
  }, []);

  const loadPersistedCards = useCallback(async () => {
    try {
      const rawCards = await fetchPersistedCardsFromBackend();
      if (!rawCards) return;
      setCards((prev) => reconcilePersistedCards(prev, rawCards));
    } catch (err) {
      console.warn("[cards] failed to load persisted cards", err);
    }
  }, []);

  return { cards, upsertCard, dismissCard, loadPersistedCards };
}

function parseServerMessage(raw: string): Extract<ServerMessage, { type: string }> | null {
  let msg: ServerMessage;
  try {
    msg = JSON.parse(raw);
  } catch {
    return null;
  }
  if (typeof msg !== "object" || msg === null || !("type" in msg)) return null;
  return msg as Extract<ServerMessage, { type: string }>;
}

function useDataChannelMessages({
  setAgentState,
  appendLine,
  upsertCard,
  idleFallback,
}: {
  setAgentState: SetAgentState;
  appendLine: AppendLine;
  upsertCard: UpsertCard;
  idleFallback: IdleFallbackControls;
}) {
  return useCallback(
    (raw: string) => {
      const msg = parseServerMessage(raw);
      if (!msg) return;
      handleTypedMessage(msg, setAgentState, appendLine, upsertCard, idleFallback);
    },
    [appendLine, idleFallback, setAgentState, upsertCard],
  );
}

function useMessageState() {
  const [agentState, setAgentState] = useState<AgentState>("idle");
  const { logLines, appendLine } = useLogState();
  const { cards, upsertCard, dismissCard, loadPersistedCards } = useCardsState();
  const idleFallback = useIdleFallback(setAgentState);
  const onDcMessage = useDataChannelMessages({
    setAgentState,
    appendLine,
    upsertCard,
    idleFallback,
  });

  return { agentState, logLines, cards, appendLine, dismissCard, loadPersistedCards, onDcMessage };
}

function usePeerConnectionControls({
  textOnly,
  connected,
  appendLine,
  onDcMessage,
  loadPersistedCards,
  showStatus,
  refs,
  setConnected,
  setConnecting,
  setRemoteStream,
  textOnlyRef,
}: {
  textOnly: boolean;
  connected: boolean;
  appendLine: AppendLine;
  onDcMessage: (raw: string) => void;
  loadPersistedCards: () => Promise<void>;
  showStatus: (text: string, persistMs?: number) => void;
  refs: RTCRefs;
  setConnected: (value: boolean) => void;
  setConnecting: (value: boolean) => void;
  setRemoteStream: (stream: MediaStream | null) => void;
  textOnlyRef: { current: boolean };
}) {
  const closePC = useCallback(() => {
    const dc = refs.dcRef.current;
    const pc = refs.pcRef.current;
    const localTracks = refs.localTracksRef.current;
    const senderTracks = refs.micSendersRef.current.map((sender) => sender.track);

    refs.dcRef.current = null;
    refs.pcRef.current = null;
    refs.micSendersRef.current = [];
    refs.localTracksRef.current = [];

    stopMediaTracks([...senderTracks, ...localTracks]);
    dc?.close();
    pc?.close();

    setConnected(false);
    setConnecting(false);
    if (refs.remoteAudioRef.current) refs.remoteAudioRef.current.srcObject = null;
    setRemoteStream(null);
  }, [refs, setConnected, setConnecting, setRemoteStream]);
  const closePCRef = useRef(closePC);

  useEffect(() => {
    closePCRef.current = closePC;
  }, [closePC]);

  const connect = useCallback(async () => {
    await runConnect(
      refs,
      {
        setConnected,
        setConnecting,
        setRemoteStream,
        showStatus,
        appendLine,
        onDcMessage,
        onDcOpen: () => {
          loadPersistedCards().catch(() => {});
        },
        closePC,
      },
      { textOnly: textOnlyRef.current },
    );
  }, [
    appendLine,
    closePC,
    loadPersistedCards,
    onDcMessage,
    refs,
    setConnected,
    setConnecting,
    setRemoteStream,
    showStatus,
    textOnlyRef,
  ]);

  useEffect(() => {
    if (textOnly || !connected || refs.micSendersRef.current.length > 0) return;
    closePC();
    connect().catch(() => {});
  }, [closePC, connect, connected, refs.micSendersRef, textOnly]);

  useEffect(() => {
    return () => {
      closePCRef.current();
    };
  }, []);

  return { closePC, connect };
}

export function useWebRTC(): WebRTCState {
  const [connected, setConnected] = useState(false);
  const [connecting, setConnecting] = useState(false);
  const [textOnly, setTextOnlyState] = useState(false);
  const [voiceStatus, setVoiceStatus] = useState("");
  const [statusVisible, setStatusVisible] = useState(false);
  const [remoteStream, setRemoteStream] = useState<MediaStream | null>(null);
  const refs: RTCRefs = {
    pcRef: useRef<RTCPeerConnection | null>(null),
    dcRef: useRef<RTCDataChannel | null>(null),
    remoteAudioRef: useRef<HTMLAudioElement | null>(null),
    micSendersRef: useRef<RTCRtpSender[]>([]),
    localTracksRef: useRef<MediaStreamTrack[]>([]),
  };
  const statusTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
  const textOnlyRef = useRef(false);
  const { sendTextMessage } = useBackendActions();
  const { agentState, logLines, cards, appendLine, dismissCard, loadPersistedCards, onDcMessage } =
    useMessageState();

  const setTextOnly = useCallback((enabled: boolean) => {
    textOnlyRef.current = enabled;
    setTextOnlyState(enabled);
  }, []);
  const showStatus = useCallback((text: string, persistMs = 0) => {
    setVoiceStatus(text);
    setStatusVisible(true);
    if (statusTimerRef.current) clearTimeout(statusTimerRef.current);
    if (persistMs > 0)
      statusTimerRef.current = setTimeout(() => setStatusVisible(false), persistMs);
  }, []);
  const sendJson = useCallback(
    (msg: ClientMessage) => {
      const dc = refs.dcRef.current;
      if (dc?.readyState === "open") dc.send(JSON.stringify(msg));
    },
    [refs.dcRef],
  );

  useCardPolling(loadPersistedCards);
  useRemoteAudioBindings({
    textOnly,
    connected,
    showStatus,
    remoteAudioRef: refs.remoteAudioRef,
    micSendersRef: refs.micSendersRef,
    dcRef: refs.dcRef,
    textOnlyRef,
  });
  const { connect } = usePeerConnectionControls({
    textOnly,
    connected,
    appendLine,
    onDcMessage,
    loadPersistedCards,
    showStatus,
    refs,
    setConnected,
    setConnecting,
    setRemoteStream,
    textOnlyRef,
  });

  return {
    connected,
    connecting,
    textOnly,
    agentState,
    logLines,
    cards,
    voiceStatus,
    statusVisible,
    remoteAudioEl: refs.remoteAudioRef.current,
    remoteStream,
    sendJson,
    sendTextMessage,
    dismissCard,
    setTextOnly,
    connect,
  };
}