clarity-helper/hooks/useLiveSession.ts

import { useState, useRef, useEffect, useCallback } from 'react';
import { GoogleGenAI, LiveServerMessage, Modality } from '@google/genai';
import { Scenario, UserSettings, Message, AnalysisResult } from '../types';
import { SYSTEM_ANALYSIS_PROMPT, CBT_ANALYSIS_SCHEMA } from '../constants';
import { floatTo16BitPCM, encodeBase64, decodeBase64, decodeAudioData } from '../services/audioUtils';

export type SessionStatus = 'connecting' | 'active' | 'analyzing' | 'finished';

export function useLiveSession(scenario: Scenario, voiceSettings: UserSettings) {
  const [status, setStatus] = useState<SessionStatus>('connecting');
  const [isMuted, setIsMuted] = useState(false);
  const [audioLevel, setAudioLevel] = useState(0);
  const [aiAudioLevel, setAiAudioLevel] = useState(0);
  const [isAiSpeaking, setIsAiSpeaking] = useState(false);
  const [isUserSpeaking, setIsUserSpeaking] = useState(false);
  const [messages, setMessages] = useState<Message[]>([]);
  const [isSilenceDetected, setIsSilenceDetected] = useState(false);
  
  // Insights state
  const [liveInsight, setLiveInsight] = useState<AnalysisResult | null>(null);
  const [isAnalyzingTurn, setIsAnalyzingTurn] = useState(false);

  // Refs for audio processing
  const sessionPromiseRef = useRef<Promise<any> | null>(null);
  const inputAudioCtxRef = useRef<AudioContext | null>(null);
  const outputAudioCtxRef = useRef<AudioContext | null>(null);
  const aiAnalyserRef = useRef<AnalyserNode | null>(null);
  const nextStartTimeRef = useRef<number>(0);
  const audioSourcesRef = useRef<Set<AudioBufferSourceNode>>(new Set());
  const streamRef = useRef<MediaStream | null>(null);
  const analysisChatRef = useRef<any>(null);
  const silenceCheckTimerRef = useRef<NodeJS.Timeout | null>(null);
  const aiLevelUpdateTimerRef = useRef<number | null>(null);

  // Buffers for transcription
  const currentInputTranscription = useRef('');
  const currentOutputTranscription = useRef('');

  const stopAllAudio = useCallback(() => {
    audioSourcesRef.current.forEach(source => {
      try { source.stop(); } catch (e) {}
    });
    audioSourcesRef.current.clear();
    nextStartTimeRef.current = 0;
    setAiAudioLevel(0);
  }, []);

  const toggleMute = useCallback(() => setIsMuted(prev => !prev), []);

  useEffect(() => {
    const ai = new GoogleGenAI({ apiKey: process.env.API_KEY });

    // Setup Live Insights Chat
    analysisChatRef.current = ai.chats.create({
      model: 'gemini-3-flash-preview',
      config: {
        systemInstruction: `${SYSTEM_ANALYSIS_PROMPT}\nSCENARIO: ${scenario.title}\nFOCUS: ${scenario.cbtFocus.join(', ')}`,
        responseMimeType: 'application/json',
        responseSchema: CBT_ANALYSIS_SCHEMA,
      }
    });

    const updateAiLevel = () => {
      if (aiAnalyserRef.current && isAiSpeaking) {
        const dataArray = new Uint8Array(aiAnalyserRef.current.frequencyBinCount);
        aiAnalyserRef.current.getByteFrequencyData(dataArray);
        let sum = 0;
        for (let i = 0; i < dataArray.length; i++) sum += dataArray[i];
        const avg = sum / dataArray.length;
        setAiAudioLevel(avg * 0.8); // Scale for visualizer
      } else {
        setAiAudioLevel(0);
      }
      aiLevelUpdateTimerRef.current = requestAnimationFrame(updateAiLevel);
    };

    const connect = async () => {
      try {
        inputAudioCtxRef.current = new (window.AudioContext || (window as any).webkitAudioContext)({ sampleRate: 16000 });
        outputAudioCtxRef.current = new (window.AudioContext || (window as any).webkitAudioContext)({ sampleRate: 24000 });
        
        aiAnalyserRef.current = outputAudioCtxRef.current.createAnalyser();
        aiAnalyserRef.current.fftSize = 256;
        aiAnalyserRef.current.connect(outputAudioCtxRef.current.destination);
        
        updateAiLevel();

        const constraints: MediaStreamConstraints = { 
          audio: voiceSettings.preferredAudioDeviceId 
            ? { deviceId: { exact: voiceSettings.preferredAudioDeviceId } } 
            : true 
        };

        const stream = await navigator.mediaDevices.getUserMedia(constraints);
        streamRef.current = stream;

        // Silence detection logic: If after 5 seconds of active session, 
        // the max audio level seen is very low, flag it.
        let maxLevelSeen = 0;
        silenceCheckTimerRef.current = setTimeout(() => {
          if (maxLevelSeen < 1.0) { // Threshold for silence
            setIsSilenceDetected(true);
          }
        }, 6000);

        const sessionPromise = ai.live.connect({
          model: 'gemini-2.5-flash-native-audio-preview-09-2025',
          callbacks: {
            onopen: () => {
              setStatus('active');
              const source = inputAudioCtxRef.current!.createMediaStreamSource(stream);
              const scriptProcessor = inputAudioCtxRef.current!.createScriptProcessor(4096, 1, 1);
              
              scriptProcessor.onaudioprocess = (e) => {
                const inputData = e.inputBuffer.getChannelData(0);
                
                // Volume calculation
                let sum = 0;
                for (let i = 0; i < inputData.length; i++) sum += inputData[i] * inputData[i];
                const level = Math.sqrt(sum / inputData.length) * 100;
                setAudioLevel(level);
                setIsUserSpeaking(level > 2.0);
                
                if (level > maxLevelSeen) maxLevelSeen = level;
                if (level > 2.0 && isSilenceDetected) setIsSilenceDetected(false);

                if (isMuted) return;

                const pcm16 = floatTo16BitPCM(inputData);
                const base64 = encodeBase64(new Uint8Array(pcm16));
                
                sessionPromise.then(session => {
                  session.sendRealtimeInput({ 
                    media: { data: base64, mimeType: 'audio/pcm;rate=16000' } 
                  });
                });
              };

              source.connect(scriptProcessor);
              scriptProcessor.connect(inputAudioCtxRef.current!.destination);
            },
            onmessage: async (message: LiveServerMessage) => {
              // Transcriptions
              if (message.serverContent?.outputTranscription) {
                currentOutputTranscription.current += message.serverContent.outputTranscription.text;
              } else if (message.serverContent?.inputTranscription) {
                currentInputTranscription.current += message.serverContent.inputTranscription.text;
              }

              if (message.serverContent?.turnComplete) {
                const userText = currentInputTranscription.current;
                const modelText = currentOutputTranscription.current;

                if (userText) {
                  setMessages(prev => [...prev, { id: Date.now().toString(), role: 'user', text: userText, timestamp: new Date() }]);
                  
                  if (analysisChatRef.current && modelText) {
                    setIsAnalyzingTurn(true);
                    analysisChatRef.current.sendMessage({ message: `User: "${userText}"\nAI: "${modelText}"` })
                      .then((res: any) => {
                        try { if (res.text) setLiveInsight(JSON.parse(res.text)); }
                        catch(e) {} finally { setIsAnalyzingTurn(false); }
                      });
                  }
                }
                
                if (modelText) {
                  setMessages(prev => [...prev, { id: (Date.now()+1).toString(), role: 'model', text: modelText, timestamp: new Date() }]);
                }

                currentInputTranscription.current = '';
                currentOutputTranscription.current = '';
              }

              // Audio Output
              const audioBase64 = message.serverContent?.modelTurn?.parts[0]?.inlineData?.data;
              if (audioBase64) {
                setIsAiSpeaking(true);
                const audioCtx = outputAudioCtxRef.current!;
                nextStartTimeRef.current = Math.max(nextStartTimeRef.current, audioCtx.currentTime);
                
                const buffer = await decodeAudioData(decodeBase64(audioBase64), audioCtx, 24000, 1);
                const source = audioCtx.createBufferSource();
                source.buffer = buffer;
                source.connect(aiAnalyserRef.current!);
                
                source.onended = () => {
                  audioSourcesRef.current.delete(source);
                  if (audioSourcesRef.current.size === 0) setIsAiSpeaking(false);
                };
                
                source.start(nextStartTimeRef.current);
                nextStartTimeRef.current += buffer.duration;
                audioSourcesRef.current.add(source);
              }

              if (message.serverContent?.interrupted) {
                stopAllAudio();
                setIsAiSpeaking(false);
              }
            },
            onclose: () => console.debug("Session closed"),
            onerror: (e) => console.error("Session error", e)
          },
          config: {
            responseModalities: [Modality.AUDIO],
            speechConfig: {
              voiceConfig: { prebuiltVoiceConfig: { voiceName: voiceSettings.voiceName } }
            },
            systemInstruction: `
              ${scenario.systemInstruction}
              
              VOCAL REALISM PROTOCOL:
              - Speak with natural human patterns (fillers like "um", "uh", "well...").
              - Use brief hesitations when thinking or emotionally activated.
              - Vary pace based on character emotion.
              - Strictly in character; no clinical meta-feedback.
            `,
            inputAudioTranscription: {},
            outputAudioTranscription: {}
          }
        });

        sessionPromiseRef.current = sessionPromise;
      } catch (err) {
        console.error("Connect failed", err);
      }
    };

    connect();

    return () => {
      stopAllAudio();
      if (silenceCheckTimerRef.current) clearTimeout(silenceCheckTimerRef.current);
      if (aiLevelUpdateTimerRef.current) cancelAnimationFrame(aiLevelUpdateTimerRef.current);
      if (streamRef.current) streamRef.current.getTracks().forEach(t => t.stop());
      if (inputAudioCtxRef.current) inputAudioCtxRef.current.close();
      if (outputAudioCtxRef.current) outputAudioCtxRef.current.close();
      if (sessionPromiseRef.current) sessionPromiseRef.current.then(s => s.close());
    };
  }, [scenario, voiceSettings, stopAllAudio, isMuted]);

  return {
    status,
    setStatus,
    isMuted,
    toggleMute,
    audioLevel,
    aiAudioLevel,
    isAiSpeaking,
    isUserSpeaking,
    messages,
    liveInsight,
    isAnalyzingTurn,
    isSilenceDetected,
    sessionPromiseRef,
    startSession: () => {
      if (sessionPromiseRef.current) {
        sessionPromiseRef.current.then(session => {
          // We send a "hidden" text input to the model to trigger its initial message
          // The Live API supports sending text turns in the clientContent
          session.send({
            clientContent: {
              turns: [{
                role: 'user',
                parts: [{ text: "Please start the scenario now. Greet me as your character." }]
              }],
              turnComplete: true
            }
          });
        });
      }
    }
  };
}