clarity-helper/hooks/useLiveSession.ts

import { useState, useRef, useEffect, useCallback } from 'react';
import { Scenario, UserSettings, Message, AnalysisResult } from '../types';
import { getOpenRouterApiKey } from '../services/apiConfig';
import { streamOpenRouterChat, analyzeTurnWithOpenRouter, synthesizeSpeechWithOpenRouter } from '../services/openRouterService';

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

export function useLiveSession(scenario: Scenario, voiceSettings: UserSettings) {
  const [status, setStatus] = useState<SessionStatus>('connecting');
  const [errorMessage, setErrorMessage] = useState<string | null>(null);
  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);

  // Audio Context & Nodes
  const isMutedRef = useRef(isMuted);
  useEffect(() => {
    isMutedRef.current = isMuted;
  }, [isMuted]);

  const inputAudioCtxRef = useRef<AudioContext | null>(null);
  const outputAudioCtxRef = useRef<AudioContext | null>(null);
  const aiAnalyserRef = useRef<AnalyserNode | null>(null);
  const userAnalyserRef = useRef<AnalyserNode | null>(null);
  const currentAudioSourceRef = useRef<AudioBufferSourceNode | null>(null);
  const streamRef = useRef<MediaStream | null>(null);
  const recognitionRef = useRef<any>(null);
  const aiLevelTimerRef = useRef<number | null>(null);
  const userLevelTimerRef = useRef<number | null>(null);
  const isProcessingTurnRef = useRef(false);

  const stopAllAudio = useCallback(() => {
    if (currentAudioSourceRef.current) {
      try { currentAudioSourceRef.current.stop(); } catch (e) {}
      currentAudioSourceRef.current = null;
    }
    if ('speechSynthesis' in window) {
      window.speechSynthesis.cancel();
    }
    setIsAiSpeaking(false);
    setAiAudioLevel(0);
  }, []);

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

  const speakText = useCallback(async (text: string) => {
    stopAllAudio();
    setIsAiSpeaking(true);

    try {
      // Attempt high-fidelity OpenRouter audio synthesis
      const arrayBuffer = await synthesizeSpeechWithOpenRouter(
        text,
        voiceSettings.voiceName.toLowerCase() === 'kore' ? 'nova' : 'alloy',
        voiceSettings.openRouterApiKey
      );

      const audioCtx = outputAudioCtxRef.current || new AudioContext();
      outputAudioCtxRef.current = audioCtx;
      if (audioCtx.state === 'suspended') await audioCtx.resume();

      const buffer = await audioCtx.decodeAudioData(arrayBuffer);
      const source = audioCtx.createBufferSource();
      source.buffer = buffer;

      if (!aiAnalyserRef.current) {
        aiAnalyserRef.current = audioCtx.createAnalyser();
        aiAnalyserRef.current.fftSize = 256;
      }
      source.connect(aiAnalyserRef.current);
      aiAnalyserRef.current.connect(audioCtx.destination);

      source.onended = () => {
        setIsAiSpeaking(false);
        setAiAudioLevel(0);
      };

      currentAudioSourceRef.current = source;
      source.start(0);
    } catch (e) {
      console.warn("OpenRouter TTS fallback to Web Speech Synthesis", e);
      // Fallback to browser SpeechSynthesis
      if ('speechSynthesis' in window) {
        const utterance = new SpeechSynthesisUtterance(text);
        utterance.rate = 1.0;
        utterance.pitch = 1.0;
        utterance.onstart = () => setIsAiSpeaking(true);
        utterance.onend = () => {
          setIsAiSpeaking(false);
          setAiAudioLevel(0);
        };
        utterance.onerror = () => {
          setIsAiSpeaking(false);
          setAiAudioLevel(0);
        };
        window.speechSynthesis.speak(utterance);
      } else {
        setIsAiSpeaking(false);
      }
    }
  }, [stopAllAudio, voiceSettings.voiceName, voiceSettings.openRouterApiKey]);

  // Turn Processing Handler
  const handleUserTurn = useCallback(async (userText: string, currentHistory: Message[]) => {
    if (isProcessingTurnRef.current || !userText.trim()) return;
    isProcessingTurnRef.current = true;

    stopAllAudio();

    const userMsg: Message = {
      id: Date.now().toString(),
      role: 'user',
      text: userText.trim(),
      timestamp: new Date()
    };

    const updatedHistory = [...currentHistory, userMsg];
    setMessages(updatedHistory);

    const botMsgId = (Date.now() + 1).toString();
    const botMsg: Message = {
      id: botMsgId,
      role: 'model',
      text: '',
      timestamp: new Date()
    };
    setMessages([...updatedHistory, botMsg]);

    const roleplaySystemInstruction = `
      ${scenario.systemInstruction}
      
      [CONTEXT]
      You are participating in a clinical CBT roleplay simulation. 
      The instructions above define your character, the scenario, and your internal CBT logic.
      
      [TASK]
      Respond to the user naturally in character to advance the roleplay scenario.
      
      [CONSTRAINTS]
      - CRITICAL: Stay strictly in character as the persona defined in the scenario.
      - CRITICAL: NEVER provide clinical feedback, advice, or meta-commentary about the user's progress.
      - CRITICAL: Keep dialogue natural and conversational (1-3 sentences per turn).
      - If the user successfully concludes the scenario, end the dialogue naturally.
    `;

    try {
      let accumulatedText = '';
      await streamOpenRouterChat(
        updatedHistory,
        roleplaySystemInstruction,
        {
          onChunk: (chunk) => {
            accumulatedText += chunk;
            setMessages(prev => prev.map(m =>
              m.id === botMsgId ? { ...m, text: accumulatedText } : m
            ));
          },
          onComplete: (fullText) => {
            speakText(fullText);

            // Parallel turn supervisor analysis
            setIsAnalyzingTurn(true);
            const turnContext = `User said: "${userMsg.text}". Roleplay partner said: "${fullText}".`;
            analyzeTurnWithOpenRouter(
              turnContext,
              scenario.title,
              scenario.description,
              scenario.cbtFocus,
              {
                apiKey: voiceSettings.openRouterApiKey,
                model: voiceSettings.supervisorModel
              }
            ).then(res => {
              setLiveInsight(res);
            }).catch(e => {
              console.error("Voice turn analysis error", e);
            }).finally(() => {
              setIsAnalyzingTurn(false);
              isProcessingTurnRef.current = false;
            });
          },
          onError: (err) => {
            console.error("OpenRouter stream failed", err);
            isProcessingTurnRef.current = false;
          }
        },
        {
          apiKey: voiceSettings.openRouterApiKey,
          model: voiceSettings.roleplayModel
        }
      );
    } catch (err) {
      console.error("User turn execution failed", err);
      isProcessingTurnRef.current = false;
    }
  }, [scenario, speakText, stopAllAudio, voiceSettings.openRouterApiKey, voiceSettings.roleplayModel, voiceSettings.supervisorModel]);

  // Main setup effect
  useEffect(() => {
    const apiKey = getOpenRouterApiKey(voiceSettings.openRouterApiKey);
    if (!apiKey) {
      setStatus('error');
      setErrorMessage('OpenRouter API Key is missing. Please set your API key in Settings.');
      return;
    }

    const initAudioAndSpeech = async () => {
      try {
        setStatus('connecting');
        setErrorMessage(null);

        // WebAudio Input Context
        inputAudioCtxRef.current = new (window.AudioContext || (window as any).webkitAudioContext)();
        outputAudioCtxRef.current = new (window.AudioContext || (window as any).webkitAudioContext)();

        userAnalyserRef.current = inputAudioCtxRef.current.createAnalyser();
        userAnalyserRef.current.fftSize = 256;

        aiAnalyserRef.current = outputAudioCtxRef.current.createAnalyser();
        aiAnalyserRef.current.fftSize = 256;
        aiAnalyserRef.current.connect(outputAudioCtxRef.current.destination);

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

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

        const micSource = inputAudioCtxRef.current.createMediaStreamSource(stream);
        micSource.connect(userAnalyserRef.current);

        // Visualizer Level Tracking Loops
        const updateUserLevel = () => {
          if (userAnalyserRef.current && !isMutedRef.current) {
            const dataArray = new Uint8Array(userAnalyserRef.current.frequencyBinCount);
            userAnalyserRef.current.getByteFrequencyData(dataArray);
            let sum = 0;
            for (let i = 0; i < dataArray.length; i++) sum += dataArray[i];
            const avg = sum / dataArray.length;
            const lvl = Math.min(100, avg * 1.5);
            setAudioLevel(lvl);
            setIsUserSpeaking(lvl > 8);
          } else {
            setAudioLevel(0);
            setIsUserSpeaking(false);
          }
          userLevelTimerRef.current = requestAnimationFrame(updateUserLevel);
        };
        updateUserLevel();

        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);
          } else {
            setAiAudioLevel(0);
          }
          aiLevelTimerRef.current = requestAnimationFrame(updateAiLevel);
        };
        updateAiLevel();

        // Browser Speech Recognition for user speech
        const SpeechRec = (window as any).SpeechRecognition || (window as any).webkitSpeechRecognition;
        if (SpeechRec) {
          const rec = new SpeechRec();
          rec.continuous = true;
          rec.interimResults = false;
          rec.lang = 'en-US';

          rec.onresult = (event: any) => {
            if (isMutedRef.current || isProcessingTurnRef.current) return;
            const last = event.results.length - 1;
            const transcript = event.results[last][0].transcript;
            if (transcript && transcript.trim().length > 0) {
              setMessages(current => {
                handleUserTurn(transcript, current);
                return current;
              });
            }
          };

          rec.onerror = (e: any) => {
            if (e.error !== 'no-speech') console.warn("Speech recognition notice:", e.error);
          };

          rec.onend = () => {
            // Auto restart recognition while active
            if (status === 'active' || status === 'connecting') {
              try { rec.start(); } catch (e) {}
            }
          };

          try {
            rec.start();
            recognitionRef.current = rec;
          } catch (e) {}
        }

        setStatus('active');
      } catch (err: any) {
        console.error("Live session initialization failed", err);
        setStatus('error');
        setErrorMessage(err?.message || 'Could not access audio hardware. Please check microphone permissions.');
      }
    };

    initAudioAndSpeech();

    return () => {
      stopAllAudio();
      if (userLevelTimerRef.current) cancelAnimationFrame(userLevelTimerRef.current);
      if (aiLevelTimerRef.current) cancelAnimationFrame(aiLevelTimerRef.current);
      if (recognitionRef.current) {
        try { recognitionRef.current.stop(); } catch (e) {}
      }
      if (streamRef.current) streamRef.current.getTracks().forEach(t => t.stop());
      if (inputAudioCtxRef.current) inputAudioCtxRef.current.close();
      if (outputAudioCtxRef.current) outputAudioCtxRef.current.close();
    };
  }, [scenario, voiceSettings.preferredAudioDeviceId, voiceSettings.openRouterApiKey, stopAllAudio, handleUserTurn, status, isAiSpeaking]);

  return {
    status,
    setStatus,
    errorMessage,
    isMuted,
    toggleMute,
    audioLevel,
    aiAudioLevel,
    isAiSpeaking,
    isUserSpeaking,
    messages,
    liveInsight,
    isAnalyzingTurn,
    isSilenceDetected,
    startSession: () => {
      const initialBotMsg: Message = {
        id: 'init',
        role: 'model',
        text: scenario.initialMessage,
        timestamp: new Date()
      };
      setMessages([initialBotMsg]);
      speakText(scenario.initialMessage);
    }
  };
}