clarity-helper/components/VoiceSession.tsx

import React, { useState, useRef, useEffect, useCallback } from 'react';
import { Scenario, UserSettings, Message, AnalysisResult } from '../types';
import { GoogleGenAI, LiveServerMessage, Modality } from '@google/genai';
import { Mic, MicOff, PhoneOff, MessageSquare, BrainCircuit, Flag, Loader2 } from 'lucide-react';
import { SYSTEM_ANALYSIS_PROMPT, SYSTEM_SUMMARY_ANALYSIS_PROMPT, CBT_ANALYSIS_SCHEMA } from '../constants';
import { floatTo16BitPCM, arrayBufferToBase64, base64ToUint8Array, decodeAudioData } from '../services/audioUtils';

// Sub-components
import AudioVisualizer from './AudioVisualizer';
import TranscriptPanel from './TranscriptPanel';
import LiveInsightPanel from './LiveInsightPanel';
import SessionReport from './SessionReport';

interface VoiceSessionProps {
  scenario: Scenario;
  onExit: () => void;
  voiceSettings: UserSettings;
}

// Manual encoding helper as per guidelines
function encode(bytes: Uint8Array) {
  let binary = '';
  const len = bytes.byteLength;
  for (let i = 0; i < len; i++) {
    binary += String.fromCharCode(bytes[i]);
  }
  return btoa(binary);
}

// Manual decoding helper as per guidelines
function decode(base64: string) {
  const binaryString = atob(base64);
  const len = binaryString.length;
  const bytes = new Uint8Array(len);
  for (let i = 0; i < len; i++) {
    bytes[i] = binaryString.charCodeAt(i);
  }
  return bytes;
}

