"""
AI Vision API Routes
Handles vision model requests with image analysis
"""
import json
import logging
import base64
import os
from datetime import datetime
from flask import Blueprint, request, jsonify
from werkzeug.utils import secure_filename

from utils import require_api_key, encrypt_response
from database import get_db_connection
from services.ai_provider import ai_generate_vision


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

# Create blueprint
ai_vision_bp = Blueprint('ai_vision', __name__)

# Allowed image extensions
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'webp', 'bmp'}

def allowed_file(filename):
    """Check if file extension is allowed"""
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

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_vision_bp.route('/api/ai/vision/analyze', methods=['POST'])
@require_api_key
@encrypt_response
def analyze_image():
    """
    Analyze image with AI vision model
    Expects multipart/form-data with:
    - image: image file
    - message: user question about the image
    - user_id: user identifier
    - model: model name (optional, defaults to vision-capable model)
    - provider: provider name (optional, defaults to Copilot)
    """
    try:
        # Validate request
        if 'image' not in request.files:
            return create_error_response("No image file provided", 400)
        
        if 'message' not in request.form:
            return create_error_response("No message provided", 400)
        
        user_id = request.form.get('user_id')
        if not user_id:
            return create_error_response("User ID is required", 400)
        
        # Get parameters
        image_file = request.files['image']
        message = request.form.get('message')
        model_name = request.form.get('model', '')  # Empty string for auto-select
        provider = request.form.get('provider', 'Copilot')  # Copilot has good vision support
        
        # Validate image
        if image_file.filename == '':
            return create_error_response("No image file selected", 400)
        
        if not allowed_file(image_file.filename):
            return create_error_response(
                f"Invalid file type. Allowed types: {', '.join(ALLOWED_EXTENSIONS)}", 
                400
            )

        # Import TokenManager
        try:
            from services.token_manager import TokenManager
        except ImportError:
            TokenManager = None

        VISION_IMAGE_TOKENS = 1000
        
        # Check balance
        if user_id and TokenManager:
            # Estimate: Image (1000) + Message words + Expected buffer (100)
            est_tokens = VISION_IMAGE_TOKENS + len(message.split()) + 100
            check = TokenManager.check_balance(user_id, 10) # Minimum entry
            if not check['has_sufficient']:
                return create_error_response("Insufficient tokens for vision analysis", 402)

        # Read image data
        image_file.seek(0)
        image_data = image_file.read()
        
        # Convert to base64 data URI
        filename = secure_filename(image_file.filename)
        file_extension = filename.rsplit('.', 1)[1].lower()
        
        # Create data URI
        base64_image = base64.b64encode(image_data).decode('utf-8')
        mime_type = f"image/{file_extension if file_extension != 'jpg' else 'jpeg'}"
        data_uri = f"data:{mime_type};base64,{base64_image}"
        
        logger.info(f"Processing vision request for user {user_id} with model {model_name or 'auto'}")
        logger.info(f"Image size: {len(image_data)} bytes, type: {mime_type}")
        
        # Enhanced system prompt for vision tasks
        system_prompt = (
            "You are 'Sathi AI', the friendly AI learning companion of Ashlya Academy, with vision capabilities. "
            "Carefully examine the provided image and describe what you see in detail. "
            "Then answer the user's question using both the visual content and your knowledge. "
            "Provide specific, accurate observations from the image to support your response. "
            "Format your response in clean Markdown with headings, bold key terms, and bullet points. "
            "When relevant, use ```mermaid code blocks for diagrams (flowcharts, mind maps, etc.) or inline SVG for geometric/visual illustrations — our app renders both as interactive graphics. "
            "Use $...$ for inline math and $$...$$ for display equations (double backslashes for LaTeX commands). "
            "Be thorough, well-structured, and genuinely helpful!"
        )
        
        # Prepare messages
        messages = [
            {
                "role": "system",
                "content": system_prompt
            },
            {
                "role": "user",
                "content": message
            }
        ]
        
        # Prepare images array [[data_uri, filename]]
        images = [[data_uri, filename]]
        
        # Call AI vision API
        try:
            result = ai_generate_vision(
                messages=messages,
                images=images,
                model=model_name or None,
            )
            
            ai_response = result.content
            
            logger.info(f"Vision analysis successful for user {user_id} via {result.provider}")
            
            # Deduct tokens
            if user_id and TokenManager:
                try:
                    input_tokens = VISION_IMAGE_TOKENS + len(message.split())
                    output_tokens = len(ai_response.split())
                    
                    token_cost = TokenManager.calculate_token_cost(
                        input_tokens=input_tokens,
                        output_tokens=output_tokens,
                        model_name=model_name or 'vision-auto'
                    )
                    tokens_to_deduct = token_cost if token_cost > 0 else output_tokens + (input_tokens * 0.1)
                    
                    TokenManager.deduct_tokens(
                        user_id=user_id,
                        amount=tokens_to_deduct,
                        endpoint='/api/ai/vision/analyze',
                        ai_model=model_name or 'vision-auto',
                        request_tokens=input_tokens,
                        response_tokens=output_tokens
                    )
                    logger.info(f"Deducted {tokens_to_deduct:.1f} tokens for vision analysis user {user_id}")
                except Exception as e:
                    logger.error(f"Error deducting vision tokens: {e}")

            # Store in database (optional)
            try:
                with get_db_connection() as conn:
                    cursor = conn.cursor()
                    cursor.execute("""
                        INSERT INTO ai_vision_requests 
                        (user_id, message, model_name, provider, response, image_size, created_at)
                        VALUES (%s, %s, %s, %s, %s, %s, NOW())
                    """, (
                        user_id,
                        message,
                        model_name or 'auto',
                        provider,
                        ai_response,
                        len(image_data)
                    ))
                    conn.commit()
            except Exception as db_error:
                logger.warning(f"Failed to store vision request in DB: {db_error}")
                # Continue even if DB storage fails
            
            return format_api_response({
                'response': ai_response,
                'model': result.model,
                'provider': result.provider,
                'image_size': len(image_data),
                'timestamp': datetime.now().isoformat()
            })
            
        except Exception as vision_error:
            logger.error(f"Vision API error: {str(vision_error)}")
            return create_error_response(
                f"Vision analysis failed: {str(vision_error)}",
                500
            )
        
    except Exception as e:
        logger.error(f"Error in analyze_image: {str(e)}")
        return create_error_response(f"Failed to analyze image: {str(e)}", 500)

