ML/TextGenerator/content_generation.py

import os
import json
import requests
import logging
from typing import List, Dict, Optional, Union, Any
from datetime import datetime
import random
from dotenv import load_dotenv

# Import Google's Generative AI library
try:
    import google.generativeai as genai
except ImportError:
    genai = None

# Load environment variables
load_dotenv()

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Get API keys from environment
ANTHROPIC_API_KEY = os.environ.get("ANTHROPIC_API_KEY")
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
DEEPSEEK_API_KEY = os.environ.get("DEEPSEEK_API_KEY")
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")

# Initialize Gemini if available
if genai and GEMINI_API_KEY:
    genai.configure(api_key=GEMINI_API_KEY)


class ConversationContext:
    """Maintains conversation history and context for more natural interactions"""
    
    def __init__(self, role: str = "speech_therapist", topic: str = "general"):
        self.role = role
        self.topic = topic
        self.history = []
        self.turn_count = 0
        self.user_interests = []
        self.conversation_style = "supportive"
        
    def add_turn(self, speaker: str, message: str):
        self.history.append({
            "speaker": speaker,
            "message": message,
            "timestamp": datetime.now().isoformat()
        })
        if speaker == "user":
            self.turn_count += 1
            
    def get_conversation_summary(self, last_n_turns: int = 5) -> str:
        """Get a summary of recent conversation for context"""
        recent = self.history[-last_n_turns:] if len(self.history) > last_n_turns else self.history
        summary = []
        for turn in recent:
            summary.append(f"{turn['speaker'].capitalize()}: {turn['message']}")
        return "\n".join(summary)


