"""
AI Study Tools Routes
Dedicated endpoints for AI-powered flashcard generation, summaries, quizzes, and other study aids.
Uses the AI generation system with specialized prompts for each study tool type.
"""

import json
import logging
import re
from flask import Blueprint, request, jsonify
from utils import require_api_key, encrypt_response
from database import get_cached_connection

# Import AI provider
from services.ai_provider import ai_generate

# Check AI availability (g4f or groq)
try:
    from services.ai_config import is_groq_enabled
    try:
        import g4f
        ai_available = True
    except ImportError:
        ai_available = False
    ai_available = ai_available or is_groq_enabled()
except ImportError:
    ai_available = False
    
# Import AI model configuration
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from aimodel_selector.models_config import get_default_model, get_model_config

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

# Create blueprint
ai_study_tools_bp = Blueprint('ai_study_tools', __name__)


# ==================== FLASHCARD GENERATION ====================

FLASHCARD_SYSTEM_PROMPT = """You are "Sathi AI", the friendly AI learning companion of Ashlya Academy, specializing in creating effective flashcards for learning.

Create flashcards from the provided content following these rules:
1. Each flashcard should have a clear, concise question and answer
2. Focus on key concepts, definitions, formulas, and important facts
3. Use simple language that aids memorization
4. Create 5-8 high-quality flashcards
5. Assign difficulty: 'easy', 'medium', or 'hard' to each card

**IMPORTANT: You MUST respond with ONLY a valid JSON array. No explanatory text before or after.**

Response format (JSON array only):
[
  {"question": "What is X?", "answer": "X is...", "difficulty": "easy"},
  {"question": "Explain Y", "answer": "Y works by...", "difficulty": "medium"}
]
"""

