dumpsterChat/web/src/stores/voice.ts

import { create } from 'zustand';
import {
  Room,
  RoomEvent,
  Track,
  AudioPresets,
  type RemoteAudioTrack,
  type RemoteParticipant,
  type Participant,
  type TrackPublication,
} from 'livekit-client';
import { api } from '../lib/api.ts';
import { useWebSocketStore } from './ws.ts';
import { BackgroundProcessor } from '@livekit/track-processors';


export interface VoiceParticipant {
  identity: string;
  username: string;
  isMuted: boolean;
  isSpeaking: boolean;
  hasVideo: boolean;
  isScreenSharing: boolean;
}

export interface WhisperNotification {
  fromUserId: string;
  fromUsername: string;
  message: string;
  timestamp: number;
}

interface VoiceTokenResponse {
  token: string;
  livekit_url: string;
}

interface VoiceState {
  currentRoom: string | null;
  currentRoomName: string | null;
  isConnected: boolean;
  isMuted: boolean;
  isDeafened: boolean;
  isVideoOn: boolean;
  isScreenSharing: boolean;
  noiseSuppression: boolean;
  bgMode: 'disabled' | 'background-blur' | 'virtual-background';
  _bgProcessor: any;
  participants: VoiceParticipant[];
  error: string | null;
  _room: Room | null;
  isJoining: boolean;
  whispers: WhisperNotification[];
  joinVoice: (channelId: string, channelName: string) => Promise<void>;
  leaveVoice: () => Promise<void>;
  toggleMute: () => void;
  toggleDeafen: () => void;
  toggleVideo: () => Promise<void>;
  toggleScreenShare: () => Promise<void>;
  toggleNoiseSuppression: () => void;
  setBgMode: (mode: 'disabled' | 'background-blur' | 'virtual-background', imagePath?: string) => Promise<void>;
  _addWhisper: (w: WhisperNotification) => void;
  whisperTo: (targetUserId: string, targetUsername: string, message?: string) => void;
  dismissWhisper: (timestamp: number) => void;
}

const DEVICE_PREFS_KEY = 'dc_voice_device_prefs';

interface DevicePrefs {
  audioInput?: string;
  audioOutput?: string;
  videoInput?: string;
}

function loadDevicePrefs(): DevicePrefs {
  try {
    return JSON.parse(localStorage.getItem(DEVICE_PREFS_KEY) || '{}');
  } catch {
    return {};
  }
}

function saveDevicePrefs(prefs: DevicePrefs) {
  localStorage.setItem(DEVICE_PREFS_KEY, JSON.stringify(prefs));
}

export function saveDevicePref(kind: MediaDeviceKind, deviceId: string) {
  const prefs = loadDevicePrefs();
  if (kind === 'audioinput') prefs.audioInput = deviceId;
  else if (kind === 'audiooutput') prefs.audioOutput = deviceId;
  else if (kind === 'videoinput') prefs.videoInput = deviceId;
  saveDevicePrefs(prefs);
}

export function getSavedDeviceId(kind: MediaDeviceKind): string | undefined {
  const prefs = loadDevicePrefs();
  if (kind === 'audioinput') return prefs.audioInput;
  if (kind === 'audiooutput') return prefs.audioOutput;
  if (kind === 'videoinput') return prefs.videoInput;
  return undefined;
}

function participantToVoice(p: Participant): VoiceParticipant {
  const micPub = p.getTrackPublication(Track.Source.Microphone);
  const camPub = p.getTrackPublication(Track.Source.Camera);
  const screenPub = p.getTrackPublication(Track.Source.ScreenShare);
  const screenAudioPub = p.getTrackPublication(Track.Source.ScreenShareAudio);
  return {
    identity: p.identity,
    username: p.name || p.identity,
    isMuted: micPub?.isMuted ?? true,
    isSpeaking: p.isSpeaking,
    hasVideo: !!(camPub && !camPub.isMuted),
    isScreenSharing: !!(screenPub && !screenPub.isMuted) || !!(screenAudioPub && !screenAudioPub.isMuted),
  };
}