class AIContentGenerator:
    """
    Class to handle AI-based content generation for speech therapy practice
    Uses Gemini, DeepSeek, Anthropic Claude, or OpenAI GPT depending on available API keys
    """

    def __init__(self):
        # Determine which AI provider to use
        self.has_gemini = bool(genai and GEMINI_API_KEY)
        self.has_deepseek = bool(DEEPSEEK_API_KEY)
        self.has_anthropic = bool(ANTHROPIC_API_KEY)
        self.has_openai = bool(OPENAI_API_KEY)
        
        # Initialize conversation contexts
        self.conversation_contexts = {}

        print(f"Has Gemini API: {self.has_gemini}")
        print(f"Has DeepSeek API: {self.has_deepseek}")
        print(f"Has Anthropic API: {self.has_anthropic}")
        print(f"Has OpenAI API: {self.has_openai}")

        if not self.has_gemini and not self.has_deepseek and not self.has_anthropic and not self.has_openai:
            logger.warning("No API keys found. Using mock content generation instead.")

    def generate_content(self, prompt: str) -> str:
        """Generate content using available AI provider"""

        # Prioritize Gemini if available
        if self.has_gemini:
            return self._generate_with_gemini(prompt)
        elif self.has_deepseek:
            return self._generate_with_deepseek(prompt)
        elif self.has_anthropic:
            return self._generate_with_anthropic(prompt)
        elif self.has_openai:
            return self._generate_with_openai(prompt)
        else:
            return self._generate_mock_content(prompt)

    def _generate_with_gemini(self, prompt: str) -> str:
        """Generate content using Google's Gemini API"""
        try:
            # Configure the API
            genai.configure(api_key=GEMINI_API_KEY)

            # Use component-specific model for reading content generation
            model_name = os.environ.get('READING_CONTENT_MODEL', 'gemini-2.5-flash')
            model = genai.GenerativeModel(model_name)
            
            # Generate content
            response = model.generate_content(
                prompt,
                generation_config={
                    "temperature": 0.7,
                    "top_p": 0.95,
                    "top_k": 40,
                    "max_output_tokens": 1000,
                }
            )
            
            if response and hasattr(response, 'text'):
                # Log cost
                try:
                    from utils.cost_tracker import log_gemini_call
                    log_gemini_call(prompt, response.text, model_name, practice_type='content_generation')
                except Exception as e:
                    logger.warning(f"Failed to log Gemini cost: {e}")
                return response.text
            elif response and hasattr(response, 'parts'):
                # Log cost
                try:
                    from utils.cost_tracker import log_gemini_call
                    log_gemini_call(prompt, response.parts[0].text, model_name, practice_type='content_generation')
                except Exception as e:
                    logger.warning(f"Failed to log Gemini cost: {e}")
                return response.parts[0].text
            else:
                raise ValueError(f"Unexpected response format from Gemini API")
                
        except Exception as e:
            logger.error(f"Error generating with Gemini: {str(e)}")
            
            # Fall back to other providers
            if self.has_deepseek:
                return self._generate_with_deepseek(prompt)
            elif self.has_anthropic:
                return self._generate_with_anthropic(prompt)
            elif self.has_openai:
                return self._generate_with_openai(prompt)
            else:
                return self._generate_mock_content(prompt)

    def _generate_with_deepseek(self, prompt: str) -> str:
        """Generate content using DeepSeek API"""
        try:
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {DEEPSEEK_API_KEY}"
            }

            payload = {
                "model": "deepseek-chat",
                "messages": [
                    {"role": "system", "content": "You are a speech therapy assistant that creates reading practice materials."},
                    {"role": "user", "content": prompt}
                ],
                "temperature": 0.7,
                "max_tokens": 1000
            }

            response = requests.post(
                "https://api.deepseek.com/v1/chat/completions",
                json=payload,
                headers=headers
            )

            response.raise_for_status()
            result = response.json()
            return result["choices"][0]["message"]["content"].strip()

        except Exception as e:
            logger.error(f"Error generating content with DeepSeek: {str(e)}")
            if self.has_anthropic:
                return self._generate_with_anthropic(prompt)
            elif self.has_openai:
                return self._generate_with_openai(prompt)
            else:
                return self._generate_mock_content(prompt)

    def _generate_with_anthropic(self, prompt: str) -> str:
        """Generate content using Anthropic's Claude API"""
        try:
            headers = {
                "Content-Type": "application/json",
                "X-Api-Key": ANTHROPIC_API_KEY,
                "anthropic-version": "2023-06-01"
            }

            data = {
                "model": "claude-3-haiku-20240307",
                "messages": [
                    {"role": "user", "content": prompt}
                ],
                "max_tokens": 1000,
                "temperature": 0.7
            }

            response = requests.post(
                "https://api.anthropic.com/v1/messages",
                headers=headers,
                json=data
            )

            response.raise_for_status()
            result = response.json()
            return result["content"][0]["text"]

        except Exception as e:
            logger.error(f"Error generating content with Anthropic: {str(e)}")
            if self.has_openai:
                return self._generate_with_openai(prompt)
            else:
                return self._generate_mock_content(prompt)

    def _generate_with_openai(self, prompt: str) -> str:
        """Generate content using OpenAI's API"""
        try:
            headers = {
                "Authorization": f"Bearer {OPENAI_API_KEY}",
                "Content-Type": "application/json"
            }

            data = {
                "model": "gpt-3.5-turbo",
                "messages": [
                    {"role": "system", "content": "You are a speech therapy assistant that creates reading practice materials."},
                    {"role": "user", "content": prompt}
                ],
                "max_tokens": 1000,
                "temperature": 0.7
            }

            response = requests.post(
                "https://api.openai.com/v1/chat/completions",
                headers=headers,
                json=data
            )

            response.raise_for_status()
            result = response.json()
            return result["choices"][0]["message"]["content"].strip()

        except Exception as e:
            logger.error(f"Error generating content with OpenAI: {str(e)}")
            return self._generate_mock_content(prompt)

    def _generate_mock_content(self, prompt: str) -> str:
        """Fallback mock content when no AI provider is available"""
        
        # Extract topic if possible
        topic = "general topics"
        if "about " in prompt:
            parts = prompt.split("about ")
            if len(parts) > 1:
                topic_part = parts[1].split(".")[0]
                topic = topic_part.strip()

        if "articulation" in prompt.lower():
            return json.dumps([
                {"word": "rainbow", "position": "initial"},
                {"word": "carrot", "position": "medial"},
                {"word": "car", "position": "final"},
                {"word": "read", "position": "initial"},
                {"word": "horror", "position": "medial"}
            ])

        elif "reading passage" in prompt.lower() or "paragraph" in prompt.lower():
            return f"""Here's a reading passage about {topic}. This contains various words designed for speech therapy practice."""

        elif "conversation" in prompt.lower():
            return json.dumps({
                "response": f"That's interesting! Tell me more about {topic}.",
                "technique_focus": "natural flow",
                "target_sounds": ["t", "m"],
                "follow_up_suggestions": [
                    "What aspects interest you most?",
                    "How did you learn about this?"
                ]
            })

        else:
            return f"Here's some content about {topic}."


# ====== ENHANCED CONVERSATION GENERATORS ======