@ai_study_tools_bp.route('/api/ai-study/generate-flashcards', methods=['POST'])
@require_api_key
@encrypt_response
def generate_flashcards():
    """Generate AI-powered flashcards from content"""
    try:
        data = request.get_json()
        logger.info(f"Generate flashcards request data: {data}")
        if not data:
            logger.error("No JSON data received")
            return jsonify({"success": False, "message": "No data provided"}), 400
        
        content = data.get('content', '')
        topic_id = data.get('topic_id')
        user_id = data.get('user_id')
        count = data.get('count', 8)  # Number of flashcards to generate
        
        if not content and topic_id:
            # Try to fetch content from database if topic_id is provided
            try:
                # Import here to avoid circular dependencies
                from utils import fetch_datawcid
                from database import get_cached_connection
                
                logger.info(f"Fetching content for topic_id: {topic_id}")
                conn = get_cached_connection()
                
                # Fetch content from 'notes' directly using a cursor to match topic_id to id
                cursor = conn.cursor(dictionary=True)
                cursor.execute("SELECT content, description FROM notes WHERE id = %s", (topic_id,))
                note_data = cursor.fetchone()
                cursor.close()
                
                if note_data:
                    content = note_data.get('content', '') or note_data.get('description', '')
                    logger.info(f"Fetched content from notes: {len(content)} chars")
                
                # If still empty, try 'videos' table (maybe it's a video transcript/description)
                if not content:
                    logger.info(f"Content empty in notes, checking videos for id {topic_id}")
                    cursor = conn.cursor(dictionary=True)
                    cursor.execute("SELECT short_description, description FROM videos WHERE id = %s", (topic_id,))
                    video_data = cursor.fetchone()
                    cursor.close()
                    
                    if video_data:
                        content = video_data.get('description', '') or video_data.get('short_description', '')
                        logger.info(f"Fetched content from videos: {len(content)} chars")

            except Exception as e:
                logger.error(f"Error fetching content from DB: {e}")
        
        if not content:
            return jsonify({"success": False, "message": "Content is required"}), 400
        
        if not ai_available:
            return jsonify({"success": False, "message": "AI service unavailable"}), 503
        
        # Import TokenManager
        try:
            from services.token_manager import TokenManager
        except ImportError:
            # If TokenManager missing (e.g. running in basic mode), skip checks
            TokenManager = None

        # Truncate content if too long
        max_content_length = 4000
        if len(content) > max_content_length:
            content = content[:max_content_length] + "..."
        
        # Get AI model
        model_name = data.get('model') or get_default_model(user_id=user_id)
        model_config = get_model_config(model_name)
        model = model_config.get('model', 'gpt-4')
        provider = model_config.get('provider', 'PollinationsAI')
        
        # Prepare messages
        messages = [
            {"role": "system", "content": FLASHCARD_SYSTEM_PROMPT},
            {"role": "user", "content": f"Create {count} flashcards from this content:\n\n{content}"}
        ]

        # Check balance if user_id present
        if user_id and TokenManager:
            # Estimate: Input + Expected Output (~50 tokens per card)
            est_input = _estimate_tokens(content)
            est_output = count * 60
            est_total = est_input + est_output
            
            # Simple pre-check without model cost calculation (safe buffer)
            check = TokenManager.check_balance(user_id, 10)  # Check minimum 10 tokens to allow start
            if not check['has_sufficient']:
                return jsonify({
                    "success": False, 
                    "message": "Insufficient tokens. Please top up."
                }), 402
        
        logger.info(f"Generating flashcards for topic {topic_id} using model {model}")
        
        # Generate with AI
        result = ai_generate(messages=messages, model=model)
        response = result.content
        
        if not response:
            return jsonify({"success": False, "message": "Empty response from AI"}), 500
            
        # Deduct actual tokens
        if user_id and TokenManager:
            # Input text is messages content
            input_text = messages[1]['content']
            _handle_token_billing(user_id, input_text, response, model, '/api/ai-study/generate-flashcards')
        
        # Parse JSON response
        flashcards = _parse_json_response(response)
        
        if not flashcards:
            logger.warning(f"Failed to parse flashcards from response: {response[:200]}")
            return jsonify({
                "success": False, 
                "message": "Failed to parse flashcards from AI response",
                "raw_response": response[:500]
            }), 500
        
        # Add IDs to flashcards
        for i, card in enumerate(flashcards):
            card['id'] = f"{topic_id or 'fc'}_{i+1}"
            if 'difficulty' not in card:
                card['difficulty'] = 'medium'
        
        logger.info(f"Generated {len(flashcards)} flashcards successfully")
        
        return jsonify({
            "success": True,
            "flashcards": flashcards,
            "count": len(flashcards),
            "model_used": model
        })
        
    except Exception as e:
        logger.error(f"Error generating flashcards: {str(e)}")
        return jsonify({"success": False, "message": f"Error: {str(e)}"}), 500


# ==================== SUMMARY GENERATION ====================

SUMMARY_SYSTEM_PROMPT = """You are "Sathi AI", the friendly AI learning companion of Ashlya Academy, specializing in creating clear, insightful summaries.

Create a well-structured summary of the provided content:
1. Start with a brief overview (2-3 sentences) capturing the core idea
2. List the key points as organized bullet points grouped by theme
3. Highlight important terms in **bold**
4. End with main takeaways and connections between concepts

Format with clean Markdown: use ## headings, **bold** for key terms, and bullet points.
When it helps understanding, include a ```mermaid diagram (e.g., mind map of concept relationships, flowchart of processes, or timeline of events).
For math content, use $...$ for inline equations and $$...$$ for display equations.

Use clear, simple language. Be concise yet thorough — capture the essence without losing important nuances.
"""