@ai_vision_bp.route('/api/ai/vision/analyze-base64', methods=['POST'])
@require_api_key
@encrypt_response
def analyze_image_base64():
    """
    Analyze image with AI vision model using base64 encoded image
    Expects JSON with:
    - image_base64: base64 encoded image data
    - message: user question about the image
    - user_id: user identifier
    - model: model name (optional)
    - provider: provider name (optional)
    - filename: original filename (optional)
    """
    try:
        data = request.get_json()
        
        if not data:
            return create_error_response("No JSON data provided", 400)
        
        # Validate required fields
        if 'image_base64' not in data:
            return create_error_response("No image_base64 provided", 400)
        
        if 'message' not in data:
            return create_error_response("No message provided", 400)
        
        user_id = data.get('user_id')
        if not user_id:
            return create_error_response("User ID is required", 400)
        
        # Get parameters
        image_base64 = data['image_base64']
        message = data['message']
        model_name = data.get('model', '')
        provider = data.get('provider', 'Copilot')
        filename = data.get('filename', 'image.jpg')
        
        # Import TokenManager
        try:
            from services.token_manager import TokenManager
        except ImportError:
            TokenManager = None

        VISION_IMAGE_TOKENS = 1000
        
        # Check balance
        if user_id and TokenManager:
            est_tokens = VISION_IMAGE_TOKENS + len(message.split()) + 100
            check = TokenManager.check_balance(user_id, 10)
            if not check['has_sufficient']:
                return create_error_response("Insufficient tokens for vision analysis", 402)
        
        # Ensure base64 string has proper data URI format
        if not image_base64.startswith('data:'):
            # Assume JPEG if not specified
            image_base64 = f"data:image/jpeg;base64,{image_base64}"
        
        logger.info(f"Processing base64 vision request for user {user_id}")
        
        # Enhanced system prompt for vision tasks
        system_prompt = (
            "You are 'Sathi AI', the friendly AI learning companion of Ashlya Academy, with vision capabilities. "
            "Carefully examine the provided image and describe what you see in detail. "
            "Then answer the user's question using both the visual content and your knowledge. "
            "Provide specific, accurate observations from the image to support your response. "
            "Format your response in clean Markdown with headings, bold key terms, and bullet points. "
            "When relevant, use ```mermaid code blocks for diagrams (flowcharts, mind maps, etc.) or inline SVG for geometric/visual illustrations — our app renders both as interactive graphics. "
            "Use $...$ for inline math and $$...$$ for display equations (double backslashes for LaTeX commands). "
            "Be thorough, well-structured, and genuinely helpful!"
        )
        
        # Prepare messages
        messages = [
            {
                "role": "system",
                "content": system_prompt
            },
            {
                "role": "user",
                "content": message
            }
        ]
        
        # Prepare images array
        images = [[image_base64, filename]]
        
        # Call AI vision API
        try:
            result = ai_generate_vision(
                messages=messages,
                images=images,
                model=model_name or None,
            )
            
            ai_response = result.content
            
            logger.info(f"Base64 vision analysis successful for user {user_id} via {result.provider}")
            
            # Deduct tokens
            if user_id and TokenManager:
                try:
                    input_tokens = VISION_IMAGE_TOKENS + len(message.split())
                    output_tokens = len(ai_response.split())
                    
                    token_cost = TokenManager.calculate_token_cost(
                        input_tokens=input_tokens,
                        output_tokens=output_tokens,
                        model_name=model_name or 'vision-auto'
                    )
                    tokens_to_deduct = token_cost if token_cost > 0 else output_tokens + (input_tokens * 0.1)
                    
                    TokenManager.deduct_tokens(
                        user_id=user_id,
                        amount=tokens_to_deduct,
                        endpoint='/api/ai/vision/analyze-base64',
                        ai_model=model_name or 'vision-auto',
                        request_tokens=input_tokens,
                        response_tokens=output_tokens
                    )
                    logger.info(f"Deducted {tokens_to_deduct:.1f} tokens for vision analysis (base64) user {user_id}")
                except Exception as e:
                    logger.error(f"Error deducting vision tokens: {e}")
            
            return format_api_response({
                'response': ai_response,
                'model': result.model,
                'provider': result.provider,
                'timestamp': datetime.now().isoformat()
            })
            
        except Exception as vision_error:
            logger.error(f"Vision API error: {str(vision_error)}")
            return create_error_response(
                f"Vision analysis failed: {str(vision_error)}",
                500
            )
        
    except Exception as e:
        logger.error(f"Error in analyze_image_base64: {str(e)}")
        return create_error_response(f"Failed to analyze image: {str(e)}", 500)

