"""
AI Chat Management Routes with Multiple Modes
Handles chat sessions, messages, and history for the Advanced AI Chat feature
Supports multiple specialized chat modes with custom prompts
"""
import json
import logging
import asyncio
import sys
import os
from datetime import datetime
from flask import Blueprint, request, jsonify

# Add parent directory to import aimodel_selector
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from utils import require_api_key, encrypt_response, add_keep_alive
from database import get_cached_connection, optimized_execute
import mysql.connector
from database import get_db_connection

# Import central AI provider instead of g4f directly
from services.ai_provider import ai_generate, ai_generate_vision

# Import AI model configuration
from aimodel_selector.models_config import get_default_model, get_model_config

# Import subscription and token management services
from services.token_manager import TokenManager
from services.subscription_manager import SubscriptionManager

# Import AI tools for tool-based chat
from routes.ai_tools import AI_TOOLS, get_tool_system_prompt, execute_tool, parse_tool_calls, create_result_card

# Set up logging
logger = logging.getLogger(__name__)

# Create blueprint
ai_chat_bp = Blueprint('ai_chat', __name__)

def create_error_response(message, status_code=400):
    """Create standardized error response"""
    return jsonify({"status": "error", "message": message}), status_code

def format_api_response(data):
    """Format successful API response"""
    return jsonify({"status": "success", "data": data})

# AI Chat Modes Configuration - Dynamically loads models from database
CHAT_MODES = {
    'general': 'General Assistant',
    'tutor': 'Study Tutor',
    'coding': 'Code Helper',
    'language': 'Language Teacher',
    'math': 'Math Solver',
    'research': 'Research Assistant',
    'writing': 'Writing Coach',
    'creative': 'Creative Partner'
}

def get_chat_mode_config(mode_name, user_id=None):
    """Get chat mode configuration with dynamic model selection"""
    # Get default model from database
    try:
        default_model = get_default_model(user_id=user_id)
        model_config = get_model_config(default_model)
        provider = model_config.get('provider', 'PollinationsAI')
        model = model_config.get('model', 'geminisearch')
    except:
        provider = 'PollinationsAI'
        model = 'geminisearch'
    
    # Shared formatting instructions appended to all mode prompts
    FORMAT_INSTRUCTIONS = (
        '\n\n**RESPONSE FORMATTING:**\n'
        '- Use clean, well-structured **Markdown** with headings (#, ##), bold, lists, and tables as appropriate.\n'
        '- For math: use `$$...$$` for display equations, `$...$` for inline math. Use double backslashes for LaTeX commands (e.g., `\\\\frac{a}{b}`).\n'
        '- For diagrams and visual explanations: use ```mermaid code blocks for flowcharts, sequence diagrams, class diagrams, state diagrams, ER diagrams, Gantt charts, pie charts, mind maps, or timelines when they help illustrate concepts. Our app renders these as interactive visuals.\n'
        '- For geometric figures, circuits, coordinate planes, or custom illustrations: use inline SVG (`<svg>...</svg>`). Our app renders SVG as actual graphics.\n'
        '- Use diagrams proactively when they clarify relationships, processes, hierarchies, or data flows better than text alone.\n'
        '- Use code blocks with language tags (```python, ```javascript, etc.) for code.\n'
    )

    modes = {
        'general': {
            'name': 'General Assistant',
            'prompt': 'You are "Sathi AI", the friendly AI learning companion of Ashlya Academy. Provide clear, accurate, and well-structured responses. Use examples, analogies, and visual aids (diagrams/charts) when they help explain concepts. Remember: "Sathi" means "companion/friend" — be warm, supportive, and genuinely helpful!' + FORMAT_INSTRUCTIONS
        },
        'tutor': {
            'name': 'Study Tutor',
            'prompt': 'You are "Sathi AI", the friendly AI learning companion of Ashlya Academy, specializing as a study tutor. Help students deeply understand concepts through clear step-by-step explanations, real-world analogies, and visual diagrams. Break complex topics into digestible parts. Use flowcharts or mind maps to show concept relationships. Create study plans when asked. Be patient, encouraging, and thorough — like a brilliant friend who genuinely wants you to succeed!' + FORMAT_INSTRUCTIONS
        },
        'coding': {
            'name': 'Code Helper',
            'prompt': 'You are "Sathi AI", the friendly AI learning companion of Ashlya Academy, specializing as a coding mentor. Help with programming problems, code review, debugging, architecture, and best practices. Provide clean, well-commented code examples with clear explanations of the reasoning. Use mermaid diagrams to illustrate system architecture, class hierarchies, data flows, and algorithm flowcharts when they aid understanding. Teach the "why" behind the code, not just the "how"!' + FORMAT_INSTRUCTIONS
        },
        'language': {
            'name': 'Language Teacher',
            'prompt': 'You are "Sathi AI", the friendly AI learning companion of Ashlya Academy, specializing as a language teacher. Help with grammar, vocabulary, pronunciation, translations, and learning strategies. Provide clear examples with contextual usage, common mistakes to avoid, and memory aids. Use tables for comparisons (e.g., verb conjugations, tense charts). Support multiple languages with warm, encouraging explanations!' + FORMAT_INSTRUCTIONS
        },
        'math': {
            'name': 'Math Solver',
            'prompt': 'You are "Sathi AI", the friendly AI learning companion of Ashlya Academy, specializing in mathematics. Solve problems with clear, step-by-step working shown. Explain the intuition behind formulas and techniques, not just the mechanics. Use SVG diagrams for geometric figures, coordinate planes, graphs, and visual proofs. Use mermaid diagrams for decision trees or problem-solving flowcharts. Always show every step and make math feel approachable!' + FORMAT_INSTRUCTIONS
        },
        'research': {
            'name': 'Research Assistant',
            'prompt': 'You are "Sathi AI", the friendly AI learning companion of Ashlya Academy, specializing as a research assistant. Help gather information, summarize content, analyze arguments, fact-check, and structure research. Use mermaid diagrams for research frameworks, timelines, and concept relationships. Organize findings with clear headings, tables, and bullet points. Provide balanced perspectives and cite reasoning. Be a supportive guide through the research process!' + FORMAT_INSTRUCTIONS
        },
        'writing': {
            'name': 'Writing Coach',
            'prompt': 'You are "Sathi AI", the friendly AI learning companion of Ashlya Academy, specializing as a writing coach. Help improve writing with constructive, specific feedback. Assist with grammar, style, structure, tone, and creative expression. Show before/after examples to illustrate improvements. Use flowcharts for essay structure and argument flow when helpful. Be encouraging and help build the student\'s confidence as a writer!' + FORMAT_INSTRUCTIONS
        },
        'creative': {
            'name': 'Creative Partner',
            'prompt': 'You are "Sathi AI", the friendly AI learning companion of Ashlya Academy, specializing as a creative partner. Generate innovative ideas, solve creative problems, and provide inspiration. Use mind maps and flowcharts to brainstorm and visualize ideas. Think outside the box, explore multiple angles, and be an enthusiastic collaborator who builds on ideas!' + FORMAT_INSTRUCTIONS
        }
    }
    
    mode_data = modes.get(mode_name, modes['general'])
    mode_data['provider'] = provider
    mode_data['model'] = model
    return mode_data

def get_user_model_preference(user_id):
    """Get user's preferred AI model from database or return default"""
    return get_default_model(user_id=user_id)

# Language code to name mapping for AI prompts
LANGUAGE_NAMES = {
    'en-US': 'English',
    'hi-IN': 'Hindi', 
    'ne-NP': 'Nepali',
    'es-ES': 'Spanish',
    'fr-FR': 'French',
    'de-DE': 'German',
    'zh-CN': 'Chinese',
    'ja-JP': 'Japanese',
    'ko-KR': 'Korean'
}

