from flask import Blueprint, request, jsonify, Response
import os
import uuid
import sys
import importlib.util
import requests
import json
import time
import random
import re
from utils import require_api_key

# Blueprint for AI-related endpoints
ai_bp = Blueprint('ai', __name__, url_prefix='/api/ai')

# Import central AI provider
from services.ai_provider import ai_generate, ai_generate_stream, ai_generate_vision

# Keep g4f availability check for backward compat (the provider handles fallback)
try:
    import g4f
    g4f_available = True
except ImportError:
    g4f_available = False

# AI features are always available when Groq is configured
from services.ai_config import is_groq_enabled
ai_available = g4f_available or is_groq_enabled()

# Refined system prompt for educational context with strict formatting rules
EDUCATIONAL_CONTEXT = '''You are "Sathi AI", the friendly AI learning companion of Ashlya Academy. Your responses must be accurate, helpful, and thoughtfully structured. "Sathi" means "companion/friend" — be warm, encouraging, and genuinely supportive while helping students learn!

**FORMATTING RULES:**
1.  **Math Delimiters:**
    - Use `$$...$$` for display equations on their own line.
    - Use `$ ... $` for inline math expressions within a sentence.
    - **NEVER** use `\\[ ... \\]` or `\\( ... \\)`.

2.  **LaTeX Backslashes:**
    - **ALWAYS** use double backslashes (`\\\\`) for LaTeX commands.
    - Correct: `\\\\frac{a}{b}`, `\\\\sqrt{x}`, `\\\\sum_{i=1}^{n}`
    - Incorrect: `\\frac{a}{b}`, `\\sqrt{x}`

3.  **Code Blocks:**
    - Always specify the language for syntax highlighting (e.g., ```python, ```javascript).

4.  **Markdown Structure:**
    - Use `#` for main headings and `##` for subheadings.
    - Use **bold** for key terms and `inline code` for variables or function names.
    - Use tables for comparisons and structured data.

5.  **Diagrams & Visual Aids (USE PROACTIVELY):**
    - Use ```mermaid code blocks for flowcharts, mind maps, sequence diagrams, class diagrams, state diagrams, ER diagrams, Gantt charts, pie charts, or timelines. Our app renders these as interactive visuals.
    - Use inline SVG (`<svg>...</svg>`) for geometric figures, coordinate planes, graphs, circuits, or custom illustrations. Our app renders SVG as actual graphics.
    - Use diagrams whenever they clarify relationships, processes, hierarchies, structures, or data flows better than text alone.
    - Prefer mermaid for conceptual diagrams (flow, hierarchy, process) and SVG for precise visual representations (geometry, graphs, custom figures).

6.  **Teaching Best Practices:**
    - Break complex topics into clear, digestible steps.
    - Use real-world analogies and examples to ground abstract concepts.
    - Anticipate common misconceptions and address them.
    - Provide "why" explanations, not just "what" and "how".
'''

def enhance_markdown_formatting(text):
    
    if not isinstance(text, str):
        return ""
    # Step 1: Normalize common LaTeX delimiters to $ and $$
    text = re.sub(r'\\\[(.*?)\\\]', r'$$\1$$', text, flags=re.DOTALL)
    text = re.sub(r'\\\((.*?)\\\)', r'$\1$', text, flags=re.DOTALL)

    # Step 2: Process content inside math delimiters
    def process_math_content(match):
        delimiters, content = match.group(1), match.group(2)
        content = content.strip()
        # Ensure LaTeX commands have double backslashes
        content = re.sub(r'(?<!\\)\\([a-zA-Z]+)', r'\\\\\1', content)
        return f"\n\n$$\n{content}\n$$\n\n" if delimiters == '$$' else f'${content}$'
    
    text = re.sub(r'(\$\$|\$)(.*?)(\1)', process_math_content, text, flags=re.DOTALL)
    
    # Step 3: Clean up spacing and newlines
    text = re.sub(r'([a-zA-Z0-9])\$', r'\1 $', text)
    text = re.sub(r'\$([a-zA-Z0-9])', r'$ \1', text)
    text = re.sub(r'\n{3,}', '\n\n', text)
    return text.strip()

def _prepare_messages(data):
    messages = data.get('messages', [])
    language_code = data.get('language', 'en-US')
    language_instructions = ""
    if language_code != 'en-US':
        language_names = {'hi-IN': 'Hindi', 'ne-NP': 'Nepali', 'es-ES': 'Spanish', 'zh-CN': 'Chinese', 'fr-FR': 'French', 'de-DE': 'German', 'ja-JP': 'Japanese', 'ko-KR': 'Korean'}
        language_name = language_names.get(language_code, language_code)
        language_instructions = f"""\n\n**CRITICAL LANGUAGE REQUIREMENT:**
You MUST respond ENTIRELY in {language_name} language.
DO NOT respond in English or any other language.
DO NOT mix languages.
Every word of your response must be in {language_name}.
This is a strict requirement - respond ONLY in {language_name}."""

    final_system_content = EDUCATIONAL_CONTEXT
    if language_instructions:
        final_system_content += language_instructions

    if not any(msg.get('role') == 'system' for msg in messages):
        messages.insert(0, {'role': 'system', 'content': final_system_content})
    else:
        for msg in messages:
            if msg.get('role') == 'system':
                msg['content'] = f"{final_system_content}\n\n{msg.get('content', '')}"
                break
    return messages