export const useVoiceStore = create<VoiceState>((set, get) => ({
  currentRoom: null,
  currentRoomName: null,
  isConnected: false,
  isMuted: false,
  isDeafened: false,
  isVideoOn: false,
  isScreenSharing: false,
  noiseSuppression: true,
  bgMode: 'disabled',
  _bgProcessor: null,
  participants: [],
  error: null,
  _room: null,
  isJoining: false,
  whispers: [],

  joinVoice: async (channelId, channelName) => {
    if (get().isJoining) {
      console.log("[voice] already joining, ignoring duplicate click");
      return;
    }
    set({ isJoining: true });
    console.log("[voice] joinVoice called:", channelId);
    const existing = get()._room;
    if (existing) {
      await get().leaveVoice();
    }

    set({ error: null });

    try {
      console.log("[voice] requesting token for:", channelId);
      const resp = await api.post<VoiceTokenResponse>('/voice/token', {
        room_name: channelId,
      });
      console.log("[voice] got token response:", resp);

      const room = new Room({
        adaptiveStream: true,
        dynacast: true,
        publishDefaults: {
          red: false,
          dtx: false,
          audioPreset: AudioPresets.speech,
        },
        audioCaptureDefaults: {
          channelCount: 1,
          echoCancellation: true,
          noiseSuppression: true,
          autoGainControl: true,
        },
      });

      room.on(RoomEvent.Connected, () => {
        set({ isConnected: true, error: null });
        syncParticipants();
      });

      room.on(RoomEvent.Disconnected, (reason?: any) => {
        console.warn('[voice] disconnected, reason:', reason);
        set({
          isConnected: false,
          currentRoom: null,
          currentRoomName: null,
          participants: [],
          isMuted: false,
          isDeafened: false,
          isVideoOn: false,
          isScreenSharing: false,
        });
      });

      room.on(RoomEvent.Reconnecting, () => {
        console.warn('[voice] reconnecting to LiveKit...');
        set({ error: 'Reconnecting...' });
      });

      room.on(RoomEvent.Reconnected, () => {
        console.log('[voice] reconnected to LiveKit');
        set({ error: null });
        syncParticipants();
      });

      room.on(RoomEvent.SignalConnected, () => {
        console.log('[voice] signal channel connected');
      });

      room.on(RoomEvent.ConnectionQualityChanged, (_quality: any, participant: Participant) => {
        if (participant.identity === room.localParticipant.identity) {
          // if quality drops to poor, log it
          if (_quality === 'poor') {
            console.warn('[voice] poor connection quality');
          }
        }
      });

      room.on(RoomEvent.ParticipantConnected, () => {
        syncParticipants();
      });

      room.on(RoomEvent.ParticipantDisconnected, () => {
        syncParticipants();
      });

      room.on(RoomEvent.TrackMuted, () => {
        syncParticipants();
      });

      room.on(RoomEvent.TrackUnmuted, () => {
        syncParticipants();
      });

      room.on(RoomEvent.TrackPublished, () => syncParticipants());
      room.on(RoomEvent.TrackUnpublished, () => syncParticipants());
      room.on(RoomEvent.TrackSubscribed, () => syncParticipants());
      room.on(RoomEvent.TrackUnsubscribed, () => syncParticipants());
      room.on(RoomEvent.LocalTrackPublished, () => syncParticipants());
      room.on(RoomEvent.LocalTrackUnpublished, () => syncParticipants());

      room.on(
        RoomEvent.ActiveSpeakersChanged,
        (speakers: Participant[]) => {
          set((state) => ({
            participants: state.participants.map((p) => ({
              ...p,
              isSpeaking: speakers.some((s) => s.identity === p.identity),
            })),
          }));
        },
      );

      function syncParticipants() {
        const r = get()._room;
        if (!r) return;

        const all: Participant[] = [
          r.localParticipant,
          ...Array.from(r.remoteParticipants.values()),
        ];
        set({ participants: all.map(participantToVoice) });
      }

      await room.connect(resp.livekit_url, resp.token);

      // Set state immediately so the UI renders the voice panel
      set({
        _room: room,
        currentRoom: channelId,
        currentRoomName: channelName,
        error: null,
      });

      // Mute mic by default on join (best effort; don't fail join if this errors)
      try {
        await room.localParticipant.setMicrophoneEnabled(false);
        set({ isMuted: true });
      } catch (micErr) {
        console.warn("[voice] failed to set initial mute state:", micErr);
      }

      // Restore saved device preferences
      const prefs = loadDevicePrefs();
      try {
        if (prefs.audioInput) await room.switchActiveDevice('audioinput', prefs.audioInput);
        if (prefs.audioOutput) await room.switchActiveDevice('audiooutput', prefs.audioOutput);
        if (prefs.videoInput) await room.switchActiveDevice('videoinput', prefs.videoInput);
      } catch (devErr) {
        console.warn("[voice] failed to restore saved devices:", devErr);
      }

      set({ isJoining: false });

      // Notify WebSocket peers
      const wsSend = useWebSocketStore.getState().send;
      wsSend({
        type: 'VOICE_JOIN',
        data: { room_name: channelId },
      });
    } catch (err) {
      set({
        error: err instanceof Error ? err.message : 'Failed to join voice',
      });
    }
  },

  leaveVoice: async () => {
    const room = get()._room;
    if (!room) return;

    const channelId = get().currentRoom;

    const wsSend = useWebSocketStore.getState().send;
    if (channelId) {
      wsSend({
        type: 'VOICE_LEAVE',
        data: { room_name: channelId },
      });
    }

    await room.disconnect();
    set({
      _room: null,
      currentRoom: null,
      currentRoomName: null,
      isConnected: false,
      isMuted: false,
      isDeafened: false,
      isVideoOn: false,
      isScreenSharing: false,
      participants: [],
      error: null,
    });
  },

  toggleMute: () => {
    const room = get()._room;
    if (!room) return;

    const newMuted = !get().isMuted;
    room.localParticipant.setMicrophoneEnabled(!newMuted);
    set({ isMuted: newMuted });

    // If un-deafening while unmuting, also restore speaker
    if (!newMuted && get().isDeafened) {
      set({ isDeafened: false });
    }

    const wsSend = useWebSocketStore.getState().send;
    wsSend({
      type: 'VOICE_MUTE',
      data: { room_name: get().currentRoom, muted: newMuted },
    });
  },

  toggleDeafen: () => {
    const room = get()._room;
    if (!room) return;

    const newDeafened = !get().isDeafened;

    if (newDeafened) {
      // Deafen: mute mic too
      room.localParticipant.setMicrophoneEnabled(false);
      set({ isDeafened: true, isMuted: true });
    } else {
      set({ isDeafened: false });
    }

    // Mute/unmute all remote audio tracks
    room.remoteParticipants.forEach((p: RemoteParticipant) => {
      p.audioTrackPublications.forEach((pub: TrackPublication) => {
        if (pub.track && pub.kind === Track.Kind.Audio) {
          (pub.track as RemoteAudioTrack).setVolume(newDeafened ? 0 : 1.0);
        }
      });
    });

    const wsSend = useWebSocketStore.getState().send;
    wsSend({
      type: 'VOICE_DEAFEN',
      data: { room_name: get().currentRoom, deafened: newDeafened },
    });
  },

  toggleVideo: async () => {
    const room = get()._room;
    if (!room) return;

    const newState = !get().isVideoOn;
    await room.localParticipant.setCameraEnabled(newState);
    set({ isVideoOn: newState });

    if (newState && get().bgMode !== 'disabled') {
      const mode = get().bgMode;
      // re-trigger setBgMode to re-attach processor to the new track
      await get().setBgMode(mode);
    }
  },

  setBgMode: async (mode: 'disabled' | 'background-blur' | 'virtual-background', imagePath?: string) => {
    const room = get()._room;
    if (!room) return;

    let processor = get()._bgProcessor;
    if (!processor && mode !== 'disabled') {
      try {
        processor = BackgroundProcessor({ mode: 'background-blur', blurRadius: 10 });
        set({ _bgProcessor: processor });
      } catch (err) {
        console.warn('Failed to initialize BackgroundProcessor:', err);
        return;
      }
    }

    if (processor) {
      try {
        if (mode === 'disabled') {
          await processor.switchTo({ mode: 'disabled' });
        } else if (mode === 'background-blur') {
          await processor.switchTo({ mode: 'background-blur', blurRadius: 10 });
        } else if (mode === 'virtual-background') {
          await processor.switchTo({ mode: 'virtual-background', imagePath: imagePath || 'https://images.unsplash.com/photo-1579546929518-9e396f3cc809' });
        }

        // attach to camera track if it's running
        const pub = room.localParticipant.getTrackPublication(Track.Source.Camera);
        const track = pub?.track as any;
        if (track && track.setProcessor) {
          await track.setProcessor(processor);
        }
        set({ bgMode: mode });
      } catch (err) {
        console.error('[voice] setBgMode failed:', err);
        alert(`Failed to set background effect: ${err instanceof Error ? err.message : err}`);
        set({ bgMode: 'disabled' });
      }
    } else {
      set({ bgMode: mode });
    }
  },

  toggleScreenShare: async () => {
    const room = get()._room;
    if (!room) return;

    const newState = !get().isScreenSharing;
    try {
      await room.localParticipant.setScreenShareEnabled(newState, { audio: true });
      set({ isScreenSharing: newState });
    } catch (err: any) {
      // user cancelled the browser picker — not an error
      if (err?.name === 'AbortError' || err?.message?.includes('cancel')) return;
      console.error('[voice] screenshare toggle failed:', err);
      set({ error: err instanceof Error ? err.message : 'Screen share failed' });
    }
  },

  toggleNoiseSuppression: async () => {
    const room = get()._room;
    if (!room) return;

    const newState = !get().noiseSuppression;
    set({ noiseSuppression: newState });

    // Apply noise suppression + echo cancellation to local audio track
    try {
      const audioTrack = room.localParticipant.getTrackPublication(Track.Source.Microphone);
      if (audioTrack?.track?.mediaStreamTrack) {
        await audioTrack.track.mediaStreamTrack.applyConstraints({
          echoCancellation: newState,
          noiseSuppression: newState,
          autoGainControl: newState,
        } as MediaTrackConstraints);
      }
    } catch (err) {
      console.warn('[voice] failed to update audio processing:', err);
    }
  },

  _addWhisper: (w) => {
    set((state) => ({ whispers: [...state.whispers, w] }));
  },

  whisperTo: (targetUserId, targetUsername, message) => {
    const wsSend = useWebSocketStore.getState().send;
    wsSend({
      type: 'VOICE_WHISPER',
      data: {
        target_user_id: targetUserId,
        target_username: targetUsername,
        message: message || '',
      },
    });
  },

  dismissWhisper: (timestamp) => {
    set((state) => ({
      whispers: state.whispers.filter((w) => w.timestamp !== timestamp),
    }));
  },
}));