# Available conversation roles
CONVERSATION_ROLES = {
    "friend": {
        "name": "Friendly Companion",
        "style": "casual, warm, and supportive",
        "topics": ["hobbies", "daily life", "interests", "weekend plans", "favorite things"],
        "personality": "enthusiastic, curious, and encouraging"
    },
    "teacher": {
        "name": "Patient Teacher",
        "style": "educational, clear, and encouraging",
        "topics": ["learning", "subjects", "study tips", "school experiences", "knowledge"],
        "personality": "knowledgeable, patient, and supportive"
    },
    "interviewer": {
        "name": "Professional Interviewer",
        "style": "professional, respectful, and interested",
        "topics": ["experience", "skills", "goals", "achievements", "background"],
        "personality": "attentive, thorough, and encouraging"
    },
    "therapist": {
        "name": "Speech Therapist",
        "style": "supportive, understanding, and helpful",
        "topics": ["progress", "practice", "feelings", "challenges", "successes"],
        "personality": "compassionate, patient, and motivating"
    },
    "shopkeeper": {
        "name": "Helpful Shop Assistant",
        "style": "helpful, friendly, and informative",
        "topics": ["products", "recommendations", "prices", "preferences", "needs"],
        "personality": "knowledgeable, patient, and service-oriented"
    },
    "doctor": {
        "name": "Medical Professional",
        "style": "professional, caring, and clear",
        "topics": ["health", "symptoms", "lifestyle", "wellness", "concerns"],
        "personality": "attentive, thorough, and reassuring"
    },
    "coach": {
        "name": "Motivational Coach",
        "style": "energetic, positive, and motivating",
        "topics": ["goals", "progress", "challenges", "strategies", "achievements"],
        "personality": "enthusiastic, supportive, and inspiring"
    },
    "tourist_guide": {
        "name": "Tour Guide",
        "style": "informative, engaging, and friendly",
        "topics": ["places", "history", "culture", "recommendations", "experiences"],
        "personality": "knowledgeable, enthusiastic, and welcoming"
    }
}


def initialize_conversation(
    role: str = "friend",
    topic: str = "general",
    user_name: Optional[str] = None,
    session_id: Optional[str] = None
) -> Dict[str, Any]:
    """
    Initialize a new conversation with specified role and topic
    
    Args:
        role: The AI's role in the conversation
        topic: Conversation topic
        user_name: Optional user name for personalization
        session_id: Unique session identifier
        
    Returns:
        Initial conversation setup with greeting and context
    """
    
    generator = AIContentGenerator()
    
    # Get role details
    role_info = CONVERSATION_ROLES.get(role, CONVERSATION_ROLES["friend"])
    
    # Create or retrieve conversation context
    if session_id:
        if session_id not in generator.conversation_contexts:
            generator.conversation_contexts[session_id] = ConversationContext(role, topic)
        context = generator.conversation_contexts[session_id]
    else:
        context = ConversationContext(role, topic)
    
    # Generate personalized greeting
    greeting_prompt = f"""
    You are playing the role of a {role_info['name']}.
    Your conversation style is {role_info['style']}.
    Your personality is {role_info['personality']}.
    
    Create a natural, welcoming greeting to start a conversation about {topic}.
    {f"The person's name is {user_name}." if user_name else ""}
    
    Make the greeting:
    - Warm and inviting
    - Appropriate for the role
    - Natural and conversational
    - Set up for easy continuation
    
    Format as JSON:
    {{
        "greeting": "Your greeting message",
        "suggested_topics": ["topic1", "topic2", "topic3"],
        "conversation_starters": ["question1", "question2"]
    }}
    """
    
    try:
        response = generator.generate_content(greeting_prompt)
        greeting_data = json.loads(response)
        
        # Add to conversation history
        context.add_turn("ai", greeting_data["greeting"])
        
        return {
            "role": role,
            "role_info": role_info,
            "topic": topic,
            "greeting": greeting_data["greeting"],
            "suggested_topics": greeting_data.get("suggested_topics", role_info["topics"][:3]),
            "conversation_starters": greeting_data.get("conversation_starters", []),
            "session_id": session_id or f"session_{datetime.now().timestamp()}"
        }
        
    except Exception as e:
        logger.error(f"Error initializing conversation: {str(e)}")
        # Fallback greeting
        greeting = f"Hello{f' {user_name}' if user_name else ''}! I'm excited to chat with you about {topic}. What would you like to discuss?"
        context.add_turn("ai", greeting)
        
        return {
            "role": role,
            "role_info": role_info,
            "topic": topic,
            "greeting": greeting,
            "suggested_topics": role_info["topics"][:3],
            "conversation_starters": [
                f"What interests you most about {topic}?",
                f"Tell me about your experience with {topic}."
            ],
            "session_id": session_id or f"session_{datetime.now().timestamp()}"
        }