def generate_ai_response(user_message, mode='general', conversation_history=None, images=None, model_override=None, provider_override=None, user_id=None, language='en-US'):
    """Generate AI response using specified mode and conversation history with dynamic model selection
    
    Args:
        user_message: User's message text
        mode: Chat mode (general, tutor, math, etc.)
        conversation_history: Previous messages for context
        images: List of image data URIs if vision request [[data_uri, filename], ...]
        model_override: Override model selection
        provider_override: Override provider selection
        user_id: User ID to fetch model preference (if authenticated)
        language: Language code for AI response (e.g., 'en-US', 'hi-IN', 'ne-NP')
    """
    try:
        mode_config = get_chat_mode_config(mode, user_id=user_id)
        system_prompt = mode_config['prompt']
        provider = provider_override or mode_config['provider']
        
        # Add language instruction to system prompt if not English
        if language and language != 'en-US':
            language_name = LANGUAGE_NAMES.get(language, language)
            system_prompt = f"{system_prompt}\n\n**IMPORTANT: Respond in {language_name} language.**"
        
        # Use model selection priority: override > user preference > system default
        if model_override:
            model_name = model_override
        else:
            # Get user's preferred model from database
            model_name = get_user_model_preference(user_id)
        
        # Enhance system prompt if images are attached
        if images and len(images) > 0:
            vision_instructions = (
                "\n\n**IMPORTANT**: The user has attached an image. "
                "First, carefully analyze the image provided. "
                "Describe what you see in detail, then answer the user's question based on both "
                "the visual content and your knowledge. "
                "If the question is about the image, provide specific details from what you observe. "
                "If the image provides context, use it to give a more accurate and relevant response."
            )
            system_prompt = system_prompt + vision_instructions
            
            # Use Copilot for better vision support if not overridden
            if not provider_override:
                provider = "Copilot"
        
        # Prepare messages for context
        messages = [{"role": "system", "content": system_prompt}]
        
        # Add conversation history for context (last 10 messages)
        if conversation_history:
            for msg in conversation_history[-10:]:
                if msg['message_type'] == 'user':
                    messages.append({"role": "user", "content": msg['content']})
                elif msg['message_type'] == 'assistant':
                    messages.append({"role": "assistant", "content": msg['content']})
        
        # Add current user message
        messages.append({"role": "user", "content": user_message})
        
        # Generate response using central AI provider
        if images and len(images) > 0:
            # Vision request with images
            logger.info(f"Vision request with {len(images)} image(s)")
            result = ai_generate_vision(
                messages=messages,
                images=images,
                model=model_name,
            )
        else:
            # Standard text request
            result = ai_generate(
                messages=messages,
                model=model_name,
            )
        
        return result.content
        
    except Exception as e:
        logger.error(f"Error generating AI response: {str(e)}")
        return f"I apologize, but I'm having trouble processing your request right now. Please try again later."


def generate_ai_response_with_tools(user_message, mode='general', conversation_history=None, images=None, 
                                     model_override=None, provider_override=None, user_id=None, 
                                     language='en-US', enable_tools=True):
    """Generate AI response with tool execution support (Cursor-style).
    
    Returns a structured response with text blocks and tool results.
    
    Args:
        user_message: User's message text
        mode: Chat mode (general, tutor, math, etc.)
        conversation_history: Previous messages for context
        images: List of image data URIs if vision request
        model_override: Override model selection
        provider_override: Override provider selection
        user_id: User ID for tool execution and model preference
        language: Language code for AI response
        enable_tools: Whether to enable tool execution
        
    Returns:
        dict: {
            'response_blocks': [
                {'type': 'text', 'content': '...'},
                {'type': 'tool_call', 'tool': 'search_notes', 'params': {...}, 'results': [...]}
            ],
            'raw_response': str,
            'tokens_used': int
        }
    """
    try:
        mode_config = get_chat_mode_config(mode, user_id=user_id)
        base_prompt = mode_config['prompt']
        provider = provider_override or mode_config['provider']
        
        # Build system prompt with tools if enabled
        if enable_tools and user_id:
            tool_prompt = get_tool_system_prompt()
            system_prompt = f"{base_prompt}\n\n{tool_prompt}"
        else:
            system_prompt = base_prompt
        
        # Add language instruction if not English
        if language and language != 'en-US':
            language_name = LANGUAGE_NAMES.get(language, language)
            system_prompt = f"{system_prompt}\n\n**IMPORTANT: Respond in {language_name} language.**"
        
        # Use model selection priority: override > user preference > system default
        if model_override:
            model_key = model_override
        else:
            model_key = get_user_model_preference(user_id)
        
        # Get actual model config from database
        model_config = get_model_config(model_key)
        actual_model_name = model_config.get('model', model_key)
        backend_type = model_config.get('backend_type', 'groq')
        
        # Handle vision requests
        if images and len(images) > 0:
            vision_instructions = (
                "\n\n**IMPORTANT**: The user has attached an image. "
                "First, carefully analyze the image provided. "
                "Describe what you see in detail, then answer the user's question."
            )
            system_prompt = system_prompt + vision_instructions
        
        # Prepare initial messages
        messages = [{"role": "system", "content": system_prompt}]
        
        # Add conversation history (last 10 messages)
        if conversation_history:
            for msg in conversation_history[-10:]:
                if msg['message_type'] == 'user':
                    messages.append({"role": "user", "content": msg['content']})
                elif msg['message_type'] == 'assistant':
                    # Strip tool tags from history for cleaner context
                    content = msg['content']
                    content = re.sub(r'<tool[^>]*>.*?</tool>', '[Tool was used]', content, flags=re.DOTALL)
                    content = re.sub(r'<tool_result>.*?</tool_result>', '', content, flags=re.DOTALL)
                    messages.append({"role": "assistant", "content": content})
        
        # Add current user message
        messages.append({"role": "user", "content": user_message})
        
        logger.info(f"Calling AI provider with model: {actual_model_name}, backend: {backend_type}")
        
        # Generate initial AI response
        raw_response = None
        try:
            if images and len(images) > 0:
                result = ai_generate_vision(
                    messages=messages,
                    images=images,
                    model=actual_model_name,
                )
            else:
                result = ai_generate(
                    messages=messages,
                    model=actual_model_name,
                )
            raw_response = result.content
        except Exception as model_error:
            logger.warning(f"Model {actual_model_name} failed: {model_error}. Falling back to default.")
            try:
                if images and len(images) > 0:
                    result = ai_generate_vision(
                        messages=messages,
                        images=images,
                    )
                else:
                    result = ai_generate(messages=messages)
                raw_response = result.content
                logger.info(f"Fallback succeeded with model: {result.model}")
            except Exception as fallback_error:
                logger.error(f"Fallback model also failed: {fallback_error}")
                raise fallback_error
        
        logger.info(f"AI raw response (first 200 chars): {raw_response[:200]}")
        
        # Parse response for tool calls
        response_blocks = []
        
        if enable_tools and user_id:
            tool_calls = parse_tool_calls(raw_response)
            
            if tool_calls:
                logger.info(f"Found {len(tool_calls)} tool calls in response")
                
                # Split response by tool calls and process
                remaining_text = raw_response
                
                for tc in tool_calls:
                    # Get text before this tool call
                    parts = remaining_text.split(tc['raw'], 1)
                    
                    if parts[0].strip():
                        response_blocks.append({
                            'type': 'text',
                            'content': parts[0].strip()
                        })
                    
                    # Execute the tool
                    logger.info(f"Executing tool: {tc['tool']} with params: {tc['params']}")
                    tool_result = execute_tool(tc['tool'], tc['params'], int(user_id))
                    
                    response_blocks.append({
                        'type': 'tool_call',
                        'tool': tc['tool'],
                        'params': tc['params'],
                        'status': 'success' if tool_result.get('success') else 'error',
                        'results': tool_result.get('results', []),
                        'error': tool_result.get('error')
                    })
                    
                    # Update remaining text
                    remaining_text = parts[1] if len(parts) > 1 else ''
                
                # Add any remaining text after all tool calls
                if remaining_text.strip():
                    response_blocks.append({
                        'type': 'text',
                        'content': remaining_text.strip()
                    })
                
                # If tools were used, we can optionally ask AI to continue with results
                # For now, we'll let the frontend display the results and AI's existing commentary
                
            else:
                # No tool calls - just text response
                response_blocks.append({
                    'type': 'text',
                    'content': raw_response
                })
        else:
            # Tools disabled - plain text response
            response_blocks.append({
                'type': 'text',
                'content': raw_response
            })
        
        # Calculate tokens used (rough estimate)
        input_tokens = len(user_message.split())
        output_tokens = len(raw_response.split())
        
        return {
            'response_blocks': response_blocks,
            'raw_response': raw_response,
            'tokens_used': input_tokens + output_tokens,
            'input_tokens': input_tokens,
            'output_tokens': output_tokens,
            'model_used': actual_model_name,
            'tools_used': [tc['tool'] for tc in tool_calls] if enable_tools and user_id else []
        }
        
    except Exception as e:
        logger.error(f"Error generating AI response with tools: {str(e)}")
        return {
            'response_blocks': [{'type': 'text', 'content': "I apologize, but I'm having trouble processing your request right now. Please try again later."}],
            'raw_response': '',
            'tokens_used': 0,
            'input_tokens': 0,
            'output_tokens': 0,
            'error': str(e)
        }

