import os
from flask import Blueprint, render_template, request, jsonify
import base64
from io import BytesIO
from PIL import Image
import numpy as np
import cv2
from utils.math_solver import solve_equation
from typing import Union, Tuple, Any

# Initialize Pix2Text with MFD (Mathematical Formula Detection) model for better accuracy
try:
    from pix2text import Pix2Text
    p2t = Pix2Text(analyzer_config=dict(model_name='mfd'))
except Exception as e:
    print(f"Warning: Could not initialize Pix2Text: {e}")
    p2t = None

scribble_bp = Blueprint('scribble_bp', __name__)

UPLOAD_FOLDER = 'static/uploads'
PROCESSED_FOLDER = 'static/processed'
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(PROCESSED_FOLDER, exist_ok=True)


def preprocess_image(image_data):
    """Preprocess image for better OCR results"""
    try:
        # Convert base64 to image
        image_data = image_data.split(',')[1]  # Remove data URL prefix
        image = Image.open(BytesIO(base64.b64decode(image_data)))
        
        # Convert to OpenCV format
        opencv_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
        
        # Convert to grayscale
        gray = cv2.cvtColor(opencv_image, cv2.COLOR_BGR2GRAY)
        
        # Apply mild Gaussian blur to reduce noise while preserving edges
        blurred = cv2.GaussianBlur(gray, (3, 3), 0)
        
        # Apply adaptive thresholding with parameters better suited for handwritten text
        thresh = cv2.adaptiveThreshold(
            blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 15, 2
        )
        
        # Save processed image
        processed_path = os.path.join(PROCESSED_FOLDER, 'scribble_processed.png')
        cv2.imwrite(processed_path, thresh)
        
        return processed_path
    except Exception as e:
        print(f"Preprocessing error: {e}")
        # Save original if preprocessing fails
        image_path = os.path.join(UPLOAD_FOLDER, 'scribble.png')
        # Reopen image to save it
        image = Image.open(BytesIO(base64.b64decode(image_data.split(',')[1])))
        image.save(image_path)
        return image_path


@scribble_bp.route("/scribble")
def scribble_page():
    return render_template("scribble.html")


@scribble_bp.route("/scribble/process", methods=["POST"])
def process_scribble():
    try:
        data = request.get_json()
        if not data:
            return jsonify({'error': 'No data provided'}), 400
        
        image_data = data.get('image', '')
        if not image_data:
            return jsonify({'error': 'No image data provided'}), 400
        
        # Save original image first
        image_path = os.path.join(UPLOAD_FOLDER, 'scribble_original.png')
        image_data_content = image_data.split(',')[1]
        image = Image.open(BytesIO(base64.b64decode(image_data_content)))
        image.save(image_path)
        
        # Process with Pix2Text if available
        if p2t:
            print(f"Processing scribble with MFD model: {image_path}")
            
            # Try with original image first (works better for handwritten math)
            result = p2t.recognize(image_path)
            print(f"Original image result: {result}")
            
            # Handle different result types
            if isinstance(result, dict):
                latex_code = result.get('text', '')
            elif isinstance(result, list):
                # If result is a list, extract text from first item
                if result and isinstance(result[0], dict):
                    latex_code = result[0].get('text', '')
                else:
                    latex_code = str(result)
            else:
                latex_code = str(result)
            
            # If we get no result or very short result, try with preprocessing
            if len(latex_code.strip()) < 2:
                print("Result too short, trying with preprocessing...")
                processed_path = preprocess_image(image_data)
                result = p2t.recognize(processed_path)
                print(f"Preprocessed image result: {result}")
                
                if isinstance(result, dict):
                    latex_code = result.get('text', '')
                elif isinstance(result, list):
                    if result and isinstance(result[0], dict):
                        latex_code = result[0].get('text', '')
                    else:
                        latex_code = str(result)
                else:
                    latex_code = str(result)
            
            print(f"Final extracted LaTeX: {latex_code}")
        else:
            latex_code = "\\text{Pix2Text not available}"
        
        return jsonify({
            'success': True,
            'latex': latex_code
        })
        
    except Exception as e:
        print(f"Error processing scribble: {e}")
        return jsonify({'error': str(e)}), 500
    return jsonify({'error': 'Unknown error'}), 500


@scribble_bp.route("/scribble/solve", methods=["POST"])
def solve_scribble_equation():
    """Solve a LaTeX equation from scribble"""
    try:
        data = request.get_json()
        if not data:
            return jsonify({'error': 'No data provided'}), 400
        
        latex_equation = data.get('latex', '')
        if not latex_equation:
            return jsonify({'error': 'No equation provided'}), 400
        
        # Solve the equation
        solution = solve_equation(latex_equation)
        
        return jsonify({
            'success': True,
            'solution': solution
        })
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500
    return jsonify({'error': 'Unknown error'}), 500