def generate_ai_conversation_response(
    user_input: str,
    role: str = "friend",
    topic: str = "general",
    technique: str = "normal",
    session_id: Optional[str] = None,
    conversation_history: Optional[List[Dict]] = None
) -> Dict[str, Any]:
    """
    Generate a natural AI response based on role and conversation context
    
    Args:
        user_input: What the user said
        role: AI's conversation role
        topic: Current conversation topic
        technique: Speech technique being practiced
        session_id: Session identifier for context
        conversation_history: Optional conversation history
        
    Returns:
        AI response with practice suggestions and follow-ups
    """
    
    generator = AIContentGenerator()
    
    # Get role information
    role_info = CONVERSATION_ROLES.get(role, CONVERSATION_ROLES["friend"])
    
    # Get or create conversation context
    if session_id and session_id in generator.conversation_contexts:
        context = generator.conversation_contexts[session_id]
    else:
        context = ConversationContext(role, topic)
        if session_id:
            generator.conversation_contexts[session_id] = context
    
    # Add user input to history
    context.add_turn("user", user_input)
    
    # Build conversation history for context
    history_context = ""
    if conversation_history:
        history_context = "Recent conversation:\n"
        for turn in conversation_history[-5:]:  # Last 5 turns
            history_context += f"{turn['speaker']}: {turn['message']}\n"
    elif context.history:
        history_context = f"Recent conversation:\n{context.get_conversation_summary()}\n"
    
    # Technique-specific guidance
    technique_guidance = {
        "normal": "natural speech patterns",
        "prolonged": "opportunities to extend vowel sounds",
        "gentle_onset": "words starting with soft sounds or vowels",
        "easy_onset": "words beginning with continuous sounds (m, n, f, s)",
        "rhythmic": "rhythmic speaking patterns and natural pauses"
    }.get(technique, "natural speech patterns")
    
    # Create the response prompt
    prompt = f"""
    You are playing the role of a {role_info['name']}.
    Your conversation style is {role_info['style']}.
    Your personality is {role_info['personality']}.
    The conversation topic is: {topic}
    
    {history_context}
    
    The user just said: "{user_input}"
    
    Generate a natural, engaging response that:
    1. Directly addresses what the user said
    2. Stays in character for your role
    3. Keeps the conversation flowing naturally
    4. Shows genuine interest and engagement
    5. Provides opportunities for {technique_guidance}
    6. Varies your responses to avoid repetition
    
    Important guidelines:
    - Make each response unique and contextual
    - React naturally to the user's emotions or experiences
    - Ask follow-up questions that build on what they said
    - Share relevant thoughts or experiences when appropriate for the role
    - Keep responses conversational, not clinical
    
    Format your response as JSON:
    {{
        "response": "Your natural conversational response",
        "emotion": "The emotional tone (e.g., curious, excited, supportive)",
        "technique_focus": "Specific speech element to practice",
        "target_sounds": ["sound1", "sound2"],
        "follow_up_options": ["natural follow-up 1", "natural follow-up 2"],
        "conversation_tips": ["tip for continuing the conversation"]
    }}
    
    Provide only the JSON.
    """
    
    try:
        response = generator.generate_content(prompt)
        
        # Parse JSON response
        json_start = response.find("{")
        json_end = response.rfind("}") + 1
        if json_start >= 0 and json_end > json_start:
            json_str = response[json_start:json_end]
            response_data = json.loads(json_str)
            
            # Add AI response to history
            context.add_turn("ai", response_data["response"])
            
            # Ensure all fields are present
            return {
                "response": response_data.get("response", "That's interesting! Tell me more."),
                "emotion": response_data.get("emotion", "interested"),
                "technique_focus": response_data.get("technique_focus", technique_guidance),
                "target_sounds": response_data.get("target_sounds", ["r", "s"]),
                "follow_up_options": response_data.get("follow_up_options", [
                    "What else would you like to share?",
                    "How does that make you feel?"
                ]),
                "conversation_tips": response_data.get("conversation_tips", [
                    "Take your time to form your thoughts"
                ]),
                "turn_count": context.turn_count,
                "role": role,
                "topic": topic
            }
            
    except Exception as e:
        logger.error(f"Error generating conversation response: {str(e)}")
        
        # Generate fallback response based on role
        fallback_responses = {
            "friend": [
                "That's really interesting! I'd love to hear more about that.",
                "Oh wow, I hadn't thought about it that way. What made you think of that?",
                "That sounds like quite an experience! How did it make you feel?"
            ],
            "teacher": [
                "That's a thoughtful observation. Can you elaborate on that idea?",
                "Excellent point! What other connections can you make?",
                "I'm impressed by your thinking. What led you to that conclusion?"
            ],
            "interviewer": [
                "That's valuable insight. Could you provide a specific example?",
                "Interesting perspective. How has that shaped your approach?",
                "I appreciate you sharing that. What was the outcome?"
            ],
            "therapist": [
                "Thank you for sharing that with me. How are you feeling about it?",
                "That's an important observation. What would you like to explore further?",
                "I hear what you're saying. What matters most to you about this?"
            ]
        }
        
        responses = fallback_responses.get(role, fallback_responses["friend"])
        selected_response = random.choice(responses)
        
        # Add to history
        context.add_turn("ai", selected_response)
        
        return {
            "response": selected_response,
            "emotion": "supportive",
            "technique_focus": technique_guidance,
            "target_sounds": ["r", "s", "m"],
            "follow_up_options": [
                "Would you like to tell me more?",
                "What are your thoughts on this?"
            ],
            "conversation_tips": [
                "Take your time to express yourself clearly"
            ],
            "turn_count": context.turn_count,
            "role": role,
            "topic": topic
        }


