chatbot.py

import os
import logging
import json
import torch
import re
from typing import List, Dict, Any, Optional, Union
from datetime import datetime
from pydantic import BaseModel, Field
import tempfile

# Model imports
from transformers import (
    pipeline, 
    AutoTokenizer, 
    AutoModelForCausalLM,
    BitsAndBytesConfig
)
from peft import PeftModel, PeftConfig
from sentence_transformers import SentenceTransformer

# LangChain imports
from langchain.llms import HuggingFacePipeline
from langchain.chains import LLMChain
from langchain.memory import ConversationBufferMemory
from langchain.prompts import PromptTemplate
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.document_loaders import TextLoader
from langchain.vectorstores import FAISS

# Import FlowManager
from conversation_flow import FlowManager

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler()]
)
logger = logging.getLogger(__name__)

# Suppress warnings
import warnings
warnings.filterwarnings('ignore', category=UserWarning)

# Set up cache directories
def setup_cache_dirs():
    # Check if running in Hugging Face Spaces
    is_spaces = os.environ.get('SPACE_ID') is not None
    
    if is_spaces:
        # Use /tmp for Hugging Face Spaces with proper permissions
        cache_dir = '/tmp/huggingface'
        os.environ.update({
            'TRANSFORMERS_CACHE': cache_dir,
            'HF_HOME': cache_dir,
            'TOKENIZERS_PARALLELISM': 'false',
            'TRANSFORMERS_VERBOSITY': 'error',
            'BITSANDBYTES_NOWELCOME': '1',
            'HF_DATASETS_CACHE': cache_dir,
            'HF_METRICS_CACHE': cache_dir,
            'HF_MODULES_CACHE': cache_dir,
            'HUGGING_FACE_HUB_TOKEN': os.environ.get('HF_TOKEN', ''),
            'HF_TOKEN': os.environ.get('HF_TOKEN', '')
        })
    else:
        # Use default cache for local development
        cache_dir = os.path.expanduser('~/.cache/huggingface')
        os.environ.update({
            'TOKENIZERS_PARALLELISM': 'false',
            'TRANSFORMERS_VERBOSITY': 'error',
            'BITSANDBYTES_NOWELCOME': '1'
        })
    
    # Create cache directory if it doesn't exist
    os.makedirs(cache_dir, exist_ok=True)
    
    return cache_dir

# Set up cache directories
CACHE_DIR = setup_cache_dirs()

# Define base directory and paths
BASE_DIR = os.path.abspath(os.path.dirname(__file__))
MODELS_DIR = os.path.join(BASE_DIR, "models")
VECTOR_DB_PATH = os.path.join(BASE_DIR, "vector_db")
SESSION_DATA_PATH = os.path.join(BASE_DIR, "session_data")
SUMMARIES_DIR = os.path.join(BASE_DIR, "session_summaries")

# Create necessary directories
for directory in [MODELS_DIR, VECTOR_DB_PATH, SESSION_DATA_PATH, SUMMARIES_DIR]:
    os.makedirs(directory, exist_ok=True)

# Pydantic models
class Message(BaseModel):
    text: str = Field(..., description="The content of the message")
    timestamp: str = Field(None, description="ISO format timestamp of the message")
    role: str = Field("user", description="The role of the message sender (user or assistant)")

class SessionSummary(BaseModel):
    session_id: str = Field(
        ...,
        description="Unique identifier for the session",
        examples=["user_789_session_20240314"]
    )
    user_id: str = Field(
        ...,
        description="Identifier of the user",
        examples=["user_123"])
    start_time: str = Field(..., description="ISO format start time of the session"
    )
    end_time: str = Field(
        ...,
        description="ISO format end time of the session"
    )
    message_count: int = Field(
        ...,
        description="Total number of messages in the session"
    )
    duration_minutes: float = Field(
        ...,
        description="Duration of the session in minutes"
    )
    primary_emotions: List[str] = Field(
        ...,
        min_items=1,
        description="List of primary emotions detected",
        examples=[
            ["anxiety", "stress"],
            ["joy", "excitement"],
            ["sadness", "loneliness"]
        ]
    )
    emotion_progression: List[Dict[str, float]] = Field(
        ...,
        description="Progression of emotions throughout the session",
        examples=[
            [
                {"anxiety": 0.8, "stress": 0.6},
                {"calm": 0.7, "anxiety": 0.3},
                {"joy": 0.9, "calm": 0.8}
            ]
        ]
    )
    summary_text: str = Field(
        ...,
        description="Text summary of the session",
        examples=[
            "The session focused on managing work-related stress and developing coping strategies. The client showed improvement in recognizing stress triggers and implementing relaxation techniques.",
            "Discussion centered around relationship challenges and self-esteem issues. The client expressed willingness to try new communication strategies."
        ]
    )
    recommendations: Optional[List[str]] = Field(
        None,
        description="Optional recommendations based on the session"
    )

