clarity-helper/components/ChatSession.tsx

import React, { useState, useRef, useEffect, useCallback } from 'react';
import { Scenario, Message, AnalysisResult } from '../types';
import { GoogleGenAI, Chat, GenerateContentResponse } from "@google/genai";
import { BrainCircuit, ChevronLeft, Flag, Loader2 } from 'lucide-react';
import { SYSTEM_ANALYSIS_PROMPT, SYSTEM_SUMMARY_ANALYSIS_PROMPT, CBT_ANALYSIS_SCHEMA } from '../constants';

// Sub-components
import ChatMessage from './ChatMessage';
import ChatInput from './ChatInput';
import LiveInsightPanel from './LiveInsightPanel';
import SessionReport from './SessionReport';

interface ChatSessionProps {
  scenario: Scenario;
  onExit: () => void;
}

const ChatSession: React.FC<ChatSessionProps> = ({ scenario, onExit }) => {
  const [messages, setMessages] = useState<Message[]>([]);
  const [isLoading, setIsLoading] = useState(false);
  const [status, setStatus] = useState<'active' | 'analyzing' | 'finished'>('active');
  
  // Clinical Feedback State
  const [liveInsight, setLiveInsight] = useState<AnalysisResult | null>(null);
  const [isAnalyzingTurn, setIsAnalyzingTurn] = useState(false);
  const [finalReport, setFinalReport] = useState<AnalysisResult | null>(null);
  
  const chatRef = useRef<Chat | null>(null);
  const analysisRef = useRef<Chat | null>(null); 
  const messagesEndRef = useRef<HTMLDivElement>(null);

  const scrollToBottom = useCallback(() => {
    messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' });
  }, []);

  useEffect(() => {
    scrollToBottom();
  }, [messages, scrollToBottom]);

  // Session Initialization
  useEffect(() => {
    const ai = new GoogleGenAI({ apiKey: process.env.API_KEY });
    
    // Roleplay Chat Setup
    const roleplaySystemInstruction = `
      ${scenario.systemInstruction}
      
      CRITICAL ROLEPLAY PROTOCOL:
      1. Stay strictly in character as the persona defined in the scenario.
      2. NEVER provide clinical feedback, advice, or meta-commentary about the user's progress.
      3. NEVER use phrases like "Great job holding your boundary" or "Feedback:".
      4. Do not congratulate the user on using CBT techniques; simply react as the character would.
      5. If the user successfully concludes the scenario, end the dialogue naturally.
    `;

    chatRef.current = ai.chats.create({
      model: 'gemini-3-flash-preview',
      config: { systemInstruction: roleplaySystemInstruction },
    });

    // Live Analysis Chat Setup
    const analysisSystemInstruction = `${SYSTEM_ANALYSIS_PROMPT}
SCENARIO CONTEXT:
Title: ${scenario.title}
Description: ${scenario.description}
CBT Focus Areas: ${scenario.cbtFocus.join(', ')}`;

    analysisRef.current = ai.chats.create({
      model: 'gemini-3-flash-preview',
      config: {
        systemInstruction: analysisSystemInstruction,
        responseMimeType: 'application/json',
        responseSchema: CBT_ANALYSIS_SCHEMA,
      }
    });

    // Initial Message
    setMessages([{
      id: 'init',
      role: 'model',
      text: scenario.initialMessage,
      timestamp: new Date()
    }]);
    
    setLiveInsight(null);
  }, [scenario]);

  const handleSend = useCallback(async (text: string) => {
    if (!chatRef.current || isLoading) return;

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

    setMessages(prev => [...prev, userMsg]);
    setIsLoading(true);

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

      const streamResponse = await chatRef.current.sendMessageStream({ message: text });
      let fullResponseText = '';

      for await (const chunk of streamResponse) {
        const c = chunk as GenerateContentResponse;
        const textChunk = c.text || "";
        fullResponseText += textChunk;
        
        setMessages(prev => prev.map(m => 
          m.id === botMsgId ? { ...m, text: fullResponseText } : m
        ));
      }

      // Live Turn Analysis
      if (analysisRef.current) {
        setIsAnalyzingTurn(true);
        const turnContext = `User said: "${userMsg.text}". Roleplay partner said: "${fullResponseText}".`;
        
        analysisRef.current.sendMessage({ message: turnContext }).then(res => {
          try {
             if (res.text) setLiveInsight(JSON.parse(res.text));
          } catch(e) { console.error("Live analysis parse error", e); } finally {
             setIsAnalyzingTurn(false);
          }
        });
      }

    } catch (error) {
      console.error("Chat transmission failed", error);
      setMessages(prev => prev.filter(m => m.text !== ''));
    } finally {
      setIsLoading(false);
    }
  }, [isLoading]);

  const handleFinishSession = async () => {
    if (messages.length < 2) return onExit();

    setStatus('analyzing');
    
    try {
      const ai = new GoogleGenAI({ apiKey: process.env.API_KEY });
      const transcript = 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${transcript}` }] }]
      });

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

  if (status === 'analyzing') {
    return (
      <div className="flex flex-col items-center justify-center min-h-[500px] h-full bg-slate-50 rounded-3xl p-8 border border-slate-200 w-full max-w-2xl mx-auto animate-in fade-in duration-500">
        <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 text-slate-800 mb-2">Compiling Report</h2>
        <p className="text-slate-500 text-center max-w-sm">Generating evidence-based synthesis of your cognitive and behavioral patterns during this session...</p>
      </div>
    );
  }

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

  return (
    <div className="flex flex-col h-[calc(100vh-100px)] max-w-6xl mx-auto w-full gap-4 px-2">
      {/* Interaction Header */}
      <div className="bg-white p-4 rounded-xl shadow-sm border border-slate-200 flex justify-between items-center">
        <div className="flex items-center gap-3">
          <button onClick={onExit} className="p-2 hover:bg-slate-50 rounded-lg text-slate-400 hover:text-teal-600 transition-colors">
            <ChevronLeft size={20} />
          </button>
          <div>
            <h2 className="font-bold text-slate-800 text-lg leading-tight">{scenario.title}</h2>
            <div className="flex gap-2">
               {scenario.cbtFocus.slice(0, 3).map(f => (
                 <span key={f} className="text-[10px] font-bold text-teal-600 uppercase tracking-tight">{f}</span>
               ))}
            </div>
          </div>
        </div>
        <button 
          onClick={handleFinishSession} 
          className="flex items-center gap-2 px-4 py-2 rounded-lg bg-teal-50 text-teal-700 hover:bg-teal-600 hover:text-white text-[10px] font-black uppercase tracking-widest transition-all border border-teal-100"
        >
          <Flag size={14} />
          Complete Interaction
        </button>
      </div>

      <div className="flex flex-col lg:flex-row gap-4 h-full overflow-hidden">
        {/* Core Chat Module */}
        <div className="flex-grow flex flex-col bg-white rounded-xl shadow-sm border border-slate-200 overflow-hidden">
          <div className="flex-grow overflow-y-auto p-6 space-y-6 scrollbar-hide">
            {messages.map((msg) => (
              <ChatMessage key={msg.id} message={msg} isLoading={isLoading && msg.id === messages[messages.length-1].id} />
            ))}
            <div ref={messagesEndRef} />
          </div>

          <ChatInput onSend={handleSend} disabled={isLoading} />
        </div>

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

export default React.memo(ChatSession);