def generate_dynamic_conversation_prompts(
    role: str = "friend",
    topic: str = "general",
    technique: str = "normal",
    difficulty: str = "medium",
    conversation_stage: str = "beginning"
) -> List[Dict[str, Any]]:
    """
    Generate conversation prompts that adapt to role and conversation stage
    
    Args:
        role: AI's conversation role
        topic: Conversation subject
        technique: Speech technique being practiced
        difficulty: Complexity level
        conversation_stage: beginning, middle, or ending
        
    Returns:
        List of dynamic conversation prompts
    """
    
    generator = AIContentGenerator()
    role_info = CONVERSATION_ROLES.get(role, CONVERSATION_ROLES["friend"])
    
    # Stage-specific guidance
    stage_guidance = {
        "beginning": "opening questions to establish rapport and interest",
        "middle": "deeper questions that build on established topics",
        "ending": "reflective questions or future-oriented topics"
    }.get(conversation_stage, "engaging questions")
    
    prompt = f"""
    Create conversation prompts for a {role_info['name']} discussing {topic}.
    
    Role characteristics:
    - Style: {role_info['style']}
    - Personality: {role_info['personality']}
    - Typical topics: {', '.join(role_info['topics'])}
    
    Generate 5-7 {stage_guidance} that:
    1. Feel natural for this role
    2. Encourage {difficulty} level responses
    3. Provide opportunities for {technique} speech practice
    4. Vary in style and approach
    5. Build engagement and rapport
    
    Each prompt should feel like something this character would naturally say.
    
    Format as JSON array:
    [
        {{
            "prompt": "The natural question or statement",
            "intent": "What the prompt aims to explore",
            "technique_focus": "Speech technique element",
            "target_sounds": ["sound1", "sound2"],
            "difficulty_notes": "Why this matches the difficulty level"
        }}
    ]
    
    Provide only the JSON array.
    """
    
    try:
        response = generator.generate_content(prompt)
        
        # Extract and parse JSON
        json_start = response.find("[")
        json_end = response.rfind("]") + 1
        if json_start >= 0 and json_end > json_start:
            json_str = response[json_start:json_end]
            prompts = json.loads(json_str)
            
            # Validate and enrich prompts
            validated_prompts = []
            for p in prompts:
                if isinstance(p, dict) and "prompt" in p:
                    validated_prompts.append({
                        "prompt": p["prompt"],
                        "intent": p.get("intent", "Explore user's thoughts"),
                        "technique_focus": p.get("technique_focus", technique),
                        "target_sounds": p.get("target_sounds", ["r", "s"]),
                        "difficulty_notes": p.get("difficulty_notes", "Appropriate for practice"),
                        "role": role,
                        "stage": conversation_stage
                    })
            
            return validated_prompts if validated_prompts else generate_fallback_dynamic_prompts(role, topic, conversation_stage)
            
    except Exception as e:
        logger.error(f"Error generating dynamic prompts: {str(e)}")
        return generate_fallback_dynamic_prompts(role, topic, conversation_stage)