@ai_study_tools_bp.route('/api/ai-study/generate-summary', methods=['POST'])
@require_api_key
@encrypt_response
def generate_summary():
    """Generate AI-powered summary from content"""
    try:
        data = request.get_json()
        if not data:
            return jsonify({"success": False, "message": "No data provided"}), 400
        
        content = data.get('content', '')
        style = data.get('style', 'detailed')  # 'brief', 'detailed', 'bullet-points'
        
        if not content:
            return jsonify({"success": False, "message": "Content is required"}), 400
        
        if not ai_available:
            return jsonify({"success": False, "message": "AI service unavailable"}), 503
        
        # Import TokenManager
        try:
            from services.token_manager import TokenManager
        except ImportError:
            TokenManager = None

        # Truncate content if too long
        max_content_length = 6000
        if len(content) > max_content_length:
            content = content[:max_content_length] + "..."
            
        user_id = data.get('user_id')
        
        # Get AI model
        model_name = data.get('model') or get_default_model(user_id=user_id)
        model_config = get_model_config(model_name)
        model = model_config.get('model', 'gpt-4')
        provider = model_config.get('provider', 'PollinationsAI')
        
        style_instruction = {
            'brief': 'Create a brief 3-4 sentence summary.',
            'detailed': 'Create a detailed summary with sections.',
            'bullet-points': 'Create a bullet-point summary of key points.'
        }.get(style, 'Create a detailed summary.')
        
        messages = [
            {"role": "system", "content": SUMMARY_SYSTEM_PROMPT},
            {"role": "user", "content": f"{style_instruction}\n\nContent:\n{content}"}
        ]
        
        # Check balance
        if user_id and TokenManager:
             # Check minimum balance to start
            check = TokenManager.check_balance(user_id, 10)
            if not check['has_sufficient']:
                return jsonify({"success": False, "message": "Insufficient tokens."}), 402
        
        result = ai_generate(messages=messages, model=model)
        response = result.content
        
        if not response:
            return jsonify({"success": False, "message": "Empty response from AI"}), 500
            
        # Deduct actual tokens
        if user_id and TokenManager:
            input_text = messages[1]['content']
            _handle_token_billing(user_id, input_text, response, model, '/api/ai-study/generate-summary')
        
        return jsonify({
            "success": True,
            "summary": response,
            "style": style,
            "model_used": model
        })
        
    except Exception as e:
        logger.error(f"Error generating summary: {str(e)}")
        return jsonify({"success": False, "message": f"Error: {str(e)}"}), 500


# ==================== QUIZ GENERATION ====================

QUIZ_SYSTEM_PROMPT = """You are "Sathi AI", the friendly AI learning companion of Ashlya Academy, specializing in creating effective quiz questions.

Create multiple choice questions from the provided content:
1. Each question should test understanding, not just memorization
2. Provide 4 options (A, B, C, D) with only one correct answer
3. Include a brief explanation for the correct answer
4. Vary difficulty levels

**IMPORTANT: You MUST respond with ONLY a valid JSON array. No explanatory text.**

Response format:
[
  {
    "question": "Question text?",
    "options": ["A) Option 1", "B) Option 2", "C) Option 3", "D) Option 4"],
    "correct_answer": "A",
    "explanation": "Brief explanation why A is correct",
    "difficulty": "medium"
  }
]
"""