# Need to import re at the top if not already
import re

@ai_chat_bp.route('/api/ai-chat/sessions', methods=['GET'])
@require_api_key
@encrypt_response
def get_user_chat_sessions():
    """Get all chat sessions for a user"""
    try:
        user_id = request.args.get('user_id')
        if not user_id:
            return create_error_response("User ID is required", 400)
        
        limit = int(request.args.get('limit', 20))
        offset = int(request.args.get('offset', 0))
        
        with get_db_connection() as conn:
            cursor = conn.cursor(dictionary=True)
            
            # Get sessions with message count and last message preview
            query = """
                SELECT 
                    s.id,
                    s.session_title,
                    s.session_type,
                    s.is_active,
                    s.created_at,
                    s.updated_at,
                    COUNT(m.id) as message_count,
                    (SELECT content FROM ai_chat_messages m2 
                     WHERE m2.session_id = s.id 
                     ORDER BY m2.created_at DESC LIMIT 1) as last_message_preview
                FROM ai_chat_sessions s
                LEFT JOIN ai_chat_messages m ON s.id = m.session_id
                WHERE s.user_id = %s AND s.is_active = 1
                GROUP BY s.id
                ORDER BY s.updated_at DESC
                LIMIT %s OFFSET %s
            """
            
            optimized_execute(cursor, query, (user_id, limit, offset))
            sessions = cursor.fetchall()
            
            # Format datetime objects
            for session in sessions:
                session['created_at'] = session['created_at'].isoformat() if session['created_at'] else None
                session['updated_at'] = session['updated_at'].isoformat() if session['updated_at'] else None
                # Truncate last message for preview
                if session['last_message_preview'] and len(session['last_message_preview']) > 100:
                    session['last_message_preview'] = session['last_message_preview'][:100] + "..."
            
            return format_api_response(sessions)
            
    except mysql.connector.Error as db_error:
        logger.error(f"Database error in get_user_chat_sessions: {str(db_error)}")
        return create_error_response(f"Database error: {str(db_error)}", 500)
    except Exception as e:
        logger.error(f"Error in get_user_chat_sessions: {str(e)}")
        return create_error_response(f"Failed to fetch chat sessions: {str(e)}", 500)

@ai_chat_bp.route('/api/ai-chat/sessions', methods=['POST'])
@require_api_key
@encrypt_response
def create_chat_session():
    """Create a new chat session"""
    try:
        data = request.get_json()
        if not data:
            return create_error_response("No data provided", 400)
        
        user_id = data.get('user_id')
        session_title = data.get('session_title', 'New Chat')
        session_type = data.get('session_type', 'text')
        
        if not user_id:
            return create_error_response("User ID is required", 400)
        
        with get_db_connection() as conn:
            cursor = conn.cursor()
            
            query = """
                INSERT INTO ai_chat_sessions (user_id, session_title, session_type)
                VALUES (%s, %s, %s)
            """
            
            optimized_execute(cursor, query, (user_id, session_title, session_type))
            session_id = cursor.lastrowid
            conn.commit()
            
            logger.info(f"Created new chat session {session_id} for user {user_id}")
            
            return format_api_response({
                "session_id": session_id,
                "session_title": session_title,
                "session_type": session_type,
                "message": "Chat session created successfully"
            })
            
    except mysql.connector.Error as db_error:
        logger.error(f"Database error in create_chat_session: {str(db_error)}")
        return create_error_response(f"Database error: {str(db_error)}", 500)
    except Exception as e:
        logger.error(f"Error in create_chat_session: {str(e)}")
        return create_error_response(f"Failed to create chat session: {str(e)}", 500)

@ai_chat_bp.route('/api/ai-chat/sessions/<int:session_id>/messages', methods=['GET'])
@require_api_key
@encrypt_response
def get_session_messages(session_id):
    """Get all messages in a chat session"""
    try:
        user_id = request.args.get('user_id')
        limit = int(request.args.get('limit', 50))
        offset = int(request.args.get('offset', 0))
        
        if not user_id:
            return create_error_response("User ID is required", 400)
        
        with get_db_connection() as conn:
            cursor = conn.cursor(dictionary=True)
            
            # Verify session belongs to user
            verify_query = """
                SELECT id FROM ai_chat_sessions 
                WHERE id = %s AND user_id = %s AND is_active = 1
            """
            optimized_execute(cursor, verify_query, (session_id, user_id))
            if not cursor.fetchone():
                return create_error_response("Session not found or access denied", 404)
            
            # Get messages
            query = """
                SELECT 
                    id,
                    message_type,
                    content,
                    metadata,
                    parent_message_id,
                    audio_url,
                    image_url,
                    tokens_used,
                    processing_time,
                    created_at
                FROM ai_chat_messages
                WHERE session_id = %s
                ORDER BY created_at ASC
                LIMIT %s OFFSET %s
            """
            
            optimized_execute(cursor, query, (session_id, limit, offset))
            messages = cursor.fetchall()
            
            # Format messages
            for message in messages:
                message['created_at'] = message['created_at'].isoformat() if message['created_at'] else None
                if message['metadata']:
                    try:
                        message['metadata'] = json.loads(message['metadata']) if isinstance(message['metadata'], str) else message['metadata']
                    except json.JSONDecodeError:
                        message['metadata'] = {}
                else:
                    message['metadata'] = {}
            
            return format_api_response(messages)
            
    except mysql.connector.Error as db_error:
        logger.error(f"Database error in get_session_messages: {str(db_error)}")
        return create_error_response(f"Database error: {str(db_error)}", 500)
    except Exception as e:
        logger.error(f"Error in get_session_messages: {str(e)}")
        return create_error_response(f"Failed to fetch messages: {str(e)}", 500)

@ai_chat_bp.route('/api/ai-chat/messages/<int:session_id>', methods=['GET'])
@require_api_key
@encrypt_response
def get_messages_by_session_id(session_id):
    """Get messages for a session - alternative endpoint for frontend compatibility"""
    # This provides the same functionality as get_session_messages but with different URL pattern
    # to match what the frontend is expecting: /api/ai-chat/messages/{sessionId}
    return get_session_messages(session_id)