const VoiceSession: React.FC<VoiceSessionProps> = ({ scenario, onExit, voiceSettings }) => {
  const [isMuted, setIsMuted] = useState(false);
  const [showTranscript, setShowTranscript] = useState(false);
  const [messages, setMessages] = useState<Message[]>([]);
  const [status, setStatus] = useState<'connecting' | 'active' | 'analyzing' | 'finished'>('connecting');
  const [audioLevel, setAudioLevel] = useState(0);
  const [isAiSpeaking, setIsAiSpeaking] = useState(false);

  // Analysis States
  const [liveInsight, setLiveInsight] = useState<AnalysisResult | null>(null);
  const [isAnalyzingTurn, setIsAnalyzingTurn] = useState(false);
  const [finalReport, setFinalReport] = useState<AnalysisResult | null>(null);

  // Refs for Audio & Live API
  const sessionPromiseRef = useRef<Promise<any> | null>(null);
  const inputAudioCtxRef = useRef<AudioContext | null>(null);
  const outputAudioCtxRef = useRef<AudioContext | 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);

  // Transcription buffers
  const currentInputTranscription = useRef('');
  const currentOutputTranscription = useRef('');

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

  const handleFinishSession = async () => {
    setStatus('analyzing');
    // Close live session
    if (sessionPromiseRef.current) {
      sessionPromiseRef.current.then(session => session.close());
    }
    
    try {
      const ai = new GoogleGenAI({ apiKey: process.env.API_KEY });
      const transcriptText = messages.map(m => `[${m.role.toUpperCase()}]: ${m.text}`).join('\n');
      
      const summaryPrompt = `${SYSTEM_SUMMARY_ANALYSIS_PROMPT}
SCENARIO CONTEXT: ${scenario.title} - ${scenario.description}`;

      const result = await ai.models.generateContent({
        model: 'gemini-3-flash-preview',
        config: {
          systemInstruction: summaryPrompt,
          responseMimeType: 'application/json',
          responseSchema: CBT_ANALYSIS_SCHEMA,
        },
        contents: [{ role: 'user', parts: [{ text: `SESSION TRANSCRIPT:\n${transcriptText}` }] }]
      });

      if (result.text) {
        setFinalReport(JSON.parse(result.text));
        setStatus('finished');
      } else {
        throw new Error("Analysis failed");
      }
    } catch (err) {
      console.error("Post-session analysis failure", err);
      onExit();
    }
  };

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

    // Setup Analysis Model for Live Insights
    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 init = async () => {
      try {
        inputAudioCtxRef.current = new (window.AudioContext || (window as any).webkitAudioContext)({ sampleRate: 16000 });
        outputAudioCtxRef.current = new (window.AudioContext || (window as any).webkitAudioContext)({ sampleRate: 24000 });
        
        const stream = await navigator.mediaDevices.getUserMedia({ audio: true });
        streamRef.current = stream;

        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);
                
                // Calculate simple audio level for visualizer
                let sum = 0;
                for (let i = 0; i < inputData.length; i++) sum += inputData[i] * inputData[i];
                setAudioLevel(Math.sqrt(sum / inputData.length) * 100);

                if (isMuted) return;

                const pcm16 = floatTo16BitPCM(inputData);
                const base64 = encode(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) => {
              // Handle 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) {
                  const userMsg: Message = { id: Date.now().toString(), role: 'user', text: userText, timestamp: new Date() };
                  setMessages(prev => [...prev, userMsg]);
                  
                  // Trigger Analysis
                  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 = '';
              }

              // Handle 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(decode(audioBase64), audioCtx, 24000, 1);
                const source = audioCtx.createBufferSource();
                source.buffer = buffer;
                source.connect(audioCtx.destination);
                
                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.log("Live session closed"),
            onerror: (e) => console.error("Live API error", e)
          },
          config: {
            responseModalities: [Modality.AUDIO],
            speechConfig: {
              voiceConfig: { prebuiltVoiceConfig: { voiceName: voiceSettings.voiceName } }
            },
            systemInstruction: `
              ${scenario.systemInstruction}
              
              VIRTUAL REALITY ROLEPLAY PROTOCOL:
              1. Stay strictly in character.
              2. DO NOT provide advice or analysis within the conversation.
              3. React naturally to the user's voice and emotional tone.
            `,
            inputAudioTranscription: {},
            outputAudioTranscription: {}
          }
        });

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

    init();

    return () => {
      stopAllAudio();
      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]);

  if (status === 'connecting') {
    return (
      <div className="flex flex-col items-center justify-center h-full text-slate-800 animate-pulse">
        <Loader2 size={48} className="animate-spin text-teal-600 mb-4" />
        <p className="font-bold uppercase tracking-widest text-xs">Opening Neural Bridge...</p>
      </div>
    );
  }

  if (status === 'analyzing') {
    return (
      <div className="flex flex-col items-center justify-center h-full text-slate-800">
        <div className="bg-white p-12 rounded-full shadow-lg border border-teal-100 mb-8 relative">
           <BrainCircuit size={64} className="text-teal-600 animate-pulse" />
           <div className="absolute inset-0 border-4 border-teal-500 border-t-transparent rounded-full animate-spin"></div>
        </div>
        <h2 className="text-2xl font-bold mb-2">Generating Report</h2>
        <p className="text-slate-500 text-center max-w-sm">Synthesizing vocal patterns and cognitive restructuring evidence...</p>
      </div>
    );
  }

  if (status === 'finished' && finalReport) {
    return <SessionReport analysis={finalReport} onExit={onExit} />;
  }

  return (
    <div className="flex flex-col lg:flex-row h-full gap-4 relative overflow-hidden">
      {/* Main Interaction Area */}
      <div className="flex-grow bg-slate-900 rounded-3xl overflow-hidden flex flex-col relative shadow-2xl border border-white/5">
        <div className="absolute inset-0 overflow-hidden pointer-events-none opacity-20">
          <div className={`absolute top-1/2 left-1/2 -translate-x-1/2 -translate-y-1/2 w-[80%] h-[80%] rounded-full blur-[100px] transition-all duration-1000 ${isAiSpeaking ? 'bg-blue-500 scale-110' : 'bg-teal-500 scale-100'}`}></div>
        </div>

        <div className="p-8 flex justify-between items-center relative z-10">
          <div className="flex items-center gap-4">
            <button onClick={onExit} className="p-3 bg-white/5 hover:bg-white/10 rounded-2xl text-white/50 hover:text-white transition-all">
              <PhoneOff size={20} />
            </button>
            <div>
              <h2 className="text-white font-bold text-lg">{scenario.title}</h2>
              <p className="text-teal-400 text-[10px] font-black uppercase tracking-widest">Active Roleplay Session</p>
            </div>
          </div>
          <button onClick={handleFinishSession} className="bg-teal-600 hover:bg-teal-700 text-white px-6 py-2.5 rounded-xl text-[10px] font-black uppercase tracking-widest transition-all shadow-lg shadow-teal-900/40 flex items-center gap-2">
            <Flag size={14} /> Finish
          </button>
        </div>

        <div className="flex-grow flex flex-col items-center justify-center p-12 relative z-10">
          <AudioVisualizer level={audioLevel} active={true} isAiSpeaking={isAiSpeaking} />
          
          <div className="mt-12 text-center">
            <p className={`text-sm transition-all duration-500 font-light ${isAiSpeaking ? 'text-blue-300' : 'text-teal-300'}`}>
              {isAiSpeaking ? `Vocalizing Response...` : (isMuted ? 'Microphone Muted' : 'Listening...')}
            </p>
            <p className="text-[10px] text-white/30 uppercase font-bold tracking-[0.3em] mt-4">Clinical Fidelity Level: 99.8%</p>
          </div>
        </div>

        <div className="p-8 flex justify-center gap-6 relative z-10 bg-black/20 backdrop-blur-md">
          <button 
            onClick={() => setIsMuted(!isMuted)}
            className={`p-5 rounded-full transition-all ${isMuted ? 'bg-rose-600 text-white' : 'bg-white/10 text-white hover:bg-white/20'}`}
          >
            {isMuted ? <MicOff size={24} /> : <Mic size={24} />}
          </button>
          <button 
            onClick={() => setShowTranscript(!showTranscript)}
            className={`p-5 rounded-full transition-all ${showTranscript ? 'bg-teal-600 text-white' : 'bg-white/10 text-white hover:bg-white/20'}`}
          >
            <MessageSquare size={24} />
          </button>
        </div>

        {showTranscript && <TranscriptPanel messages={messages} onClose={() => setShowTranscript(false)} />}
      </div>

      {/* Real-time Insights Module */}
      <LiveInsightPanel analysis={liveInsight} isAnalyzing={isAnalyzingTurn} />
    </div>
  );
};

export default React.memo(VoiceSession);