def generate_fallback_dynamic_prompts(role: str, topic: str, stage: str) -> List[Dict[str, Any]]:
    """Generate fallback prompts when AI generation fails"""
    
    role_prompts = {
        "friend": {
            "beginning": [
                f"Hey! I've been really curious about {topic} lately. What got you interested in it?",
                f"So, what's your take on {topic}? I'd love to hear your thoughts!",
                f"You know, I was just thinking about {topic} the other day. Do you have any experience with it?"
            ],
            "middle": [
                f"That's so cool! What's been the most surprising thing about {topic} for you?",
                f"I love how passionate you are about this! What would you like to try next with {topic}?",
                f"Based on what you've shared, what advice would you give someone just starting with {topic}?"
            ]
        },
        "teacher": {
            "beginning": [
                f"Welcome! Today we're exploring {topic}. What do you already know about this subject?",
                f"Let's discuss {topic}. What questions do you have that you'd like us to address?",
                f"Before we dive deeper into {topic}, what aspects interest you the most?"
            ],
            "middle": [
                f"Excellent observations! How do you think this concept applies to real-world situations?",
                f"You're making great connections! What patterns have you noticed in {topic}?",
                f"That's a thoughtful analysis. Can you think of any examples that illustrate this principle?"
            ]
        },
        "interviewer": {
            "beginning": [
                f"Thank you for joining me today. Could you tell me about your background with {topic}?",
                f"I'd like to learn about your experience. How did you first become involved with {topic}?",
                f"Let's start with your journey. What initially attracted you to {topic}?"
            ],
            "middle": [
                f"That's impressive. Could you walk me through a specific situation where you applied this knowledge?",
                f"Building on that, what challenges have you encountered and how did you overcome them?",
                f"Your approach is interesting. What results have you seen from implementing these strategies?"
            ]
        }
    }
    
    # Get prompts for the specific role and stage
    prompts_list = role_prompts.get(role, role_prompts["friend"]).get(stage, role_prompts["friend"]["beginning"])
    
    # Format as proper response objects
    formatted_prompts = []
    for i, prompt in enumerate(prompts_list):
        formatted_prompts.append({
            "prompt": prompt,
            "intent": "Build rapport and explore topic",
            "technique_focus": "natural conversation flow",
            "target_sounds": ["r", "s", "t", "m"][i:i+2],
            "difficulty_notes": "Appropriate for speech practice",
            "role": role,
            "stage": stage
        })
    
    return formatted_prompts


# ====== ARTICULATION PRACTICE GENERATORS (keeping existing functions) ======

def generate_articulation_words(
    sound: str,
    difficulty: str = "medium",
    count: int = 5
) -> List[Dict[str, str]]:
    """Generate words that target a specific sound for articulation practice"""
    
    generator = AIContentGenerator()

    difficulty_desc = {
        "easy": "simple words with the sound in initial position mostly, suitable for young children",
        "medium": "moderately complex words with the sound in various positions",
        "hard": "more complex and longer words, including consonant clusters with the target sound"
    }.get(difficulty, "moderately complex words")

    prompt = f"""
    Generate {count} words that include the '{sound}' sound for speech therapy articulation practice.
    The words should be {difficulty_desc}.

    For each word, specify whether the target sound appears in the:
    - "initial" position (beginning of word)
    - "medial" position (middle of word)
    - "final" position (end of word)

    Format your response as a JSON array of objects, each with 'word' and 'position' fields:
    [
        {{"word": "example", "position": "initial"}},
        ...
    ]

    Provide only the JSON array with no additional text or explanations.
    """

    try:
        response = generator.generate_content(prompt)

        json_start = response.find("[")
        json_end = response.rfind("]") + 1
        if json_start >= 0 and json_end > json_start:
            json_str = response[json_start:json_end]
            words_list = json.loads(json_str)

            result = []
            for item in words_list[:count]:
                if isinstance(item, dict) and "word" in item and "position" in item:
                    result.append({
                        "word": item["word"].lower().strip(),
                        "position": item["position"].lower().strip()
                    })

            return result

        else:
            logger.error("Failed to extract JSON from AI response")
            return generate_fallback_articulation_words(sound, count)

    except Exception as e:
        logger.error(f"Error generating articulation words: {str(e)}")
        return generate_fallback_articulation_words(sound, count)


def generate_fallback_articulation_words(sound: str, count: int = 5) -> List[Dict[str, str]]:
    """Generate fallback articulation words when AI generation fails"""
    
    sound_map = {
        "r": [
            {"word": "red", "position": "initial"},
            {"word": "car", "position": "final"},
            {"word": "parent", "position": "medial"},
            {"word": "train", "position": "initial"},
            {"word": "borrow", "position": "medial"}
        ],
        "s": [
            {"word": "sun", "position": "initial"},
            {"word": "pass", "position": "final"},
            {"word": "messy", "position": "medial"},
            {"word": "snake", "position": "initial"},
            {"word": "bus", "position": "final"}
        ],
        "l": [
            {"word": "light", "position": "initial"},
            {"word": "ball", "position": "final"},
            {"word": "yellow", "position": "medial"},
            {"word": "link", "position": "initial"},
            {"word": "pillow", "position": "medial"}
        ],
        "sh": [
            {"word": "ship", "position": "initial"},
            {"word": "fish", "position": "final"},
            {"word": "washing", "position": "medial"},
            {"word": "shape", "position": "initial"},
            {"word": "brush", "position": "final"}
        ]
    }

    if sound in sound_map:
        return sound_map[sound][:count]
    else:
        return [
            {"word": f"{sound}ample", "position": "initial"},
            {"word": f"e{sound}ample", "position": "medial"},
            {"word": f"bas{sound}", "position": "final"},
            {"word": f"{sound}onder", "position": "initial"},
            {"word": f"ca{sound}e", "position": "medial"}
        ][:count]


