import cv2
import numpy as np

def enhance_image_for_math_ocr(image_bytes: bytes) -> bytes:
    """
    Upscales and binarizes an image to improve OCR accuracy for math equations.
    V5.13.14: Added OpenCV Pre-processing layer as requested.
    """
    try:
        # Decode the image bytes to OpenCV format
        nparr = np.frombuffer(image_bytes, np.uint8)
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

        if img is None:
            return image_bytes # Fallback if decode fails

        # 1. Upscale 2x using Cubic Interpolation for smoother edges
        img = cv2.resize(img, None, fx=2, fy=2, interpolation=cv2.INTER_CUBIC)

        # 2. Convert to Grayscale
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        # 3. Apply slight Gaussian Blur to remove noise before thresholding
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)

        # 4. Adaptive Thresholding - turns paper to pure white and ink to pure black
        binary = cv2.adaptiveThreshold(
            blurred, 
            255, 
            cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
            cv2.THRESH_BINARY, 
            11, 2
        )

        # Encode back to JPEG bytes (using .jpg as requested)
        success, encoded_image = cv2.imencode('.jpg', binary)
        if success:
            return encoded_image.tobytes()
        return image_bytes

    except Exception as e:
        print(f"⚠️ [IMAGE ENHANCEMENT FAILED]: {e}. Using original image.")
        return image_bytes