src/utils/text_explainer.py

"""
Text Explanation utilities using Mistral AI.
Splits text by markdown headings and generates contextual explanations for each section.
Maintains chat history to provide coherent explanations that build upon previous sections.
"""

import os
import re
from typing import List, Dict, Tuple, Optional
from mistralai import Mistral


class TextExplainer:
    """Generate explanations for text sections using Mistral AI."""
    
    def __init__(self):
        """Initialize the text explainer with Mistral AI client."""
        self.api_key = os.environ.get("MISTRAL_API_KEY")
        if not self.api_key:
            raise ValueError("MISTRAL_API_KEY environment variable is required")
        self.client = Mistral(api_key=self.api_key)
        self.chat_history = []

    def get_topic(self, text: str) -> Optional[str]:
        """
        Extract the main topic from the text using Mistral AI with structured output.
        
        Args:
            text: Input text to analyze
            
        Returns:
            Main topic as a string or None if not found
        """
        try:
            # Define the JSON schema for structured output
            topic_schema = {
                "type": "json_schema",
                "json_schema": {
                    "schema": {
                        "type": "object",
                        "properties": {
                            "main_topic": {
                                "type": "string",
                                "title": "Main Topic",
                                "description": "The primary / general topic or subject of the text"
                            },
                        },
                        "required": ["main_topic"],
                        "additionalProperties": False
                    },
                    "name": "topic_extraction",
                    "strict": True
                }
            }
            
            response = self.client.chat.complete(
                model="ministral-8b-2410",  # Using a more recent model that supports structured output
                messages=[
                    {
                        "role": "system",
                        "content": "You are an expert in summarizing texts. Extract the main topic from the provided text."
                    },
                    {
                        "role": "user",
                        "content": f"Analyze this text and extract the main topic:\n\n{text[:2000]}..."  # Limit to first 2000 characters for performance
                    }
                ],
                temperature=0.3,  # Lower temperature for more consistent structured output
                max_tokens=200,
                response_format=topic_schema
            )
            
            if hasattr(response, 'choices') and response.choices:
                # Parse the structured JSON response
                import json
                try:
                    topic_data = json.loads(response.choices[0].message.content)
                    main_topic = topic_data.get("main_topic", "").strip()
                    confidence = topic_data.get("confidence", 0.0)
                    secondary_topics = topic_data.get("secondary_topics", [])
                    
                    # Log the structured output for debugging
                    print(f"📊 Topic extraction - Main: '{main_topic}', Confidence: {confidence:.2f}")
                    if secondary_topics:
                        print(f"🔍 Secondary topics: {', '.join(secondary_topics)}")
                    
                    return main_topic if main_topic else None
                except json.JSONDecodeError as json_err:
                    print(f"Error parsing JSON response: {json_err}")
                    # Fallback to raw content if JSON parsing fails
                    return response.choices[0].message.content.strip()
            return None
        except Exception as e:
            print(f"Error extracting topic: {str(e)}")
            return None

    def split_text_by_headings(self, text: str) -> List[Dict[str, str]]:
        """
        Split text into sections based on markdown headings.
        
        Args:
            text: Input text with markdown headings
            
        Returns:
            List of dictionaries with 'heading' and 'content' keys
        """
        if not text:
            return []
        
        # Split by markdown headings (# ## ### etc.)
        sections = []
        
        # Regex to find headings and their content
        # Matches: # Heading, ## Heading, ### Heading, etc.
        heading_pattern = r'^(#{1,6})\s+(.+?)$'
        
        lines = text.split('\n')
        current_heading = None
        current_content = []
        current_level = 0
        
        for line in lines:
            heading_match = re.match(heading_pattern, line.strip())
            
            if heading_match:
                # Save previous section if it exists
                if current_heading and current_content:
                    content_text = '\n'.join(current_content).strip()
                    if content_text:  # Only add if there's actual content
                        sections.append({
                            'heading': current_heading,
                            'content': content_text,
                            'level': current_level
                        })
                
                # Start new section
                level = len(heading_match.group(1))  # Count the # characters
                current_heading = heading_match.group(2).strip()
                current_level = level
                current_content = []
            else:
                # Add line to current content if we have a heading
                if current_heading is not None:
                    current_content.append(line)
        
        # Don't forget the last section
        if current_heading and current_content:
            content_text = '\n'.join(current_content).strip()
            if content_text:
                sections.append({
                    'heading': current_heading,
                    'content': content_text,
                    'level': current_level
                })
        
        # If no headings found, treat entire text as one section
        if not sections and text.strip():
            sections.append({
                'heading': 'Document Content',
                'content': text.strip(),
                'level': 1
            })
        return sections

    def generate_explanation(self, topic: str, heading: str, content: str, section_number: int = 1, total_sections: int = 1) -> str:
        """
        Generate an explanation for a text section using Mistral AI with chat history context.
        
        Args:
            topic: General topic of the document
            heading: Section heading
            content: Section content
            section_number: Current section number (for context)
            total_sections: Total number of sections (for context)
            
        Returns:
            Generated explanation in simple terms
        """
        try:
            # Build the current user message
            prompt = f"""
            **Section {section_number} of {total_sections}**
            **Section Heading:** {heading}

            **Section Content:**
            {content}

            **Your Explanation:**"""

            # If this is the first section, initialize with system prompt
            if section_number == 1:
                system_prompt = f"""You are an expert teacher who explains complex topics in simple, easy-to-understand terms.

                I will give you sections of text with their headings on the topic of "{topic}", and I want you to explain what each section is about in simple language, by breaking down any complex concepts or terminology. You should also explain why this information might be important or useful, use examples or analogies when helpful, and keep the explanation engaging and educational.

                Make your explanation clear enough for someone without prior knowledge of the topic to understand. As you explain each section, consider how it relates to the previous sections you've already explained to provide coherent, contextual explanations throughout the document.

                Do not mention anything far irrelevant from the topic of "{topic}". Do not repeat information unnecessarily, but build on previous explanations to create a comprehensive understanding of the topic. Avoid using the term 'section' and use the actual section heading instead. No need to mention the section number in your explanation.
                """
                
                # Initialize chat history with system message
                self.chat_history = [
                    {
                        "role": "system",
                        "content": system_prompt
                    }
                ]

            # Check if content is too small (less than 200 characters)
            if len(content) < 200:
                print(f"📋 Skipping API call for short content in '{heading}' ({len(content)} chars < 200)")
                # Add the user prompt to chat history for context in subsequent queries
                self.chat_history.append({
                    "role": "user",
                    "content": prompt
                })
                # Return a simple message indicating the content was too short
                return f"This section contains minimal content ({len(content)} characters). The information has been noted for context in subsequent explanations."

            # Add the current user message to chat history
            self.chat_history.append({
                "role": "user",
                "content": prompt
            })

            # Call Mistral AI for explanation with full chat history
            response = self.client.chat.complete(
                model="mistral-small-2503",
                messages=self.chat_history,
                temperature=0.7,  # Some creativity but still focused
                # max_tokens=1000   # Reasonable explanation length
            )
            
            # Extract the explanation from response
            if hasattr(response, 'choices') and response.choices:
                explanation = response.choices[0].message.content
                
                # Add the assistant's response to chat history
                self.chat_history.append({
                    "role": "assistant",
                    "content": explanation
                })
                
                return explanation.strip()
            else:
                return f"Could not generate explanation for section: {heading}"
        
        except Exception as e:
            print(f"Error generating explanation for '{heading}': {str(e)}")
            return f"Error generating explanation for this section: {str(e)}"
        
    def explain_all_sections(self, text: str) -> List[Dict[str, str]]:
        """
        Split text by headings and generate explanations for all sections with chat history context.
        
        Args:
            text: Input text with markdown headings
            
        Returns:
            List of dictionaries with 'heading', 'content', 'explanation', and 'level' keys
        """
        sections = self.split_text_by_headings(text)
        
        if not sections:
            return []
        
        print(f"🔍 Found {len(sections)} sections to explain...")
        
        # Extract the main topic from the text
        print("🎯 Extracting main topic...")
        topic = self.get_topic(text)
        if topic:
            print(f"📋 Main topic identified: {topic}")
        else:
            topic = "General Content"  # Fallback topic
            print("⚠️ Could not identify main topic, using fallback")
        
        # Reset chat history for new document
        self.chat_history = []
        
        explained_sections = []
        
        for i, section in enumerate(sections, 1):
            print(f"📝 Generating explanation for section {i}/{len(sections)}: {section['heading'][:50]}...")
            
            # Pass topic, section content, and context information
            explanation = self.generate_explanation(
                topic,
                section['heading'], 
                section['content'], 
                section_number=i, 
                total_sections=len(sections)
            )
            
            explained_sections.append({
                'heading': section['heading'],
                'content': section['content'],
                'explanation': explanation,
                'level': section['level']
            })
        
        print(f"✅ Generated explanations for all {len(explained_sections)} sections")
        return explained_sections
    
    def reset_chat_history(self):
        """Reset the chat history for a new document or conversation."""
        self.chat_history = []
    
    def get_chat_history(self) -> List[Dict[str, str]]:
        """Get the current chat history for debugging purposes."""
        return self.chat_history.copy()
    
    def get_chat_history_summary(self) -> str:
        """Get a summary of the current chat history."""
        if not self.chat_history:
            return "No chat history available."
        
        summary = f"Chat history contains {len(self.chat_history)} messages:\n"
        for i, message in enumerate(self.chat_history, 1):
            role = message['role']
            content_preview = message['content'][:100] + "..." if len(message['content']) > 100 else message['content']
            summary += f"{i}. {role.upper()}: {content_preview}\n"
        
        return summary

    def format_explanations_for_display(self, explained_sections: List[Dict[str, str]]) -> str:
        """
        Concatenate only the explanations from all sections for display, filtering out placeholder explanations for minimal content.
        Args:
            explained_sections: List of sections with explanations
        Returns:
            Concatenated explanations as a single string
        """
        if not explained_sections:
            return "No sections found to explain."
        skip_phrase = "This section contains minimal content"
        return "\n\n".join(
            section['explanation']
            for section in explained_sections
            if section.get('explanation') and not section['explanation'].strip().startswith(skip_phrase)
        )