@ai_chat_bp.route('/api/ai-chat/messages', methods=['POST'])
@require_api_key
@encrypt_response
def save_chat_message():
    """Save a chat message to the database"""
    try:
        data = request.get_json()
        if not data:
            return create_error_response("No data provided", 400)
        
        required_fields = ['session_id', 'message_type', 'content']
        for field in required_fields:
            if field not in data:
                return create_error_response(f"Missing required field: {field}", 400)
        
        session_id = data['session_id']
        message_type = data['message_type']
        content = data['content']
        metadata = data.get('metadata', {})
        parent_message_id = data.get('parent_message_id')
        audio_url = data.get('audio_url')
        image_url = data.get('image_url')
        tokens_used = data.get('tokens_used', 0)
        processing_time = data.get('processing_time', 0)
        user_id = data.get('user_id')
        
        if not user_id:
            return create_error_response("User ID is required", 400)
        
        with get_db_connection() as conn:
            cursor = conn.cursor()
            
            # Verify session belongs to user
            verify_query = """
                SELECT id FROM ai_chat_sessions 
                WHERE id = %s AND user_id = %s AND is_active = 1
            """
            optimized_execute(cursor, verify_query, (session_id, user_id))
            if not cursor.fetchone():
                return create_error_response("Session not found or access denied", 404)
            
            # Insert message
            query = """
                INSERT INTO ai_chat_messages 
                (session_id, message_type, content, metadata, parent_message_id, 
                 audio_url, image_url, tokens_used, processing_time)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
            """
            
            optimized_execute(cursor, query, (
                session_id, message_type, content, json.dumps(metadata),
                parent_message_id, audio_url, image_url, tokens_used, processing_time
            ))
            
            message_id = cursor.lastrowid
            
            # Update session's updated_at timestamp
            update_session_query = """
                UPDATE ai_chat_sessions 
                SET updated_at = CURRENT_TIMESTAMP
                WHERE id = %s
            """
            optimized_execute(cursor, update_session_query, (session_id,))
            
            # Insert into search index
            search_index_query = """
                INSERT INTO ai_chat_search_index 
                (session_id, message_id, user_id, searchable_content)
                VALUES (%s, %s, %s, %s)
            """
            optimized_execute(cursor, search_index_query, (session_id, message_id, user_id, content))
            
            conn.commit()
            
            logger.info(f"Saved message {message_id} to session {session_id}")
            
            return format_api_response({
                "message_id": message_id,
                "session_id": session_id,
                "message": "Message saved successfully"
            })
            
    except mysql.connector.Error as db_error:
        logger.error(f"Database error in save_chat_message: {str(db_error)}")
        return create_error_response(f"Database error: {str(db_error)}", 500)
    except Exception as e:
        logger.error(f"Error in save_chat_message: {str(e)}")
        return create_error_response(f"Failed to save message: {str(e)}", 500)

@ai_chat_bp.route('/api/ai-chat/sessions/<int:session_id>', methods=['PUT'])
@require_api_key
@encrypt_response
def update_chat_session(session_id):
    """Update chat session (e.g., change title)"""
    try:
        data = request.get_json()
        user_id = data.get('user_id')
        
        if not user_id:
            return create_error_response("User ID is required", 400)
        
        with get_db_connection() as conn:
            cursor = conn.cursor()
            
            # Verify session belongs to user
            verify_query = """
                SELECT id FROM ai_chat_sessions 
                WHERE id = %s AND user_id = %s AND is_active = 1
            """
            optimized_execute(cursor, verify_query, (session_id, user_id))
            if not cursor.fetchone():
                return create_error_response("Session not found or access denied", 404)
            
            # Build update query dynamically
            update_fields = []
            update_values = []
            
            if 'session_title' in data:
                update_fields.append("session_title = %s")
                update_values.append(data['session_title'])
            
            if 'session_type' in data:
                update_fields.append("session_type = %s")
                update_values.append(data['session_type'])
            
            if not update_fields:
                return create_error_response("No fields to update", 400)
            
            update_values.append(session_id)
            
            query = f"""
                UPDATE ai_chat_sessions 
                SET {', '.join(update_fields)}, updated_at = CURRENT_TIMESTAMP
                WHERE id = %s
            """
            
            optimized_execute(cursor, query, update_values)
            conn.commit()
            
            return format_api_response({"message": "Session updated successfully"})
            
    except mysql.connector.Error as db_error:
        logger.error(f"Database error in update_chat_session: {str(db_error)}")
        return create_error_response(f"Database error: {str(db_error)}", 500)
    except Exception as e:
        logger.error(f"Error in update_chat_session: {str(e)}")
        return create_error_response(f"Failed to update session: {str(e)}", 500)

@ai_chat_bp.route('/api/ai-chat/sessions/<int:session_id>', methods=['DELETE'])
@require_api_key
@encrypt_response
def delete_chat_session(session_id):
    """Delete (deactivate) a chat session"""
    try:
        user_id = request.args.get('user_id')
        
        if not user_id:
            return create_error_response("User ID is required", 400)
        
        with get_db_connection() as conn:
            cursor = conn.cursor()
            
            # Verify session belongs to user
            verify_query = """
                SELECT id FROM ai_chat_sessions 
                WHERE id = %s AND user_id = %s AND is_active = 1
            """
            optimized_execute(cursor, verify_query, (session_id, user_id))
            if not cursor.fetchone():
                return create_error_response("Session not found or access denied", 404)
            
            # Soft delete by setting is_active = 0
            query = """
                UPDATE ai_chat_sessions 
                SET is_active = 0, updated_at = CURRENT_TIMESTAMP
                WHERE id = %s
            """
            
            optimized_execute(cursor, query, (session_id,))
            conn.commit()
            
            logger.info(f"Deleted session {session_id} for user {user_id}")
            
            return format_api_response({"message": "Session deleted successfully"})
            
    except mysql.connector.Error as db_error:
        logger.error(f"Database error in delete_chat_session: {str(db_error)}")
        return create_error_response(f"Database error: {str(db_error)}", 500)
    except Exception as e:
        logger.error(f"Error in delete_chat_session: {str(e)}")
        return create_error_response(f"Failed to delete session: {str(e)}", 500)

@ai_chat_bp.route('/api/ai-chat/search', methods=['GET'])
@require_api_key
@encrypt_response
def search_chat_history():
    """Search through user's chat history"""
    try:
        user_id = request.args.get('user_id')
        query_text = request.args.get('q', '').strip()
        limit = int(request.args.get('limit', 20))
        offset = int(request.args.get('offset', 0))
        
        if not user_id:
            return create_error_response("User ID is required", 400)
        
        if not query_text:
            return create_error_response("Search query is required", 400)
        
        with get_db_connection() as conn:
            cursor = conn.cursor(dictionary=True)
            
            # Full-text search with relevance scoring
            search_query = """
                SELECT 
                    i.session_id,
                    i.message_id,
                    s.session_title,
                    m.message_type,
                    m.content,
                    m.created_at,
                    MATCH(i.searchable_content) AGAINST(%s IN NATURAL LANGUAGE MODE) as relevance
                FROM ai_chat_search_index i
                JOIN ai_chat_sessions s ON i.session_id = s.id
                JOIN ai_chat_messages m ON i.message_id = m.id
                WHERE i.user_id = %s 
                AND s.is_active = 1
                AND MATCH(i.searchable_content) AGAINST(%s IN NATURAL LANGUAGE MODE)
                ORDER BY relevance DESC, m.created_at DESC
                LIMIT %s OFFSET %s
            """
            
            optimized_execute(cursor, search_query, (query_text, user_id, query_text, limit, offset))
            results = cursor.fetchall()
            
            # Format results
            for result in results:
                result['created_at'] = result['created_at'].isoformat() if result['created_at'] else None
                # Highlight search terms in content (simple approach)
                if query_text.lower() in result['content'].lower():
                    content = result['content']
                    start_idx = content.lower().find(query_text.lower())
                    if start_idx != -1:
                        # Get context around the match
                        context_start = max(0, start_idx - 50)
                        context_end = min(len(content), start_idx + len(query_text) + 50)
                        result['content_snippet'] = content[context_start:context_end]
                        if context_start > 0:
                            result['content_snippet'] = "..." + result['content_snippet']
                        if context_end < len(content):
                            result['content_snippet'] = result['content_snippet'] + "..."
                    else:
                        result['content_snippet'] = result['content'][:100] + "..."
                else:
                    result['content_snippet'] = result['content'][:100] + "..."
            
            return format_api_response(results)
            
    except mysql.connector.Error as db_error:
        logger.error(f"Database error in search_chat_history: {str(db_error)}")
        return create_error_response(f"Database error: {str(db_error)}", 500)
    except Exception as e:
        logger.error(f"Error in search_chat_history: {str(e)}")
        return create_error_response(f"Failed to search chat history: {str(e)}", 500)