@ai_study_tools_bp.route('/api/ai-study/generate-quiz', methods=['POST'])
@require_api_key
@encrypt_response
def generate_quiz():
    """Generate AI-powered quiz questions from content"""
    try:
        data = request.get_json()
        if not data:
            return jsonify({"success": False, "message": "No data provided"}), 400
        
        content = data.get('content', '')
        count = data.get('count', 5)
        
        if not content:
            return jsonify({"success": False, "message": "Content is required"}), 400
        
        if not ai_available:
            return jsonify({"success": False, "message": "AI service unavailable"}), 503
            
        # Import TokenManager
        try:
            from services.token_manager import TokenManager
        except ImportError:
            TokenManager = None
            
        user_id = data.get('user_id')
        
        # Truncate content
        max_content_length = 4000
        if len(content) > max_content_length:
            content = content[:max_content_length] + "..."
        
        model_name = data.get('model') or get_default_model(user_id=user_id)
        model_config = get_model_config(model_name)
        model = model_config.get('model', 'gpt-4')
        provider = model_config.get('provider', 'PollinationsAI')
        
        messages = [
            {"role": "system", "content": QUIZ_SYSTEM_PROMPT},
            {"role": "user", "content": f"Create {count} quiz questions from this content:\n\n{content}"}
        ]
        
        # Check balance
        if user_id and TokenManager:
            check = TokenManager.check_balance(user_id, 10)
            if not check['has_sufficient']:
                return jsonify({"success": False, "message": "Insufficient tokens."}), 402
        
        result = ai_generate(messages=messages, model=model)
        response = result.content
        
        if not response:
            return jsonify({"success": False, "message": "Empty response from AI"}), 500
            
        # Deduct actual tokens
        if user_id and TokenManager:
            input_text = messages[1]['content']
            _handle_token_billing(user_id, input_text, response, model, '/api/ai-study/generate-quiz')
        
        questions = _parse_json_response(response)
        
        if not questions:
            return jsonify({
                "success": False,
                "message": "Failed to parse quiz questions",
                "raw_response": response[:500]
            }), 500
        
        # Add IDs
        for i, q in enumerate(questions):
            q['id'] = f"q_{i+1}"
        
        return jsonify({
            "success": True,
            "questions": questions,
            "count": len(questions),
            "model_used": model
        })
        
    except Exception as e:
        logger.error(f"Error generating quiz: {str(e)}")
        return jsonify({"success": False, "message": f"Error: {str(e)}"}), 500


# ==================== KEY CONCEPTS EXTRACTION ====================

CONCEPTS_SYSTEM_PROMPT = """You are "Sathi AI", the friendly AI learning companion of Ashlya Academy, specializing in identifying and explaining key concepts.

Extract the most important concepts from the content:
1. Identify 5-10 key concepts, terms, or ideas
2. Provide a brief, clear definition for each
3. Rate importance: 'essential', 'important', or 'supplementary'

**IMPORTANT: You MUST respond with ONLY a valid JSON array.**

Response format:
[
  {"concept": "Term", "definition": "Clear definition...", "importance": "essential"},
  {"concept": "Another Term", "definition": "Definition...", "importance": "important"}
]
"""

@ai_study_tools_bp.route('/api/ai-study/extract-concepts', methods=['POST'])
@require_api_key
@encrypt_response
def extract_concepts():
    """Extract key concepts from content using AI"""
    try:
        data = request.get_json()
        if not data:
            return jsonify({"success": False, "message": "No data provided"}), 400
        
        content = data.get('content', '')
        
        if not content:
            return jsonify({"success": False, "message": "Content is required"}), 400
        
        if not ai_available:
            return jsonify({"success": False, "message": "AI service unavailable"}), 503
            
        # Import TokenManager
        try:
            from services.token_manager import TokenManager
        except ImportError:
            TokenManager = None
            
        user_id = data.get('user_id')
        
        max_content_length = 4000
        if len(content) > max_content_length:
            content = content[:max_content_length] + "..."
        
        model_name = data.get('model') or get_default_model(user_id=user_id)
        model_config = get_model_config(model_name)
        model = model_config.get('model', 'gpt-4')
        provider = model_config.get('provider', 'PollinationsAI')
        
        messages = [
            {"role": "system", "content": CONCEPTS_SYSTEM_PROMPT},
            {"role": "user", "content": f"Extract key concepts from:\n\n{content}"}
        ]
        
        # Check balance
        if user_id and TokenManager:
            check = TokenManager.check_balance(user_id, 10)
            if not check['has_sufficient']:
                return jsonify({"success": False, "message": "Insufficient tokens."}), 402
        
        result = ai_generate(messages=messages, model=model)
        response = result.content
        
        if not response:
            return jsonify({"success": False, "message": "Empty response from AI"}), 500
            
        # Deduct actual tokens
        if user_id and TokenManager:
            input_text = messages[1]['content']
            _handle_token_billing(user_id, input_text, response, model, '/api/ai-study/extract-concepts')
        
        concepts = _parse_json_response(response)
        
        if not concepts:
            return jsonify({
                "success": False,
                "message": "Failed to parse concepts",
                "raw_response": response[:500]
            }), 500
        
        return jsonify({
            "success": True,
            "concepts": concepts,
            "count": len(concepts),
            "model_used": model
        })
        
    except Exception as e:
        logger.error(f"Error extracting concepts: {str(e)}")
        return jsonify({"success": False, "message": f"Error: {str(e)}"}), 500