# ====== READING PRACTICE GENERATORS (keeping existing functions) ======

def generate_reading_passage(
    topic: str,
    difficulty: str = "medium",
    format_type: str = "paragraph",
    vowel_focus: Optional[List[str]] = None,
    consonant_focus: Optional[List[str]] = None,
    sentence_count: Optional[int] = None,
    technique: Optional[str] = None
) -> str:
    """Generate a reading passage for fluency practice"""
    
    generator = AIContentGenerator()

    if not sentence_count and format_type == "sentences":
        sentence_count = 5

    reading_level = {
        "easy": "simple vocabulary and short sentences (1st-2nd grade level)",
        "medium": "moderate vocabulary and sentence structure (3rd-5th grade level)",
        "hard": "more complex vocabulary and varied sentence structure (6th-8th grade level)"
    }.get(difficulty, "moderate vocabulary and sentence structure")

    technique_guidance = ""
    if technique:
        if technique == "prolonged":
            technique_guidance = "Include words with long vowel sounds that can be stretched out for prolonged speech practice."
        elif technique == "gentle_onset" or technique == "easy_onset":
            technique_guidance = "Include words that begin with vowels or soft consonants for gentle/easy onset practice."
        elif technique == "rhythmic":
            technique_guidance = "Use a natural rhythm and include some repeated phrases for rhythmic speaking practice."

    vowel_instruction = ""
    if vowel_focus and len(vowel_focus) > 0:
        vowel_str = ", ".join([f"'{v}'" for v in vowel_focus])
        vowel_instruction = f"Emphasize words containing these vowel sounds: {vowel_str}."

    consonant_instruction = ""
    if consonant_focus and len(consonant_focus) > 0:
        consonant_str = ", ".join([f"'{c}'" for c in consonant_focus])
        consonant_instruction = f"Emphasize words containing these consonant sounds: {consonant_str}."

    format_instruction = {
        "paragraph": f"Create a cohesive paragraph about {topic} with 4-6 sentences.",
        "sentences": f"Create {sentence_count} individual, standalone sentences about {topic}.",
        "story": f"Create a short story about {topic} with a beginning, middle, and end."
    }.get(format_type, f"Write about {topic}")

    prompt = f"""
    {format_instruction}

    The text should use {reading_level}.
    {vowel_instruction}
    {consonant_instruction}
    {technique_guidance}

    Make the content engaging and appropriate for speech therapy practice.
    Provide only the text with no additional explanations or commentary.
    """

    try:
        content = generator.generate_content(prompt)
        if not content or len(content.strip()) < 10:
            logger.error("AI returned empty or very short content")
            return generate_fallback_reading_passage(topic, format_type, sentence_count)
        return content
    except Exception as e:
        logger.error(f"Error generating reading passage: {str(e)}")
        return generate_fallback_reading_passage(topic, format_type, sentence_count)


def generate_fallback_reading_passage(
    topic: str,
    format_type: str,
    sentence_count: Optional[int] = 5
) -> str:
    """Generate fallback reading passage when AI generation fails"""
    
    topics = {
        "technology": [
            "Smartphones have revolutionized how we communicate and access information.",
            "Artificial intelligence is being integrated into many everyday devices.",
            "Cloud computing allows us to store vast amounts of data remotely.",
            "Virtual reality creates immersive experiences for gaming and education.",
            "Robotics is advancing rapidly in manufacturing and healthcare."
        ],
        "travel": [
            "Paris attracts millions of visitors to see the Eiffel Tower and Louvre Museum.",
            "Japan's bullet trains make traveling between cities incredibly efficient.",
            "The Great Barrier Reef in Australia is the world's largest coral reef system.",
            "Machu Picchu reveals the impressive engineering skills of the Inca civilization.",
            "Venice is famous for its canals, with boats as the main transportation method."
        ],
        "animals": [
            "Elephants are highly intelligent and have complex social structures.",
            "Dolphins communicate using a series of clicks, whistles, and body movements.",
            "Chameleons can change color to match their surroundings and regulate temperature.",
            "Eagles have incredible eyesight and can spot prey from great distances.",
            "Octopuses have three hearts and can solve complex puzzles."
        ]
    }

    sentences = topics.get(topic.lower(), [
        f"{topic} is a fascinating subject that continues to evolve over time.",
        f"Many people are interested in learning more about {topic} through books and online resources.",
        f"Experts in {topic} often share their knowledge through lectures and publications.",
        f"The history of {topic} reveals interesting patterns and developments.",
        f"Modern advances in {topic} have changed how we think about many aspects of life."
    ])

    if format_type == "sentences":
        return " ".join(sentences[:sentence_count])
    elif format_type == "story":
        return f"""Once upon a time, there was a curious student who became interested in {topic}.
        They discovered that {sentences[0].lower()}
        As they learned more, they found that {sentences[1].lower()}
        Their research showed that {sentences[2].lower()}
        The most surprising thing they learned was that {sentences[3].lower()}
        They shared their knowledge with friends, explaining how {sentences[4].lower()}
        Everyone was impressed by how much they had learned about {topic}."""
    else:  # paragraph
        return " ".join(sentences[:5])