@ai_bp.route('/generate', methods=['POST'])
@require_api_key
def generate_response():
    # Support both JSON and multipart form-data with images
    if request.content_type and request.content_type.startswith('multipart/form-data'):
        # Handle prompt + optional images in a single endpoint
        form_prompt = request.form.get('prompt', '')
        language = request.form.get('language', 'en-US')
        original_model = request.form.get('model_', None)
        is_stream = request.form.get('stream', 'false').lower() == 'true'
        
        # Get model dynamically from database
        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
        
        # Get user_id for model preference
        user_id = request.form.get('user_id') or request.headers.get('X-User-ID') or request.args.get('user_id')
        
        if not original_model:
            original_model = get_default_model(user_id=user_id)
        model_config = get_model_config(original_model)
        model = model_config.get('model', 'gpt-5-mini')

        # Collect all images from common keys
        files = []
        # images list
        files.extend(request.files.getlist('images'))
        # single image
        single_image = request.files.get('image')
        if single_image:
            files.append(single_image)
        # any other uploaded files
        for key in request.files.keys():
            for f in request.files.getlist(key):
                if f not in files:
                    files.append(f)

        # If no images, fall back to JSON-style processing using provided prompt
        if not files:
            data = {
                'messages': [{'role': 'user', 'content': form_prompt}],
                'language': language,
                'model_': original_model,
                'stream': is_stream
            }
            return generate_regular_response(data)

        if not ai_available:
            return jsonify({'success': False, 'message': 'AI image analysis unavailable'}), 503

        # Prepare messages with system prompt
        data = {
            'messages': [{'role': 'user', 'content': form_prompt or 'Analyze the attached images and respond with insights.'}],
            'language': language
        }
        messages = _prepare_messages(data)

        # Prepare image payload [[file_obj, filename], ...]
        images_payload = []
        for f in files:
            if getattr(f, 'filename', ''):
                try:
                    images_payload.append([f.stream, f.filename or f"image_{uuid.uuid4().hex}.png"])
                except Exception:
                    pass

        if not images_payload:
            return jsonify({'success': False, 'message': 'No readable images provided'}), 400

        try:
            result = ai_generate_vision(
                messages=messages,
                images=images_payload,
                model=model,
            )
            content_text = result.content

            if not content_text:
                return jsonify({'success': False, 'message': 'Empty response from AI'}), 500

            formatted = enhance_markdown_formatting(content_text)
            return jsonify({'success': True, 'content': formatted, 'model': result.model, 'message': 'Response generated successfully'})
        except Exception as e:
            return jsonify({'success': False, 'message': f'Failed to generate response with images: {str(e)}'}), 500

    # JSON path (existing behavior)
    data = request.json
    if not data or 'messages' not in data:
        return jsonify({'error': 'Missing required parameter: messages'}), 400
    return generate_stream_response(data) if data.get('stream', False) else generate_regular_response(data)

@ai_bp.route('/generate-stream', methods=['POST'])
@require_api_key
def generate_stream_endpoint():
    data = request.json
    if not data or 'messages' not in data:
        return jsonify({'error': 'Missing required parameter: messages'}), 400
    return generate_stream_response(data)