# ==================== EXPLAIN CONTENT ====================

@ai_study_tools_bp.route('/api/ai-study/explain', methods=['POST'])
@require_api_key
@encrypt_response
def explain_content():
    """Get AI explanation of content in simpler terms"""
    try:
        data = request.get_json()
        if not data:
            return jsonify({"success": False, "message": "No data provided"}), 400
        
        content = data.get('content', '')
        level = data.get('level', 'simple')  # 'simple', 'intermediate', 'detailed'
        
        if not content:
            return jsonify({"success": False, "message": "Content is required"}), 400
        
        if not ai_available:
            return jsonify({"success": False, "message": "AI service unavailable"}), 503
            
        # Import TokenManager
        try:
            from services.token_manager import TokenManager
        except ImportError:
            TokenManager = None
            
        user_id = data.get('user_id')
        
        level_instructions = {
            'simple': 'Explain this like I\'m a beginner. Use simple words and analogies.',
            'intermediate': 'Provide a clear, structured explanation with examples.',
            'detailed': 'Give a comprehensive explanation with technical details and examples.'
        }
        
        model_name = data.get('model') or get_default_model(user_id=user_id)
        model_config = get_model_config(model_name)
        model = model_config.get('model', 'gpt-4')
        provider = model_config.get('provider', 'PollinationsAI')
        
        messages = [
            {"role": "system", "content": "You are 'Sathi AI', the friendly AI learning companion of Ashlya Academy, who explains complex topics clearly and engagingly. Use real-world analogies, step-by-step breakdowns, and visual aids to make concepts stick. Format with clean Markdown: use headings (##), **bold** key terms, and bullet points. Use ```mermaid diagrams (flowcharts, mind maps, etc.) when they help visualize relationships or processes. Use inline SVG for geometric figures or graphs. For math, use $...$ inline and $$...$$ display equations with double backslashes for LaTeX commands. Anticipate common misconceptions and address them proactively."},
            {"role": "user", "content": f"{level_instructions.get(level, level_instructions['simple'])}\n\nContent to explain:\n{content}"}
        ]
        
        # Check balance
        if user_id and TokenManager:
            check = TokenManager.check_balance(user_id, 10)
            if not check['has_sufficient']:
                return jsonify({"success": False, "message": "Insufficient tokens."}), 402
        
        result = ai_generate(messages=messages, model=model)
        response = result.content
        
        if not response:
            return jsonify({"success": False, "message": "Empty response from AI"}), 500
            
        # Deduct actual tokens
        if user_id and TokenManager:
            input_text = messages[1]['content']
            _handle_token_billing(user_id, input_text, response, model, '/api/ai-study/explain')
        
        return jsonify({
            "success": True,
            "explanation": response,
            "level": level,
            "model_used": model
        })
        
    except Exception as e:
        logger.error(f"Error explaining content: {str(e)}")
        return jsonify({"success": False, "message": f"Error: {str(e)}"}), 500


# ==================== CONVERSATIONAL AUDIO SCRIPT GENERATION ====================