class Conversation(BaseModel):
    user_id: str = Field(
        ...,
        description="Identifier of the user",
        examples=["user_123"]
    )
    session_id: str = Field(
        "",
        description="Identifier of the current session"
    )
    start_time: str = Field(
        "",
        description="ISO format start time of the conversation"
    )
    messages: List[Message] = Field(
        [],
        description="List of messages in the conversation",
        examples=[
            [
                Message(text="I'm feeling anxious", role="user"),
                Message(text="I understand you're feeling anxious. Can you tell me more about what's causing this?", role="assistant")
            ]
        ]
    )
    emotion_history: List[Dict[str, float]] = Field(
        [],
        description="History of emotions detected",
        examples=[
            [
                {"anxiety": 0.8, "stress": 0.6},
                {"calm": 0.7, "anxiety": 0.3}
            ]
        ]
    )
    context: Dict[str, Any] = Field(
        {},
        description="Additional context for the conversation",
        examples=[
            {
                "last_emotion": "anxiety",
                "conversation_topic": "work stress",
                "previous_sessions": 3
            }
        ]
    )
    is_active: bool = Field(
        True,
        description="Whether the conversation is currently active",
        examples=[True, False]
    )

class MentalHealthChatbot:
    def __init__(
        self,
        model_name: str = "meta-llama/Llama-3.2-3B-Instruct",
        peft_model_path: str = "nada013/mental-health-chatbot", 
        therapy_guidelines_path: str = None,
        use_4bit: bool = True,
        device: str = None,
        max_response_length: int = 500,  # Maximum characters in response
        max_response_words: int = 100,   # Maximum words in response
        min_response_words: int = 10,    # Minimum words in response
        max_consecutive_responses: int = 3  # Max consecutive responses without user input
    ):
        # Set device (cuda if available, otherwise cpu)
        if device is None:
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            self.device = device

        # Set memory optimization for T4
        if self.device == "cuda":
            torch.cuda.empty_cache()  # Clear GPU cache
            # Set smaller batch size for T4
            self.batch_size = 4
            # Enable memory efficient attention
            os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
        else:
            self.batch_size = 8

        logger.info(f"Using device: {self.device}")
        
        # Set response limits
        self.max_response_length = max_response_length
        self.max_response_words = max_response_words
        self.min_response_words = min_response_words
        self.max_consecutive_responses = max_consecutive_responses
        self.consecutive_response_count = 0  # Track consecutive responses
        
        # Initialize models
        self.peft_model_path = peft_model_path
        
        # Initialize emotion detection model
        logger.info("Loading emotion detection model")
        self.emotion_classifier = self._load_emotion_model()
        
        # Initialize LLAMA model
        logger.info(f"Loading LLAMA model: {model_name}")
        self.llama_model, self.llama_tokenizer, self.llm = self._initialize_llm(model_name, use_4bit)
        
        # Initialize summary model
        logger.info("Loading summary model")
        self.summary_model = pipeline(
            "summarization",
            model="philschmid/bart-large-cnn-samsum",
            device=0 if self.device == "cuda" else -1,
            model_kwargs={
                "cache_dir": CACHE_DIR,
                "torch_dtype": torch.float16,
                "max_memory": {0: "2GB"} if self.device == "cuda" else None
            }
        )
        logger.info("Summary model loaded successfully")
        
        # Initialize FlowManager
        logger.info("Initializing FlowManager")
        self.flow_manager = FlowManager(self.llm)
        
        # Setup conversation memory with LangChain
        self.memory = ConversationBufferMemory(
            return_messages=True,
            input_key="input"
        )
        
        # Create conversation prompt template
        self.prompt_template = PromptTemplate(
            input_variables=["history", "input", "past_context", "emotion_context", "guidelines"],
            template="""You are a supportive and empathetic mental health conversational AI. Your role is to provide therapeutic support while maintaining professional boundaries.

Previous conversation:
{history}

EMOTIONAL CONTEXT:
{emotion_context}

Past context: {past_context}

Relevant therapeutic guidelines:
{guidelines}

Current message: {input}

Provide a supportive response that:
1. Validates the user's feelings without using casual greetings
2. Asks relevant follow-up questions
3. Maintains a conversational tone , professional and empathetic tone
4. Focuses on understanding and support
5. Avoids repeating previous responses

Response:"""
        )
        
        # Create the conversation chain
        self.conversation = LLMChain(
            llm=self.llm,
            prompt=self.prompt_template,
            memory=self.memory,
            verbose=False
        )
        
        # Setup embeddings for vector search
        self.embeddings = HuggingFaceEmbeddings(
            model_name="sentence-transformers/all-MiniLM-L6-v2"
        )
        
        # Setup vector database for retrieving relevant past conversations
        if therapy_guidelines_path and os.path.exists(therapy_guidelines_path):
            self.setup_vector_db(therapy_guidelines_path)
        else:
            self.setup_vector_db(None)
        
        # Initialize conversation storage
        self.conversations = {}
        
        # Load existing session summaries
        self.session_summaries = {}
        self._load_existing_summaries()
        
        logger.info("All models and components initialized successfully")
        
    def _load_emotion_model(self):
        try:
            # Load emotion model directly from Hugging Face
            return pipeline(
                "text-classification",
                model="SamLowe/roberta-base-go_emotions",
                top_k=None,
                device_map="auto" if torch.cuda.is_available() else None,
                model_kwargs={
                    "cache_dir": CACHE_DIR,
                    "torch_dtype": torch.float16,  # Use float16
                    "max_memory": {0: "2GB"} if torch.cuda.is_available() else None  # Limit memory usage
                },
            )
        except Exception as e:
            logger.error(f"Error loading emotion model: {e}")
            # Fallback to a simpler model
            try:
                return pipeline(
                    "text-classification", 
                    model="j-hartmann/emotion-english-distilroberta-base", 
                    return_all_scores=True,
                    device_map="auto" if torch.cuda.is_available() else None,
                    model_kwargs={
                        "cache_dir": CACHE_DIR,
                        "torch_dtype": torch.float16,
                        "max_memory": {0: "2GB"} if torch.cuda.is_available() else None
                    },
                )
            except Exception as e:
                logger.error(f"Error loading fallback emotion model: {e}")
                # Return a simple pipeline that always returns neutral
                return lambda text: [{"label": "neutral", "score": 1.0}]
    
    def _initialize_llm(self, model_name: str, use_4bit: bool):
        try:
            # Configure quantization only if CUDA is available
            if use_4bit and torch.cuda.is_available():
                quantization_config = BitsAndBytesConfig(
                    load_in_4bit=True,
                    bnb_4bit_compute_dtype=torch.float16,
                    bnb_4bit_quant_type="nf4",
                    bnb_4bit_use_double_quant=True,
                )
                
                # Set max memory for T4 GPU
                max_memory = {0: "14GB"}  # Leave 2GB buffer for other operations
            else:
                quantization_config = None
                max_memory = None
                logger.info("CUDA not available, running without quantization")

            # Load base model
            logger.info(f"Loading base model: {model_name}")
            base_model = AutoModelForCausalLM.from_pretrained(
                model_name,
                quantization_config=quantization_config,
                device_map="auto",
                max_memory=max_memory,
                trust_remote_code=True,
                cache_dir=CACHE_DIR,
                use_auth_token=os.environ.get('HF_TOKEN'),
                torch_dtype=torch.float16  # Use float16 for better memory efficiency
            )

            # Load tokenizer
            logger.info("Loading tokenizer")
            tokenizer = AutoTokenizer.from_pretrained(
                model_name,
                cache_dir=CACHE_DIR,
                use_auth_token=os.environ.get('HF_TOKEN')  # Add auth token for gated models
            )
            tokenizer.pad_token = tokenizer.eos_token

            # Load PEFT model
            logger.info(f"Loading PEFT model from {self.peft_model_path}")
            model = PeftModel.from_pretrained(
                base_model, 
                self.peft_model_path,
                cache_dir=CACHE_DIR,
                use_auth_token=os.environ.get('HF_TOKEN')  # Add auth token for gated models
            )
            logger.info("Successfully loaded PEFT model")

            # Create text generation pipeline
            text_generator = pipeline(
                "text-generation",
                model=model,
                tokenizer=tokenizer,
                max_new_tokens=512,
                temperature=0.7,
                top_p=0.95,
                repetition_penalty=1.1,
                do_sample=True,
                device_map="auto" if torch.cuda.is_available() else None
            )
            
            # Create LangChain wrapper
            llm = HuggingFacePipeline(pipeline=text_generator)
            
            return model, tokenizer, llm
            
        except Exception as e:
            logger.error(f"Error initializing LLM: {str(e)}")
            raise
    
    def setup_vector_db(self, guidelines_path: str = None):
        
        logger.info("Setting up FAISS vector database")
        
        # Check if vector DB exists
        vector_db_exists = os.path.exists(os.path.join(VECTOR_DB_PATH, "index.faiss"))
        
        if not vector_db_exists:
            # Load therapy guidelines 
            if guidelines_path and os.path.exists(guidelines_path):
                loader = TextLoader(guidelines_path)
                documents = loader.load()
                
                # Split documents into chunks with better overlap for context
                text_splitter = RecursiveCharacterTextSplitter(
                    chunk_size=500,  # Smaller chunks for more precise retrieval
                    chunk_overlap=100,
                    separators=["\n\n", "\n", " ", ""]
                )
                chunks = text_splitter.split_documents(documents)
                
                # Create and save the vector store
                self.vector_db = FAISS.from_documents(chunks, self.embeddings)
                self.vector_db.save_local(VECTOR_DB_PATH)
                logger.info("Successfully loaded and indexed therapy guidelines")
            else:
                # Initialize with empty vector DB
                self.vector_db = FAISS.from_texts(["Initial empty vector store"], self.embeddings)
                self.vector_db.save_local(VECTOR_DB_PATH)
                logger.warning("No guidelines file provided, using empty vector store")
        else:
            # Load existing vector DB
            self.vector_db = FAISS.load_local(VECTOR_DB_PATH, self.embeddings, allow_dangerous_deserialization=True)
            logger.info("Loaded existing vector database")
    
    def _load_existing_summaries(self):
        if not os.path.exists(SUMMARIES_DIR):
            return
            
        for filename in os.listdir(SUMMARIES_DIR):
            if filename.endswith('.json'):
                try:
                    with open(os.path.join(SUMMARIES_DIR, filename), 'r') as f:
                        summary_data = json.load(f)
                        session_id = summary_data.get('session_id')
                        if session_id:
                            self.session_summaries[session_id] = summary_data
                except Exception as e:
                    logger.warning(f"Failed to load summary from {filename}: {e}")

    def detect_emotion(self, text: str) -> Dict[str, float]:
        try:
            results = self.emotion_classifier(text)[0]
            return {result['label']: result['score'] for result in results}
        except Exception as e:
            logger.error(f"Error detecting emotions: {e}")
            return {"neutral": 1.0}

    def _validate_and_limit_response(self, response: str, user_message: str) -> str:
        """
        Validate and limit response length and content.
        Returns a properly limited response.
        """
        if not response or not response.strip():
            return "I understand. Could you tell me more about that?"
        
        # Clean the response
        response = response.strip()
        
        # Remove any LLM commentary or instructions
        response = re.sub(r"(Your response|This response|Response:|Note:).*", "", response, flags=re.IGNORECASE).strip()
        response = re.sub(r"---.*", "", response).strip()
        
        # Remove casual greetings
        response = re.sub(r'^(Hey|Hi|Hello|Hi there|Hey there),\s*', '', response)
        
        # Count words and characters
        words = response.split()
        word_count = len(words)
        char_count = len(response)
        
        # Check if response is too short
        if word_count < self.min_response_words:
            logger.info(f"Response too short ({word_count} words), adding follow-up question")
            if not response.endswith('?'):
                response += " Could you tell me more about that?"
        
        # Check if response is too long
        if char_count > self.max_response_length or word_count > self.max_response_words:
            logger.info(f"Response too long ({char_count} chars, {word_count} words), truncating")
            
            # Try to find a good breaking point
            if word_count > self.max_response_words:
                # Truncate to max words
                truncated_words = words[:self.max_response_words]
                response = ' '.join(truncated_words)
                
                # Try to end at a sentence
                last_period = response.rfind('.')
                last_question = response.rfind('?')
                last_exclamation = response.rfind('!')
                
                end_point = max(last_period, last_question, last_exclamation)
                if end_point > len(response) * 0.7:  # If we can end at a sentence within 70% of the limit
                    response = response[:end_point + 1]
                else:
                    # Add ellipsis if we can't end naturally
                    response = response.rstrip() + "..."
            
            elif char_count > self.max_response_length:
                # Truncate to max characters
                response = response[:self.max_response_length]
                
                # Try to end at a word boundary
                last_space = response.rfind(' ')
                if last_space > len(response) * 0.8:  # If we can end at a word within 80% of the limit
                    response = response[:last_space]
                else:
                    # Add ellipsis
                    response = response.rstrip() + "..."
        
        # Check for repetitive content
        if self._is_repetitive(response, user_message):
            logger.info("Response detected as repetitive, generating alternative")
            return "I hear what you're saying. Could you help me understand this better?"
        
        # Ensure response ends properly
        if not response.endswith(('.', '!', '?')):
            response = response.rstrip() + '.'
        
        return response.strip()

    def _is_repetitive(self, response: str, user_message: str) -> bool:
        """
        Check if response is repetitive or too similar to user message.
        """
        # Convert to lowercase for comparison
        response_lower = response.lower()
        user_lower = user_message.lower()
        
        # Check if response contains too much of the user's message
        user_words = set(user_lower.split())
        response_words = set(response_lower.split())
        
        if len(user_words) > 3:  # Only check if user message has enough words
            common_words = user_words.intersection(response_words)
            if len(common_words) / len(user_words) > 0.6:  # If more than 60% of user words are in response
                return True
        
        # Check for repetitive phrases
        repetitive_phrases = [
            "i understand", "i hear you", "that sounds", "i can see",
            "thank you for sharing", "i appreciate", "that must be"
        ]
        
        phrase_count = sum(1 for phrase in repetitive_phrases if phrase in response_lower)
        if phrase_count > 2:  # If more than 2 repetitive phrases
            return True
        
        return False

    def retrieve_relevant_context(self, query: str, k: int = 3) -> str:
        # Retrieve relevant past conversations using vector similarity
        if not hasattr(self, 'vector_db'):
            return ""
        
        try:
            # Retrieve similar documents from vector DB
            docs = self.vector_db.similarity_search(query, k=k)
            
            # Combine the content of retrieved documents
            relevant_context = "\n".join([doc.page_content for doc in docs])
            return relevant_context
        except Exception as e:
            logger.error(f"Error retrieving context: {e}")
            return ""

    def retrieve_relevant_guidelines(self, query: str, emotion_context: str) -> str:
        if not hasattr(self, 'vector_db'):
            return ""
        
        try:
            # Combine query and emotion context for better relevance
            search_query = f"{query} {emotion_context}"
            
            # Retrieve similar documents from vector DB
            docs = self.vector_db.similarity_search(search_query, k=2)
            
            # Combine the content of retrieved documents
            relevant_guidelines = "\n".join([doc.page_content for doc in docs])
            return relevant_guidelines
        except Exception as e:
            logger.error(f"Error retrieving guidelines: {e}")
            return ""

    def generate_response(self, prompt: str, emotion_data: Dict[str, float], conversation_history: List[Dict]) -> str:
        
        # Get primary and secondary emotions
        sorted_emotions = sorted(emotion_data.items(), key=lambda x: x[1], reverse=True)
        primary_emotion = sorted_emotions[0][0] if sorted_emotions else "neutral"
        
        # Get secondary emotions (if any)
        secondary_emotions = []
        for emotion, score in sorted_emotions[1:3]:  # Get 2nd and 3rd strongest emotions
            if score > 0.2:  # Only include if reasonably strong
                secondary_emotions.append(emotion)
        
        # Create emotion context string
        emotion_context = f"User is primarily feeling {primary_emotion}"
        if secondary_emotions:
            emotion_context += f" with elements of {' and '.join(secondary_emotions)}"
        emotion_context += "."
        
        # Retrieve relevant guidelines
        guidelines = self.retrieve_relevant_guidelines(prompt, emotion_context)
        
        # Retrieve past context
        past_context = self.retrieve_relevant_context(prompt)
        
        # Generate response using the conversation chain
        response = self.conversation.predict(
            input=prompt,
            past_context=past_context,
            emotion_context=emotion_context,
            guidelines=guidelines
        )
        
        # Validate and limit the response
        response = self._validate_and_limit_response(response, prompt)
        
        # Ensure the response is unique and not repeating previous messages
        if len(conversation_history) > 0:
            last_responses = [msg["text"] for msg in conversation_history[-4:] if msg["role"] == "assistant"]
            if response in last_responses:
                logger.info("Response detected as duplicate, generating alternative")
                # Generate a new response with a different angle
                alternative_response = self.conversation.predict(
                    input=f"{prompt} (Please provide a different perspective)",
                    past_context=past_context,
                    emotion_context=emotion_context,
                    guidelines=guidelines
                )
                alternative_response = self._validate_and_limit_response(alternative_response, prompt)
                response = alternative_response
        
        return response

    def generate_session_summary(
        self,
        flow_manager_session: Dict = None
    ) -> Dict:
        
        if not flow_manager_session:
            return {
                "session_id": "",
                "user_id": "",
                "start_time": "",
                "end_time": datetime.now().isoformat(),
                "duration_minutes": 0,
                "current_phase": "unknown",
                "primary_emotions": [],
                "emotion_progression": [],
                "summary": "Error: No session data provided",
                "recommendations": ["Unable to generate recommendations"],
                "session_characteristics": {}
            }
            
        # Get session data from FlowManager
        session_id = flow_manager_session.get('session_id', '')
        user_id = flow_manager_session.get('user_id', '')
        current_phase = flow_manager_session.get('current_phase')
        
        if current_phase:
            # Convert ConversationPhase to dict
            current_phase = {
                'name': current_phase.name,
                'description': current_phase.description,
                'goals': current_phase.goals,
                'started_at': current_phase.started_at,
                'ended_at': current_phase.ended_at,
                'completion_metrics': current_phase.completion_metrics
            }
        
        session_start = flow_manager_session.get('started_at')
        if isinstance(session_start, str):
            session_start = datetime.fromisoformat(session_start)
        session_duration = (datetime.now() - session_start).total_seconds() / 60 if session_start else 0
        
        # Get emotion progression and primary emotions
        emotion_progression = flow_manager_session.get('emotion_progression', [])
        emotion_history = flow_manager_session.get('emotion_history', [])
        
        # Extract primary emotions from emotion history
        primary_emotions = []
        if emotion_history:
            # Get the most frequent emotions
            emotion_counts = {}
            for entry in emotion_history:
                emotions = entry.get('emotions', {})
                if isinstance(emotions, dict):
                    primary = max(emotions.items(), key=lambda x: x[1])[0]
                    emotion_counts[primary] = emotion_counts.get(primary, 0) + 1
            
            # sort by frequency and get top 3
            primary_emotions = sorted(emotion_counts.items(), key=lambda x: x[1], reverse=True)[:3]
            primary_emotions = [emotion for emotion, _ in primary_emotions]
        
        # get session 
        session_characteristics = flow_manager_session.get('llm_context', {}).get('session_characteristics', {})
        
        # prepare the text for summarization
        summary_text = f"""
        Session Overview:
        - Session ID: {session_id}
        - User ID: {user_id}
        - Phase: {current_phase.get('name', 'unknown') if current_phase else 'unknown'}
        - Duration: {session_duration:.1f} minutes
        
        Emotional Analysis:
        - Primary Emotions: {', '.join(primary_emotions) if primary_emotions else 'No primary emotions detected'}
        - Emotion Progression: {', '.join(emotion_progression) if emotion_progression else 'No significant emotion changes noted'}
        
        Session Characteristics:
        - Therapeutic Alliance: {session_characteristics.get('alliance_strength', 'N/A')}
        - Engagement Level: {session_characteristics.get('engagement_level', 'N/A')}
        - Emotional Pattern: {session_characteristics.get('emotional_pattern', 'N/A')}
        - Cognitive Pattern: {session_characteristics.get('cognitive_pattern', 'N/A')}
        
        Key Observations:
        - The session focused on {current_phase.get('description', 'general discussion') if current_phase else 'general discussion'}
        - Main emotional themes: {', '.join(primary_emotions) if primary_emotions else 'not identified'}
        - Session progress: {session_characteristics.get('progress_quality', 'N/A')}
        """
        
        # Generate summary using BART
        summary = self.summary_model(
            summary_text,
            max_length=150,
            min_length=50,
            do_sample=False
        )[0]['summary_text']
        
        # Generate recommendations using Llama
        recommendations_prompt = f"""
        Based on the following session summary, provide 2-3 specific recommendations for follow-up:
        
        {summary}
        
        Session Characteristics:
        - Therapeutic Alliance: {session_characteristics.get('alliance_strength', 'N/A')}
        - Engagement Level: {session_characteristics.get('engagement_level', 'N/A')}
        - Emotional Pattern: {session_characteristics.get('emotional_pattern', 'N/A')}
        - Cognitive Pattern: {session_characteristics.get('cognitive_pattern', 'N/A')}
        
        Recommendations should be:
        1. Actionable and specific
        2. Based on the session content
        3. Focused on next steps
        """
        
        recommendations = self.llm.invoke(recommendations_prompt)
        
       
        recommendations = recommendations.split('\n')
        recommendations = [r.strip() for r in recommendations if r.strip()]
        recommendations = [r for r in recommendations if not r.startswith(('Based on', 'Session', 'Recommendations'))]
        
        
        return {
            "session_id": session_id,
            "user_id": user_id,
            "start_time": session_start.isoformat() if isinstance(session_start, datetime) else str(session_start),
            "end_time": datetime.now().isoformat(),
            "duration_minutes": session_duration,
            "current_phase": current_phase.get('name', 'unknown') if current_phase else 'unknown',
            "primary_emotions": primary_emotions,
            "emotion_progression": emotion_progression,
            "summary": summary,
            "recommendations": recommendations,
            "session_characteristics": session_characteristics
        }

    def start_session(self, user_id: str) -> tuple[str, str]:
        # Generate session id
        session_id = f"{user_id}_{datetime.now().strftime('%Y%m%d%H%M%S')}"
        
        # Initialize FlowManager session
        self.flow_manager.initialize_session(user_id)
        
        # Create a new conversation
        self.conversations[user_id] = Conversation(
            user_id=user_id,
            session_id=session_id,
            start_time=datetime.now().isoformat(),
            is_active=True
        )
        
        # Clear conversation memory
        self.memory.clear()
        
        # Generate initial greeting and question
        initial_message = """Hello! I'm here to support you today. How have you been feeling lately?"""
        
        # Add the initial message to conversation history
        assistant_message = Message(
            text=initial_message,
            timestamp=datetime.now().isoformat(),
            role="assistant"
        )
        self.conversations[user_id].messages.append(assistant_message)
        
        logger.info(f"Session started for user {user_id}")
        return session_id, initial_message

    def end_session(
        self, 
        user_id: str, 
        flow_manager: Optional[Any] = None
    ) -> Optional[Dict]:

        if user_id not in self.conversations or not self.conversations[user_id].is_active:
            return None
        
        conversation = self.conversations[user_id]
        conversation.is_active = False
        
        # Get FlowManager session data
        flow_manager_session = self.flow_manager.user_sessions.get(user_id)
        
        # Generate session summary
        try:
            session_summary = self.generate_session_summary(flow_manager_session)
            
            # Save summary to disk
            summary_path = os.path.join(SUMMARIES_DIR, f"{session_summary['session_id']}.json")
            with open(summary_path, 'w') as f:
                json.dump(session_summary, f, indent=2)
            
            # Store in memory
            self.session_summaries[session_summary['session_id']] = session_summary
            
            # Clear conversation memory
            self.memory.clear()
            
            return session_summary
        except Exception as e:
            logger.error(f"Failed to generate session summary: {e}")
            return None

    def process_message(self, user_id: str, message: str) -> str:
        
        # Check for risk flags 
        risk_keywords = ["suicide", "kill myself", "end my life", "self-harm", "hurt myself"]
        risk_detected = any(keyword in message.lower() for keyword in risk_keywords)
        
        # Create or get conversation
        if user_id not in self.conversations or not self.conversations[user_id].is_active:
            self.start_session(user_id)
        
        conversation = self.conversations[user_id]
        
        #  user message -> conversation history
        new_message = Message(
            text=message,
            timestamp=datetime.now().isoformat(),
            role="user"
        )
        conversation.messages.append(new_message)
        
        # For crisis 
        if risk_detected:
            logger.warning(f"Risk flag detected in session {user_id}")
            
            crisis_response = """ I'm really sorry you're feeling this way — it sounds incredibly heavy, and I want you to know that you're not alone.

You don't have to face this by yourself. Our app has licensed mental health professionals who are ready to support you. I can connect you right now if you'd like.

In the meantime, I'm here to listen and talk with you. You can also do grounding exercises or calming techniques with me if you prefer. Just say "help me calm down" or "I need a break."

Would you like to connect with a professional now, or would you prefer to keep talking with me for a bit? Either way, I'm here for you."""
            
            #  assistant response -> conversation history
            assistant_message = Message(
                text=crisis_response,
                timestamp=datetime.now().isoformat(),
                role="assistant"
            )
            conversation.messages.append(assistant_message)
            
            return crisis_response
        
        # Reset consecutive response counter when user sends a message
        self.consecutive_response_count = 0
        
        # Detect emotions
        emotions = self.detect_emotion(message)
        conversation.emotion_history.append(emotions)
        
        # Process message with FlowManager
        flow_context = self.flow_manager.process_message(user_id, message, emotions)
        
        # Format conversation history 
        conversation_history = []
        for msg in conversation.messages:
            conversation_history.append({
                "text": msg.text,
                "timestamp": msg.timestamp,
                "role": msg.role
            })
        
        # Check rate limiting for consecutive responses
        if self.consecutive_response_count >= self.max_consecutive_responses:
            logger.warning(f"Rate limit reached for user {user_id}, sending brief response")
            response_text = "I'm here to listen. Take your time to share what's on your mind."
            self.consecutive_response_count = 0  # Reset counter
        else:
            # Generate response
            response_text = self.generate_response(message, emotions, conversation_history)
            
            # Increment consecutive response counter
            self.consecutive_response_count += 1
            
            # Generate a follow-up question if the response is too short and we haven't hit limits
            if (len(response_text.split()) < self.min_response_words and 
                not response_text.endswith('?') and 
                self.consecutive_response_count < self.max_consecutive_responses):
                
                follow_up_prompt = f"""
Recent conversation:
{chr(10).join([f"{msg['role']}: {msg['text']}" for msg in conversation_history[-3:]])}

Now, write a single empathetic and open-ended question to encourage the user to share more. 
Respond with just the question, no explanation.
"""
                follow_up = self.llm.invoke(follow_up_prompt).strip()
                # Clean and extract only the actual question (first sentence ending with '?')
                matches = re.findall(r'([^\n.?!]*\?)', follow_up)
                if matches:
                    question = matches[0].strip()
                else:
                    question = follow_up.strip().split('\n')[0]
                
                # Validate the follow-up question
                question = self._validate_and_limit_response(question, message)
                
                # If the main response is very short, return just the question
                if len(response_text.split()) < 5:
                    response_text = question
                else:
                    response_text = f"{response_text}\n\n{question}"
        
        #  assistant response -> conversation history
        assistant_message = Message(
            text=response_text,
            timestamp=datetime.now().isoformat(),
            role="assistant"
        )
        conversation.messages.append(assistant_message)
        
        # Update context
        conversation.context.update({
            "last_emotion": emotions,
            "last_interaction": datetime.now().isoformat(),
            "flow_context": flow_context
        })
        
        # Store this interaction in vector database
        current_interaction = f"User: {message}\nChatbot: {response_text}"
        self.vector_db.add_texts([current_interaction])
        self.vector_db.save_local(VECTOR_DB_PATH)
        
        return response_text

    def get_session_summary(self, session_id: str) -> Optional[Dict[str, Any]]:
        
        return self.session_summaries.get(session_id)

    def get_user_replies(self, user_id: str) -> List[Dict[str, Any]]:
        if user_id not in self.conversations:
            return []
            
        conversation = self.conversations[user_id]
        user_replies = []
        
        for message in conversation.messages:
            if message.role == "user":
                user_replies.append({
                    "text": message.text,
                    "timestamp": message.timestamp,
                    "session_id": conversation.session_id
                })
                
        return user_replies

if __name__ == "__main__":
    pass