# Enhanced AI Chat Endpoints with Multiple Modes

@ai_chat_bp.route('/api/ai-chat/sessions/create', methods=['POST'])
@require_api_key
@encrypt_response
def create_enhanced_chat_session():
    """Create a new chat session with mode support"""
    try:
        data = request.get_json()
        if not data:
            return create_error_response("No data provided", 400)
        
        user_id = data.get('user_id')
        session_title = data.get('session_title', 'New Chat')
        session_type = data.get('session_type', 'text')
        mode = data.get('mode', 'general')
        
        if not user_id:
            return create_error_response("User ID is required", 400)
            
        # Validate mode
        if mode not in CHAT_MODES:
            mode = 'general'
        
        with get_db_connection() as conn:
            cursor = conn.cursor(dictionary=True)
            
            # Create session
            insert_query = """
                INSERT INTO ai_chat_sessions (user_id, session_title, session_type, is_active, created_at, updated_at)
                VALUES (%s, %s, %s, 1, NOW(), NOW())
            """
            optimized_execute(cursor, insert_query, (user_id, session_title, session_type))
            session_id = cursor.lastrowid
            
            # Add system message with mode prompt if not general mode
            if mode != 'general':
                mode_config = get_chat_mode_config(mode)
                system_message_query = """
                    INSERT INTO ai_chat_messages (session_id, message_type, content, metadata, created_at)
                    VALUES (%s, 'system', %s, %s, NOW())
                """
                metadata = json.dumps({
                    'mode': mode,
                    'mode_name': mode_config['name'],
                    'is_mode_init': True
                })
                optimized_execute(cursor, system_message_query, (session_id, mode_config['prompt'], metadata))
            
            conn.commit()
            
            # Return session details
            session_data = {
                'id': session_id,
                'session_title': session_title,
                'session_type': session_type,
                'mode': mode,
                'mode_name': get_chat_mode_config(mode)['name'],
                'is_active': True,
                'created_at': datetime.now().isoformat(),
                'message_count': 1 if mode != 'general' else 0
            }
            
            return format_api_response({'session': session_data})
            
    except mysql.connector.Error as db_error:
        logger.error(f"Database error in create_enhanced_chat_session: {str(db_error)}")
        return create_error_response(f"Database error: {str(db_error)}", 500)
    except Exception as e:
        logger.error(f"Error in create_enhanced_chat_session: {str(e)}")
        return create_error_response(f"Failed to create chat session: {str(e)}", 500)