AUDIO_SCRIPT_SYSTEM_PROMPT = """You are "Sathi AI", the friendly AI learning companion of Ashlya Academy, creating an engaging audio lesson script. Your task is to transform educational content into a natural, spoken-word script that sounds like a real teacher explaining a topic.

IMPORTANT GUIDELINES:

1. **Speak like a real teacher**: Start with a warm greeting and topic introduction. Example: "Hello everyone! I'm Sathi, your learning companion. Today, we're going to explore a fascinating topic..."

2. **Teaching structure**:
   - Begin with WHY this topic matters (hook the listener)
   - Explain concepts step by step, building on each other
   - Use analogies and real-world examples to clarify abstract ideas
   - Pause to check understanding: "Does this make sense so far?", "Let's think about this..."
   - Summarize key points at the end

3. **Natural speech patterns**:
   - Use contractions: "it's", "we're", "that's" instead of "it is", "we are"
   - Include thinking phrases: "Now, here's the interesting part...", "You might be wondering..."
   - Add emphasis words: "This is really important...", "Pay close attention to this..."
   - Use "..." for natural pauses

4. **Engagement techniques**:
   - Ask rhetorical questions: "Have you ever noticed...?", "What do you think happens when...?"
   - Give memory tips: "Here's a trick to remember this..."
   - Share interesting facts or stories related to the topic

5. **Formatting rules**:
   - NO bullet points, numbers, or markdown - this is pure spoken text
   - NO asterisks, hashtags, or special formatting characters
   - Write everything as natural flowing speech
   - Keep paragraphs short (2-3 sentences each)

6. **Pacing**:
   - Don't rush through complex concepts
   - Give the listener time to absorb information
   - Aim for a comfortable 150-160 words per minute pace

Create a script that sounds like a passionate, friendly companion who loves the subject and wants students to truly understand and remember the material.

**Output ONLY the audio script text - no titles, no formatting, just the spoken content.**
"""

@ai_study_tools_bp.route('/api/ai-study/generate-audio-script', methods=['POST'])
@require_api_key
@encrypt_response
def generate_audio_script():
    """Generate a conversational audio script from educational content for TTS"""
    try:
        data = request.get_json()
        if not data:
            return jsonify({"success": False, "message": "No data provided"}), 400
        
        content = data.get('content', '')
        topic_title = data.get('topic_title', 'this topic')
        style = data.get('style', 'conversational')  # 'conversational', 'formal', 'brief'
        
        if not content:
            return jsonify({"success": False, "message": "Content is required"}), 400
        
        if not ai_available:
            return jsonify({"success": False, "message": "AI service unavailable"}), 503
            
        # Import TokenManager
        try:
            from services.token_manager import TokenManager
        except ImportError:
            TokenManager = None
            
        user_id = data.get('user_id')
        
        # Truncate content if too long (leave room for AI response)
        max_content_length = 5000
        if len(content) > max_content_length:
            content = content[:max_content_length] + "..."
        
        # Get AI model
        model_name = data.get('model') or get_default_model(user_id=user_id)
        model_config = get_model_config(model_name)
        model = model_config.get('model', 'gpt-4')
        provider = model_config.get('provider', 'PollinationsAI')
        
        # Style-specific instructions
        style_instructions = {
            'conversational': f'Create a friendly, engaging audio script about "{topic_title}". Make it feel like a casual but informative explanation.',
            'formal': f'Create a professional, structured audio script about "{topic_title}". Keep it clear and authoritative.',
            'brief': f'Create a short, concise audio summary of "{topic_title}" (about 2-3 minutes when spoken).'
        }.get(style, f'Create an engaging audio script about "{topic_title}".')
        
        messages = [
            {"role": "system", "content": AUDIO_SCRIPT_SYSTEM_PROMPT},
            {"role": "user", "content": f"{style_instructions}\n\nContent to convert:\n\n{content}"}
        ]
        
        # Check balance
        if user_id and TokenManager:
            check = TokenManager.check_balance(user_id, 10)
            if not check['has_sufficient']:
                return jsonify({"success": False, "message": "Insufficient tokens."}), 402
        
        logger.info(f"Generating audio script for '{topic_title}' using model {model}")
        
        result = ai_generate(messages=messages, model=model)
        response = result.content
        
        if not response:
            return jsonify({"success": False, "message": "Empty response from AI"}), 500
            
        # Deduct actual tokens
        if user_id and TokenManager:
            input_text = messages[1]['content']
            _handle_token_billing(user_id, input_text, response, model, '/api/ai-study/generate-audio-script')
        
        # Clean up the script - remove any markdown artifacts
        script = response
        script = script.replace('**', '')
        script = script.replace('*', '')
        script = script.replace('#', '')
        script = script.replace('`', '')
        script = script.strip()
        
        # Estimate duration (average speaking rate is ~150 words per minute)
        word_count = len(script.split())
        estimated_duration_minutes = round(word_count / 150, 1)
        
        logger.info(f"Generated audio script: {word_count} words, ~{estimated_duration_minutes} min")
        
        return jsonify({
            "success": True,
            "script": script,
            "word_count": word_count,
            "estimated_duration_minutes": estimated_duration_minutes,
            "style": style,
            "model_used": model
        })
        
    except Exception as e:
        logger.error(f"Error generating audio script: {str(e)}")
        return jsonify({"success": False, "message": f"Error: {str(e)}"}), 500


