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')

# Check if the g4f library is available, otherwise use a fallback
try:
    import g4f
    g4f_available = True
    print("g4f package is available. AI features enabled.")
except ImportError:
    g4f_available = False
    print("Warning: g4f package not found. Using fallback for AI responses.")

# Optional high-level client for image inputs
try:
    from g4f.client import Client as G4FClient
    g4f_client = G4FClient()
    g4f_client_available = True
except Exception:
    g4f_client_available = False

# Refined system prompt for educational context with strict formatting rules
EDUCATIONAL_CONTEXT = '''You are an educational AI assistant for Ashlya Academy. Your responses must be accurate, helpful, and formatted in clean, structured Markdown for our custom MathJax-powered renderer.

**CRITICAL 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.
'''

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'}
        language_name = language_names.get(language_code, language_code)
        language_instructions = f"Respond in {language_name}."

    final_system_content = EDUCATIONAL_CONTEXT
    if language_instructions:
        final_system_content += f"\n\n{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')
        model = request.form.get('model', 'gpt-5-nano')

        # 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': model,
                'stream': False
            }
            return generate_regular_response(data)

        if not g4f_available or not g4f_client_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:
            resp = g4f_client.chat.completions.create(
                messages=messages,
                model=model,
                images=images_payload,
                provider="PollinationsAI"
            )
            # Extract text content
            content_text = None
            try:
                if hasattr(resp, 'choices') and resp.choices:
                    msg = getattr(resp.choices[0], 'message', None) or {}
                    content_text = getattr(msg, 'content', None) or msg.get('content')
            except Exception:
                pass
            if not content_text and isinstance(resp, dict) and resp.get('choices'):
                content_text = resp['choices'][0]['message']['content']
            if not content_text:
                content_text = str(resp) if resp is not None else ''

            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': 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 g4f_available:
            return jsonify({
                'success': False, 
                'message': "AI service is unavailable",
                'content': "AI service is unavailable."
            }), 503
        
        messages = _prepare_messages(data)
        model = data.get('model', 'gpt-5-nano')
        
        response_text = g4f.ChatCompletion.create(model=model, messages=messages, stream=False, provider="PollinationsAI")
        
        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)
        
        return jsonify({
            'success': True,
            'content': formatted_response,
            'model': 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 g4f_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', 'gpt-5-nano')

    def generate():
        yield 'data: {"type":"start"}\n\n'
        
        complete_response = ""
        
        # List of noise patterns to ignore from the g4f stream
        noise_patterns = [
            "[WinError", 
            "Response 429", 
            "文字过长",
            "refused the network connection",
            "Generating summary...",
        ]

        try:
            stream_response = g4f.ChatCompletion.create(model=model, messages=messages, stream=True, provider="PollinationsAI")
            
            for chunk in stream_response:
                chunk_str = str(chunk) if chunk else ""

                # Filter out known noise patterns from the stream
                if any(pattern in chunk_str for pattern in noise_patterns):
                    continue  # Skip this chunk and move to the next one

                if chunk_str:
                    complete_response += chunk_str
                    yield f'data: {json.dumps({"type": "chunk", "content": chunk_str})}\n\n'
            
            # Format the complete, cleaned 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'

        except Exception as e:
            print(f"g4f 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')