@ai_chat_bp.route('/api/ai-chat/send', methods=['POST'])
@require_api_key
@encrypt_response
def send_chat_message():
    """Send a message and get AI response with optional image support + Token Billing"""
    try:
        data = request.get_json()
        if not data:
            return create_error_response("No data provided", 400)
        
        user_id = data.get('user_id')
        session_id = data.get('session_id')
        message = data.get('message', '').strip()
        message_type = data.get('message_type', 'user')
        mode = data.get('mode', 'general')
        system_prompt = data.get('system_prompt', '')
        
        # Image support: images array in format [[data_uri, filename], ...]
        images = data.get('images', [])
        model_override = data.get('model')
        provider_override = data.get('provider')
        
        if not user_id or not session_id:
            return create_error_response("User ID and Session ID are required", 400)
        
        if not message:
            return create_error_response("Message is required", 400)
        
        # ========== TOKEN & SUBSCRIPTION CHECKS ==========
        # Estimate tokens needed (rough estimate: 1 token per 4 characters)
        estimated_tokens = max(10, len(message) / 4)
        
        # Check if user has any positive token balance (actual deduction happens after response)
        balance_check = TokenManager.check_balance(user_id, 1)  # Check for at least 1 token
        if balance_check['balance'] <= 0:
            logger.warning(f"User {user_id} has no tokens remaining: {balance_check['balance']}")
            # Return dict tuple - encrypt_response decorator will handle jsonify
            return {
                "status": "error",
                "error_code": "INSUFFICIENT_TOKENS",
                "message": "You have no AI tokens remaining",
                "current_balance": balance_check['balance'],
                "required_tokens": 1,
                "shortfall": 1,
                "upgrade_url": "/pricing"
            }, 402  # 402 Payment Required
        
        # Get the model that will be used
        mode_cfg = get_chat_mode_config(mode)
        actual_model = model_override or mode_cfg.get('model', 'geminisearch')
        
        # Check if user has subscription access to this model
        access_check = SubscriptionManager.check_subscription_access(
            user_id=user_id,
            required_benefit={'type': 'ai_tokens', 'model': actual_model}
        )
        
        if not access_check['has_access']:
            logger.warning(f"User {user_id} does not have access to model {actual_model}")
            # Return dict tuple - encrypt_response decorator will handle jsonify
            return {
                "status": "error",
                "error_code": "SUBSCRIPTION_REQUIRED",
                "message": access_check.get('reason', 'Model not available in your plan'),
                "required_model": actual_model,
                "suggested_plan": access_check.get('suggested_plan', 'ai-starter'),
                "upgrade_url": f"/pricing?highlight={access_check.get('suggested_plan', 'ai-starter')}"
            }, 403  # 403 Forbidden
        
        logger.info(f"User {user_id} authorized for {actual_model}, balance: {balance_check['balance']} tokens")
        # ========== END TOKEN & SUBSCRIPTION CHECKS ==========
        
        with get_db_connection() as conn:
            cursor = conn.cursor(dictionary=True)
            
            # Verify session belongs to user
            verify_query = "SELECT id, session_title FROM ai_chat_sessions WHERE id = %s AND user_id = %s"
            optimized_execute(cursor, verify_query, (session_id, user_id))
            session = cursor.fetchone()
            
            if not session:
                return create_error_response("Session not found or access denied", 404)
            
            # Get conversation history for context
            history_query = """
                SELECT message_type, content, created_at 
                FROM ai_chat_messages 
                WHERE session_id = %s AND message_type IN ('user', 'assistant')
                ORDER BY created_at ASC
            """
            optimized_execute(cursor, history_query, (session_id,))
            conversation_history = cursor.fetchall()
            
            # Save user message (if not system message)
            user_message_id = None
            if message_type == 'user':
                user_message_query = """
                    INSERT INTO ai_chat_messages (session_id, message_type, content, metadata, created_at)
                    VALUES (%s, %s, %s, %s, NOW())
                """
                user_metadata = json.dumps({
                    'mode': mode,
                    'provider': provider_override or 'PollinationsAI',
                    'model': model_override or 'auto',
                    'tokens_used': len(message.split()),
                    'processing_time': 0,
                    'has_images': len(images) > 0 if images else False,
                    'image_count': len(images) if images else 0
                })
                optimized_execute(cursor, user_message_query, (session_id, 'user', message, user_metadata))
                user_message_id = cursor.lastrowid
                
                # Add to search index
                search_index_query = """
                    INSERT INTO ai_chat_search_index (session_id, message_id, user_id, searchable_content, created_at)
                    VALUES (%s, %s, %s, %s, NOW())
                """
                optimized_execute(cursor, search_index_query, (session_id, user_message_id, user_id, message))
            
            # Generate AI response (if user message)
            ai_response = None
            ai_message_id = None
            if message_type == 'user':
                try:
                    # Generate AI response with image support
                    ai_response = generate_ai_response(
                        message, 
                        mode, 
                        conversation_history,
                        images=images if images else None,
                        model_override=model_override,
                        provider_override=provider_override,
                        user_id=user_id  # Pass user_id to use their preferred model
                    )
                    
                    # ========== TOKEN DEDUCTION ==========
                    logger.info(f"[TOKEN DEBUG] Starting token calculation for user {user_id}")
                    logger.info(f"[TOKEN DEBUG] Message: '{message[:50]}...'")
                    logger.info(f"[TOKEN DEBUG] Response: '{ai_response[:50] if ai_response else 'None'}...'")
                    
                    # Calculate actual tokens used (word-based estimation)
                    input_tokens = len(message.split())
                    output_tokens = len(ai_response.split()) if ai_response else 0
                    total_tokens_used = input_tokens + output_tokens
                    
                    logger.info(f"[TOKEN DEBUG] Input tokens (words): {input_tokens}")
                    logger.info(f"[TOKEN DEBUG] Output tokens (words): {output_tokens}")
                    logger.info(f"[TOKEN DEBUG] Total tokens: {total_tokens_used}")
                    
                    # Calculate cost based on model pricing
                    model_name_for_cost = model_override or mode_cfg.get('model', 'geminisearch')
                    logger.info(f"[TOKEN DEBUG] Model for cost calculation: {model_name_for_cost}")
                    
                    token_cost = TokenManager.calculate_token_cost(
                        input_tokens=input_tokens,
                        output_tokens=output_tokens,
                        model_name=model_name_for_cost
                    )
                    
                    logger.info(f"[TOKEN DEBUG] Calculated token cost: {token_cost}")
                    
                    # If token cost is 0 (free model or error), discount input tokens (10% weight)
                    tokens_to_deduct = token_cost if token_cost > 0 else output_tokens + (input_tokens * 0.1)
                    logger.info(f"[TOKEN DEBUG] Tokens to deduct: {tokens_to_deduct}")
                    
                    # Deduct tokens from user balance
                    logger.info(f"[TOKEN DEBUG] Calling TokenManager.deduct_tokens for user {user_id}")
                    deduction_result = TokenManager.deduct_tokens(
                        user_id=user_id,
                        amount=tokens_to_deduct,
                        endpoint='/api/ai-chat/send',
                        ai_model=model_override or mode_cfg.get('model', 'geminisearch'),
                        provider=provider_override or mode_cfg.get('provider', 'PollinationsAI'),
                        session_id=session_id,
                        message_id=None,  # Will update after message saved
                        request_tokens=input_tokens,
                        response_tokens=output_tokens,
                        metadata={
                            'mode': mode,
                            'has_images': len(images) > 0 if images else False,
                            'image_count': len(images) if images else 0
                        }
                    )
                    
                    logger.info(f"[TOKEN DEBUG] Deduction result: {deduction_result}")
                    
                    if not deduction_result['success']:
                        logger.error(f"[TOKEN ERROR] Token deduction FAILED for user {user_id}: {deduction_result.get('message')}")
                        logger.error(f"[TOKEN ERROR] Full deduction result: {deduction_result}")
                        # Continue anyway, log the error
                    else:
                        logger.info(f"[TOKEN SUCCESS] ✅ Deducted {tokens_to_deduct:.2f} tokens from user {user_id}")
                        logger.info(f"[TOKEN SUCCESS] New balance: {deduction_result['balance_after']:.2f}")
                        logger.info(f"[TOKEN SUCCESS] Previous balance: {deduction_result.get('balance_before', 'unknown')}")
                    # ========== END TOKEN DEDUCTION ==========
                    
                    # Save AI response
                    ai_message_query = """
                        INSERT INTO ai_chat_messages (session_id, message_type, content, metadata, parent_message_id, tokens_used, created_at)
                        VALUES (%s, %s, %s, %s, %s, %s, NOW())
                    """
                    mode_cfg = get_chat_mode_config(mode)
                    ai_metadata = json.dumps({
                        'mode': mode,
                        'mode_name': mode_cfg.get('name', 'General Assistant'),
                        'provider': provider_override or mode_cfg.get('provider', 'PollinationsAI'),
                        'model_used': model_override or mode_cfg.get('model', 'geminisearch'),
                        'tokens_used': tokens_to_deduct,
                        'input_tokens': input_tokens,
                        'output_tokens': output_tokens,
                        'processing_time': 1.0,  # Placeholder
                        'was_vision_request': len(images) > 0 if images else False,
                        'image_count': len(images) if images else 0,
                        'token_cost': token_cost,
                        'transaction_id': deduction_result.get('transaction_id') if deduction_result['success'] else None
                    })
                    optimized_execute(cursor, ai_message_query, (session_id, 'assistant', ai_response, ai_metadata, user_message_id, tokens_to_deduct))
                    ai_message_id = cursor.lastrowid
                    
                    # Add AI response to search index
                    optimized_execute(cursor, search_index_query, (session_id, ai_message_id, user_id, ai_response))
                    
                    # Update session title if first message
                    if len(conversation_history) == 0:
                        title = message[:50] + "..." if len(message) > 50 else message
                        update_title_query = "UPDATE ai_chat_sessions SET session_title = %s, updated_at = NOW() WHERE id = %s"
                        optimized_execute(cursor, update_title_query, (title, session_id))
                    else:
                        # Just update the timestamp
                        update_timestamp_query = "UPDATE ai_chat_sessions SET updated_at = NOW() WHERE id = %s"
                        optimized_execute(cursor, update_timestamp_query, (session_id,))
                    
                except Exception as ai_error:
                    logger.error(f"AI generation error: {str(ai_error)}")
                    ai_response = "I apologize, but I'm having trouble generating a response right now. Please try again."
                    
                    # Save fallback response
                    ai_message_query = """
                        INSERT INTO ai_chat_messages (session_id, message_type, content, metadata, parent_message_id, created_at)
                        VALUES (%s, %s, %s, %s, %s, NOW())
                    """
                    mode_cfg_fallback = get_chat_mode_config(mode)
                    ai_metadata = json.dumps({
                        'mode': mode,
                        'provider': mode_cfg_fallback.get('provider', 'PollinationsAI'),
                        'model_used': mode_cfg_fallback.get('model', 'geminisearch'),
                        'error': str(ai_error),
                        'is_fallback': True
                    })
                    optimized_execute(cursor, ai_message_query, (session_id, 'assistant', ai_response, ai_metadata, user_message_id))
                    ai_message_id = cursor.lastrowid
            
            # Log analytics
            try:
                analytics_query = """
                    INSERT INTO ai_chat_analytics (user_id, session_id, interaction_type, language_used, success, mode_used, created_at)
                    VALUES (%s, %s, %s, %s, %s, %s, NOW())
                """
                optimized_execute(cursor, analytics_query, (user_id, session_id, 'text_query', 'en-US', ai_response is not None, mode))
            except Exception as analytics_error:
                logger.warning(f"Analytics logging failed: {str(analytics_error)}")
            
            conn.commit()
            
            # Prepare response
            response_data = {
                'session_id': session_id,
                'user_message_id': user_message_id,
                'ai_message_id': ai_message_id,
                'ai_response': ai_response,
                'mode': mode,
                'mode_name': get_chat_mode_config(mode).get('name', 'General Assistant'),
                'timestamp': datetime.now().isoformat()
            }
            
            return format_api_response(response_data)
            
    except mysql.connector.Error as db_error:
        logger.error(f"Database error in send_chat_message: {str(db_error)}")
        return create_error_response(f"Database error: {str(db_error)}", 500)
    except Exception as e:
        logger.error(f"Error in send_chat_message: {str(e)}")
        return create_error_response(f"Failed to send message: {str(e)}", 500)