# ==================== HELPER FUNCTIONS ====================

def _parse_json_response(response_text):
    """Parse JSON from AI response, handling markdown code blocks"""
    if not response_text:
        return None
    
    try:
        # Try direct JSON parse first
        return json.loads(response_text)
    except json.JSONDecodeError:
        pass
    
    # Try to extract JSON from markdown code blocks
    json_patterns = [
        r'```json\s*([\s\S]*?)\s*```',
        r'```\s*([\s\S]*?)\s*```',
        r'\[\s*\{[\s\S]*\}\s*\]'
    ]
    
    for pattern in json_patterns:
        match = re.search(pattern, response_text)
        if match:
            try:
                json_str = match.group(1) if '```' in pattern else match.group(0)
                return json.loads(json_str)
            except (json.JSONDecodeError, IndexError):
                continue
    
    # Try to find array start/end
    start_idx = response_text.find('[')
    end_idx = response_text.rfind(']')
    
    if start_idx != -1 and end_idx != -1 and end_idx > start_idx:
        try:
            return json.loads(response_text[start_idx:end_idx + 1])
        except json.JSONDecodeError:
            pass
    
    return None

def _estimate_tokens(text):
    """Simple estimation of tokens (words + 20% buffer)"""
    if not text:
        return 0
    return int(len(text.split()) * 1.3)  # Rough estimate

def _handle_token_billing(user_id, input_text, output_text, model_name, endpoint):
    """Handle actual token billing calculation and deduction"""
    if not user_id:
        return {'success': True}
        
    try:
        try:
            from services.token_manager import TokenManager
        except ImportError:
            return {'success': True, 'message': 'TokenManager unavailable'}
            
        input_tokens = len(input_text.split())
        output_tokens = len(output_text.split()) if output_text else 0
        total_tokens = input_tokens + output_tokens
        
        # Calculate cost
        token_cost = TokenManager.calculate_token_cost(
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            model_name=model_name
        )
        
        # Determine actual deduction amount (cost or discounted if cost is 0)
        deduction_amount = token_cost if token_cost > 0 else output_tokens + (input_tokens * 0.1)
        
        # Deduct
        result = TokenManager.deduct_tokens(
            user_id=user_id,
            amount=deduction_amount,
            endpoint=endpoint,
            ai_model=model_name,
            request_tokens=input_tokens,
            response_tokens=output_tokens
        )
        
        logger.info(f"Billed user {user_id} {deduction_amount} tokens for {endpoint}")
        return result
        
    except Exception as e:
        logger.error(f"Error handling token billing: {e}")
        # Don't block response on billing error
        return {'success': False, 'message': str(e)}