@ai_vision_bp.route('/api/ai/vision/models', methods=['GET'])
@require_api_key
@encrypt_response
def get_vision_models():
    """
    Get list of available vision-capable models
    """
    try:
        # List of known vision-capable models
        vision_models = [
            {
                'name': 'gpt-4-vision-preview',
                'provider': 'OpenAI',
                'description': 'GPT-4 with vision capabilities',
                'supports_vision': True
            },
            {
                'name': 'claude-3-opus',
                'provider': 'Anthropic',
                'description': 'Claude 3 Opus with vision',
                'supports_vision': True
            },
            {
                'name': 'gemini-pro-vision',
                'provider': 'Google',
                'description': 'Gemini Pro Vision',
                'supports_vision': True
            },
            {
                'name': '',  # Auto-select
                'provider': 'Copilot',
                'description': 'Microsoft Copilot (auto-select best model)',
                'supports_vision': True
            }
        ]
        
        return format_api_response({
            'models': vision_models,
            'total': len(vision_models)
        })
        
    except Exception as e:
        logger.error(f"Error in get_vision_models: {str(e)}")
        return create_error_response(f"Failed to get vision models: {str(e)}", 500)

# Database migration SQL (run this if table doesn't exist)
"""
CREATE TABLE IF NOT EXISTS ai_vision_requests (
    id INT AUTO_INCREMENT PRIMARY KEY,
    user_id VARCHAR(50) NOT NULL,
    message TEXT NOT NULL,
    model_name VARCHAR(100),
    provider VARCHAR(100),
    response TEXT,
    image_size INT,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    INDEX idx_user_id (user_id),
    INDEX idx_created_at (created_at)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci;
"""