@ai_chat_bp.route('/api/ai-chat/send-with-tools', methods=['POST'])
@require_api_key
@encrypt_response
def send_chat_message_with_tools():
    """Send a message and get AI response with tool execution support.
    
    This endpoint returns structured response blocks that can include:
    - text: Regular markdown text
    - tool_call: Tool execution with results (cards with navigation)
    
    The AI can use tools to search the database and return interactive content.
    """
    try:
        data = request.get_json()
        if not data:
            return create_error_response("No data provided", 400)
        
        user_id = data.get('user_id')
        session_id = data.get('session_id')
        message = data.get('message', '').strip()
        mode = data.get('mode', 'general')
        enable_tools = data.get('enable_tools', True)  # Allow disabling tools
        language = data.get('language', 'en-US')
        
        # Image support
        images = data.get('images', [])
        model_override = data.get('model')
        provider_override = data.get('provider')
        
        if not user_id:
            return create_error_response("User ID is required", 400)
        
        if not message:
            return create_error_response("Message is required", 400)
        
        # ========== TOKEN & SUBSCRIPTION CHECKS ==========
        balance_check = TokenManager.check_balance(user_id, 1)
        if balance_check['balance'] <= 0:
            return {
                "status": "error",
                "error_code": "INSUFFICIENT_TOKENS",
                "message": "You have no AI tokens remaining",
                "current_balance": balance_check['balance']
            }, 402
        
        mode_cfg = get_chat_mode_config(mode)
        actual_model = model_override or mode_cfg.get('model', 'geminisearch')
        
        access_check = SubscriptionManager.check_subscription_access(
            user_id=user_id,
            required_benefit={'type': 'ai_tokens', 'model': actual_model}
        )
        
        if not access_check['has_access']:
            return {
                "status": "error",
                "error_code": "SUBSCRIPTION_REQUIRED",
                "message": access_check.get('reason', 'Model not available in your plan')
            }, 403
        
        with get_db_connection() as conn:
            cursor = conn.cursor(dictionary=True)
            
            # Auto-create session if not provided
            if not session_id:
                # Try to find an existing active session for this user, or create one
                find_session_query = """
                    SELECT id FROM ai_chat_sessions 
                    WHERE user_id = %s AND is_active = 1 
                    ORDER BY updated_at DESC LIMIT 1
                """
                optimized_execute(cursor, find_session_query, (user_id,))
                existing_session = cursor.fetchone()
                
                if existing_session:
                    session_id = existing_session['id']
                    logger.info(f"Using existing session {session_id} for user {user_id}")
                else:
                    # Create a new session
                    session_title = message[:50] + "..." if len(message) > 50 else message
                    create_session_query = """
                        INSERT INTO ai_chat_sessions (user_id, session_title, session_type, is_active, created_at, updated_at)
                        VALUES (%s, %s, %s, 1, NOW(), NOW())
                    """
                    optimized_execute(cursor, create_session_query, (user_id, session_title, 'text'))
                    session_id = cursor.lastrowid
                    conn.commit()
                    logger.info(f"Created new session {session_id} for user {user_id}")
            else:
                # Verify provided session belongs to user
                verify_query = "SELECT id, session_title FROM ai_chat_sessions WHERE id = %s AND user_id = %s"
                optimized_execute(cursor, verify_query, (session_id, user_id))
                session = cursor.fetchone()
                
                if not session:
                    return create_error_response("Session not found or access denied", 404)
            
            # Get conversation history
            history_query = """
                SELECT message_type, content, metadata, created_at 
                FROM ai_chat_messages 
                WHERE session_id = %s AND message_type IN ('user', 'assistant')
                ORDER BY created_at ASC
            """
            optimized_execute(cursor, history_query, (session_id,))
            conversation_history = cursor.fetchall()
            
            # Save user message
            user_message_query = """
                INSERT INTO ai_chat_messages (session_id, message_type, content, metadata, created_at)
                VALUES (%s, %s, %s, %s, NOW())
            """
            user_metadata = json.dumps({
                'mode': mode,
                'has_images': len(images) > 0,
                'tools_enabled': enable_tools
            })
            optimized_execute(cursor, user_message_query, (session_id, 'user', message, user_metadata))
            user_message_id = cursor.lastrowid
            
            # Add to search index
            search_index_query = """
                INSERT INTO ai_chat_search_index (session_id, message_id, user_id, searchable_content, created_at)
                VALUES (%s, %s, %s, %s, NOW())
            """
            optimized_execute(cursor, search_index_query, (session_id, user_message_id, user_id, message))
            
            # Generate AI response with tools
            ai_result = generate_ai_response_with_tools(
                user_message=message,
                mode=mode,
                conversation_history=conversation_history,
                images=images if images else None,
                model_override=model_override,
                provider_override=provider_override,
                user_id=user_id,
                language=language,
                enable_tools=enable_tools
            )
            
            # Extract data from tool response
            response_blocks = ai_result.get('response_blocks', [])
            raw_response = ai_result.get('raw_response', '')
            tools_used = ai_result.get('tools_used', [])
            input_tokens = ai_result.get('input_tokens', 0)
            output_tokens = ai_result.get('output_tokens', 0)
            
            # Token deduction
            token_cost = TokenManager.calculate_token_cost(
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                model_name=ai_result.get('model_used', actual_model)
            )
            tokens_to_deduct = token_cost if token_cost > 0 else output_tokens + (input_tokens * 0.1)
            
            deduction_result = TokenManager.deduct_tokens(
                user_id=user_id,
                amount=tokens_to_deduct,
                endpoint='/api/ai-chat/send-with-tools',
                ai_model=ai_result.get('model_used', actual_model),
                session_id=session_id,
                request_tokens=input_tokens,
                response_tokens=output_tokens,
                metadata={'tools_used': tools_used, 'mode': mode}
            )
            
            # Save AI response (store both raw and blocks)
            ai_message_query = """
                INSERT INTO ai_chat_messages (session_id, message_type, content, metadata, parent_message_id, tokens_used, created_at)
                VALUES (%s, %s, %s, %s, %s, %s, NOW())
            """
            ai_metadata = json.dumps({
                'mode': mode,
                'model_used': ai_result.get('model_used'),
                'tools_used': tools_used,
                'response_blocks': response_blocks,  # Store structured blocks
                'tokens_used': tokens_to_deduct,
                'transaction_id': deduction_result.get('transaction_id') if deduction_result.get('success') else None
            })
            optimized_execute(cursor, ai_message_query, (
                session_id, 'assistant', raw_response, ai_metadata, user_message_id, tokens_to_deduct
            ))
            ai_message_id = cursor.lastrowid
            
            # Add AI response to search index
            optimized_execute(cursor, search_index_query, (session_id, ai_message_id, user_id, raw_response))
            
            # Update session title if first message
            if len(conversation_history) == 0:
                title = message[:50] + "..." if len(message) > 50 else message
                update_title_query = "UPDATE ai_chat_sessions SET session_title = %s, updated_at = NOW() WHERE id = %s"
                optimized_execute(cursor, update_title_query, (title, session_id))
            else:
                update_timestamp_query = "UPDATE ai_chat_sessions SET updated_at = NOW() WHERE id = %s"
                optimized_execute(cursor, update_timestamp_query, (session_id,))
            
            # Log analytics
            try:
                analytics_query = """
                    INSERT INTO ai_chat_analytics (user_id, session_id, interaction_type, language_used, success, mode_used, created_at)
                    VALUES (%s, %s, %s, %s, %s, %s, NOW())
                """
                optimized_execute(cursor, analytics_query, (
                    user_id, session_id, 
                    'tool_chat' if tools_used else 'text_query',
                    language, True, mode
                ))
            except Exception as analytics_error:
                logger.warning(f"Analytics logging failed: {analytics_error}")
            
            conn.commit()
            
            # Return structured response (plain dict for @encrypt_response decorator)
            return {
                'success': True,
                'session_id': session_id,
                'user_message_id': user_message_id,
                'ai_message_id': ai_message_id,
                'response_blocks': response_blocks,  # Main structured response
                'ai_response': raw_response,  # Fallback raw text
                'tools_used': tools_used,
                'mode': mode,
                'mode_name': mode_cfg.get('name', 'General Assistant'),
                'tokens_used': tokens_to_deduct,
                'timestamp': datetime.now().isoformat()
            }
            
    except mysql.connector.Error as db_error:
        logger.error(f"Database error in send_chat_message_with_tools: {db_error}")
        return create_error_response(f"Database error: {str(db_error)}", 500)
    except Exception as e:
        logger.error(f"Error in send_chat_message_with_tools: {e}")
        return create_error_response(f"Failed to send message: {str(e)}", 500)