# ====== UTILITY FUNCTIONS ======

def get_conversation_summary(session_id: str) -> Dict[str, Any]:
    """
    Get a summary of the conversation session
    
    Args:
        session_id: Session identifier
        
    Returns:
        Summary with statistics and key points
    """
    
    generator = AIContentGenerator()
    
    if session_id not in generator.conversation_contexts:
        return {
            "session_id": session_id,
            "status": "not_found",
            "message": "No conversation found with this session ID"
        }
    
    context = generator.conversation_contexts[session_id]
    
    # Calculate statistics
    user_messages = [turn for turn in context.history if turn["speaker"] == "user"]
    ai_messages = [turn for turn in context.history if turn["speaker"] == "ai"]
    
    # Extract topics discussed
    topics_discussed = []
    for turn in user_messages:
        # Simple topic extraction (could be enhanced with NLP)
        words = turn["message"].lower().split()
        for word in words:
            if len(word) > 5 and word not in topics_discussed:
                topics_discussed.append(word)
    
    return {
        "session_id": session_id,
        "role": context.role,
        "main_topic": context.topic,
        "turn_count": context.turn_count,
        "total_exchanges": len(context.history),
        "user_messages": len(user_messages),
        "ai_messages": len(ai_messages),
        "topics_discussed": topics_discussed[:5],
        "conversation_duration": None,  # Could calculate from timestamps
        "last_activity": context.history[-1]["timestamp"] if context.history else None
    }


def reset_conversation(session_id: str) -> bool:
    """
    Reset a conversation session
    
    Args:
        session_id: Session identifier
        
    Returns:
        Success status
    """
    
    generator = AIContentGenerator()
    
    if session_id in generator.conversation_contexts:
        del generator.conversation_contexts[session_id]
        return True
    return False


def list_available_roles() -> List[Dict[str, Any]]:
    """
    Get list of available conversation roles
    
    Returns:
        List of role information
    """
    
    roles = []
    for role_id, role_info in CONVERSATION_ROLES.items():
        roles.append({
            "id": role_id,
            "name": role_info["name"],
            "description": role_info["style"],
            "suggested_topics": role_info["topics"],
            "personality": role_info["personality"]
        })
    return roles


# ====== TEST FUNCTIONS ======

def test_conversation_system():
    """Test the conversation system with different roles"""
    
    print("Testing Enhanced Conversation System")
    print("=" * 50)
    
    # Test 1: Initialize conversations with different roles
    print("\n1. Testing role initialization:")
    for role in ["friend", "teacher", "interviewer"]:
        result = initialize_conversation(
            role=role,
            topic="hobbies",
            user_name="Sam"
        )
        print(f"\n{role.upper()} Role:")
        print(f"Greeting: {result['greeting']}")
        print(f"Suggested topics: {result['suggested_topics']}")
    
    # Test 2: Generate responses
    print("\n\n2. Testing conversation responses:")
    session_id = "test_session_123"
    
    # Initialize a friend conversation
    init_result = initialize_conversation(
        role="friend",
        topic="cooking",
        session_id=session_id
    )
    
    # Simulate user responses
    user_inputs = [
        "I love trying new recipes, especially Italian food!",
        "My favorite dish to make is homemade pasta with fresh tomatoes.",
        "I learned from my grandmother who was an amazing cook."
    ]
    
    for user_input in user_inputs:
        print(f"\nUser: {user_input}")
        
        response = generate_ai_conversation_response(
            user_input=user_input,
            role="friend",
            topic="cooking",
            technique="normal",
            session_id=session_id
        )
        
        print(f"AI: {response['response']}")
        print(f"Emotion: {response['emotion']}")
        print(f"Follow-up options: {response['follow_up_options']}")
    
    # Test 3: Get conversation summary
    print("\n\n3. Testing conversation summary:")
    summary = get_conversation_summary(session_id)
    print(f"Total turns: {summary['turn_count']}")
    print(f"Topics discussed: {summary['topics_discussed']}")
    
    print("\n" + "=" * 50)
    print("Testing complete!")


if __name__ == "__main__":
    # Run tests
    test_conversation_system()