def generate_regular_response(data):
    try:
        if not ai_available:
            return jsonify({
                'success': False, 
                'message': "AI service is unavailable",
                'content': "AI service is unavailable."
            }), 503
        
        messages = _prepare_messages(data)
        
        # Get user_id from request headers for token deduction and model preference
        user_id = request.headers.get('X-User-ID') or request.args.get('user_id') or data.get('user_id')
        
        # Use user's preferred model instead of hardcoding
        model = data.get('model_') or get_default_model(user_id=user_id)
        
        # Import TokenManager for billing
        import sys
        import os
        sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
        from services.token_manager import TokenManager
        import logging
        logger = logging.getLogger(__name__)
        
        # Get user message for token calculation
        user_message = messages[-1].get('content', '') if messages else ''
        
        result = ai_generate(messages=messages, model=model)
        response_text = result.content
        
        if not response_text or not response_text.strip():
            return jsonify({
                'success': False,
                'message': "Failed to get a response from AI service",
                'content': "Failed to get a response."
            }), 500

        formatted_response = enhance_markdown_formatting(response_text)
        
        # ========== TOKEN DEDUCTION FOR REGULAR RESPONSES ==========
        if user_id:
            try:
                logger.info(f"[REGULAR TOKEN DEBUG] Starting token calculation for user {user_id}")
                
                input_tokens = result.input_tokens or len(user_message.split())
                output_tokens = result.output_tokens or len(response_text.split())
                
                logger.info(f"[REGULAR TOKEN DEBUG] Input tokens: {input_tokens}, Output tokens: {output_tokens}")
                
                token_cost = TokenManager.calculate_token_cost(
                    input_tokens=input_tokens,
                    output_tokens=output_tokens,
                    model_name=model
                )
                
                tokens_to_deduct = token_cost if token_cost > 0 else output_tokens + (input_tokens * 0.5)
                logger.info(f"[REGULAR TOKEN DEBUG] Tokens to deduct: {tokens_to_deduct}")
                
                deduction_result = TokenManager.deduct_tokens(
                    user_id=int(user_id),
                    amount=tokens_to_deduct,
                    endpoint='/api/ai/generate',
                    ai_model=result.model,
                    provider=result.provider,
                    request_tokens=input_tokens,
                    response_tokens=output_tokens,
                    metadata={'streaming': False}
                )
                
                if deduction_result['success']:
                    logger.info(f"[REGULAR TOKEN SUCCESS] ✅ Deducted {tokens_to_deduct:.2f} tokens from user {user_id}")
                else:
                    logger.error(f"[REGULAR TOKEN ERROR] Deduction failed: {deduction_result.get('message')}")
            except Exception as token_error:
                logger.error(f"[REGULAR TOKEN ERROR] Exception during deduction: {token_error}")
        # ========== END TOKEN DEDUCTION ==========
        
        return jsonify({
            'success': True,
            'content': formatted_response,
            'model': result.model,
            'message': 'Response generated successfully'
        })
    except Exception as e:
        print(f"Error in generate_regular_response: {e}")
        return jsonify({
            'success': False,
            'error': 'Failed to generate response', 
            'message': f"AI service error: {str(e)}",
            'details': str(e)
        }), 500

def generate_stream_response(data):
    if not ai_available:
        def error_stream():
            yield 'data: {"type": "error", "content": "AI service unavailable"}\n\n'
        return Response(error_stream(), mimetype='text/event-stream')

    messages = _prepare_messages(data)
    model = data.get('model_', None)
    
    # Get user_id from request headers or query params for token deduction
    user_id = request.headers.get('X-User-ID') or request.args.get('user_id')
    
    # Import TokenManager
    import sys
    import os
    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    from services.token_manager import TokenManager
    import logging
    logger = logging.getLogger(__name__)

    def generate():
        yield 'data: {"type":"start"}\n\n'
        
        complete_response = ""
        user_message = messages[-1].get('content', '') if messages else ''

        try:
            for chunk_str in ai_generate_stream(messages=messages, model=model):
                if chunk_str:
                    complete_response += chunk_str
                    yield f'data: {json.dumps({"type": "chunk", "content": chunk_str})}\n\n'
            
            # Format the complete response at the end
            formatted_final_response = enhance_markdown_formatting(complete_response)
            yield f'data: {json.dumps({"type": "end", "content": formatted_final_response})}\n\n'
            
            # ========== TOKEN DEDUCTION FOR STREAMING ==========
            if user_id:
                try:
                    input_tokens = len(user_message.split())
                    output_tokens = len(complete_response.split())
                    
                    token_cost = TokenManager.calculate_token_cost(
                        input_tokens=input_tokens,
                        output_tokens=output_tokens,
                        model_name=model or 'llama-3.3-70b-versatile'
                    )
                    
                    tokens_to_deduct = token_cost if token_cost > 0 else output_tokens + (input_tokens * 0.5)
                    
                    deduction_result = TokenManager.deduct_tokens(
                        user_id=int(user_id),
                        amount=tokens_to_deduct,
                        endpoint='/api/ai/generate-stream',
                        ai_model=model or 'llama-3.3-70b-versatile',
                        provider='groq',
                        request_tokens=input_tokens,
                        response_tokens=output_tokens,
                        metadata={'streaming': True}
                    )
                    
                    if deduction_result['success']:
                        logger.info(f"[STREAM TOKEN SUCCESS] ✅ Deducted {tokens_to_deduct:.2f} tokens from user {user_id}")
                    else:
                        logger.error(f"[STREAM TOKEN ERROR] Deduction failed: {deduction_result.get('message')}")
                except Exception as token_error:
                    logger.error(f"[STREAM TOKEN ERROR] Exception during deduction: {token_error}")
            # ========== END TOKEN DEDUCTION ==========

        except Exception as e:
            logger.error(f"AI streaming error: {e}")
            error_content = f"An AI provider error occurred: {str(e)}"
            yield f'data: {json.dumps({"type": "error", "content": error_content})}\n\n'
            yield f'data: {json.dumps({"type": "end", "content": ""})}\n\n'

    return Response(generate(), mimetype='text/event-stream')