@ai_chat_bp.route('/api/ai-chat/modes', methods=['GET'])
@require_api_key
@encrypt_response
def get_chat_modes():
    """Get available chat modes"""
    try:
        modes_list = []
        for mode_id, config in CHAT_MODES.items():
            modes_list.append({
                'id': mode_id,
                'name': config['name'],
                'description': config['prompt'][:100] + "..." if len(config['prompt']) > 100 else config['prompt']
            })
        
        return format_api_response({'modes': modes_list})
    
    except Exception as e:
        logger.error(f"Error in get_chat_modes: {str(e)}")
        return create_error_response(f"Failed to get chat modes: {str(e)}", 500)

@ai_chat_bp.route('/api/ai-chat/messages/<int:message_id>/favorite', methods=['POST'])
@require_api_key  
@encrypt_response
def add_message_to_favorites(message_id):
    """Add a message to favorites"""
    try:
        data = request.get_json()
        user_id = data.get('user_id') if data else request.args.get('user_id')
        
        if not user_id:
            return create_error_response("User ID is required", 400)
        
        with get_db_connection() as conn:
            cursor = conn.cursor(dictionary=True)
            
            # Get message and session info
            message_query = """
                SELECT m.id, m.session_id, s.user_id as session_user_id
                FROM ai_chat_messages m
                JOIN ai_chat_sessions s ON m.session_id = s.id
                WHERE m.id = %s
            """
            optimized_execute(cursor, message_query, (message_id,))
            message_info = cursor.fetchone()
            
            if not message_info:
                return create_error_response("Message not found", 404)
            
            if str(message_info['session_user_id']) != str(user_id):
                return create_error_response("Access denied", 403)
            
            # Add to favorites (ignore if already exists)
            favorite_query = """
                INSERT IGNORE INTO ai_chat_favorites (user_id, session_id, message_id, created_at)
                VALUES (%s, %s, %s, NOW())
            """
            optimized_execute(cursor, favorite_query, (user_id, message_info['session_id'], message_id))
            conn.commit()
            
            return format_api_response({'message': 'Added to favorites successfully'})
            
    except mysql.connector.Error as db_error:
        logger.error(f"Database error in add_message_to_favorites: {str(db_error)}")
        return create_error_response(f"Database error: {str(db_error)}", 500)
    except Exception as e:
        logger.error(f"Error in add_message_to_favorites: {str(e)}")
        return create_error_response(f"Failed to add to favorites: {str(e)}", 500)

@ai_chat_bp.route('/api/ai-chat/messages/<int:message_id>/favorite', methods=['DELETE'])
@require_api_key
@encrypt_response  
def remove_message_from_favorites(message_id):
    """Remove a message from favorites"""
    try:
        user_id = request.args.get('user_id')
        if not user_id:
            return create_error_response("User ID is required", 400)
        
        with get_db_connection() as conn:
            cursor = conn.cursor(dictionary=True)
            
            # Remove from favorites
            remove_query = """
                DELETE f FROM ai_chat_favorites f
                JOIN ai_chat_sessions s ON f.session_id = s.id
                WHERE f.message_id = %s AND s.user_id = %s
            """
            optimized_execute(cursor, remove_query, (message_id, user_id))
            conn.commit()
            
            return format_api_response({'message': 'Removed from favorites successfully'})
            
    except mysql.connector.Error as db_error:
        logger.error(f"Database error in remove_message_from_favorites: {str(db_error)}")
        return create_error_response(f"Database error: {str(db_error)}", 500)
    except Exception as e:
        logger.error(f"Error in remove_message_from_favorites: {str(e)}")
        return create_error_response(f"Failed to remove from favorites: {str(e)}", 500)

# Alternative favorites endpoints for frontend compatibility
@ai_chat_bp.route('/api/ai-chat/favorites', methods=['POST'])
@require_api_key
@encrypt_response
def add_to_favorites_alt():
    """Add a message to favorites - alternative endpoint for frontend compatibility"""
    try:
        data = request.get_json()
        if not data:
            return create_error_response("No data provided", 400)
        
        message_id = data.get('message_id')
        if not message_id:
            return create_error_response("Message ID is required", 400)
        
        # Call the existing function with the message_id
        return add_message_to_favorites(message_id)
        
    except Exception as e:
        logger.error(f"Error in add_to_favorites_alt: {str(e)}")
        return create_error_response(f"Failed to add to favorites: {str(e)}", 500)

@ai_chat_bp.route('/api/ai-chat/favorites/<int:message_id>', methods=['DELETE'])
@require_api_key
@encrypt_response
def remove_from_favorites_alt(message_id):
    """Remove a message from favorites - alternative endpoint for frontend compatibility"""
    # Call the existing function with the message_id
    return remove_message_from_favorites(message_id)

@ai_chat_bp.route('/api/ai-chat/analytics', methods=['POST'])
@require_api_key
@encrypt_response
def log_chat_analytics():
    """Log analytics data for chat interactions"""
    try:
        data = request.get_json()
        if not data:
            return create_error_response("No data provided", 400)
        
        required_fields = ['user_id', 'session_id', 'interaction_type']
        for field in required_fields:
            if field not in data:
                return create_error_response(f"Missing required field: {field}", 400)
        
        with get_db_connection() as conn:
            cursor = conn.cursor()
            
            query = """
                INSERT INTO ai_chat_analytics 
                (user_id, session_id, interaction_type, language_used, voice_used, 
                 success, error_message)
                VALUES (%s, %s, %s, %s, %s, %s, %s)
            """
            
            optimized_execute(cursor, query, (
                data['user_id'],
                data['session_id'],
                data['interaction_type'],
                data.get('language_used', 'en-US'),
                data.get('voice_used'),
                data.get('success', True),
                data.get('error_message')
            ))
            
            conn.commit()
            
            return format_api_response({"message": "Analytics logged successfully"})
            
    except mysql.connector.Error as db_error:
        logger.error(f"Database error in log_chat_analytics: {str(db_error)}")
        return create_error_response(f"Database error: {str(db_error)}", 500)
    except Exception as e:
        logger.error(f"Error in log_chat_analytics: {str(e)}")
        return create_error_response(f"Failed to log analytics: {str(e)}", 500)