import io
import base64
import requests
from PIL import Image

class SnapXamMathOCR:
    def __init__(self, api_key=None, base_url="https://api.snapxam.com/v1/ocr/recognize"):
        """
        Initialize the MathOCR client.

        :param api_key: Required API key for authentication.
        :param base_url: Endpoint of your OCR API.
        """
        self.api_key = api_key
        self.base_url = base_url

    def recognize(self, image, image_format="PNG"):
        """
        Recognize math expressions from a PIL image.

        :param image: PIL.Image object
        :return: JSON response from API
        """
        if not isinstance(image, Image.Image):
            raise TypeError("Expected a PIL.Image.Image instance")

        # Convert PIL image to bytes
        buffer = io.BytesIO()
        image.save(buffer, format=image_format)
        img_bytes = buffer.getvalue()

        # Encode to Base64
        encoded = base64.b64encode(img_bytes).decode("utf-8")

        # Prepare payload
        data = {"image_data": encoded}
        headers = {}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"  # optional

        # Send request
        response = requests.post(self.base_url, data=data, headers=headers)
        response.raise_for_status()

        return response.json()
