import base64
import mimetypes
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage

vision_llm = ChatOpenAI(model="gpt-4o")

def extract_text(image_path: str) -> str:
    # Read the image
    with open(image_path, "rb") as file:
        image_bytes = file.read()

    # Convert to Base64
    encoded = base64.b64encode(
        image_bytes
    ).decode("utf-8")

    # Determine image MIME type
    mime_type, _ = mimetypes.guess_type(image_path)

    if mime_type not in ("image/png", "image/jpeg", "image/webp"):
        raise ValueError("Unsupported image format")

    # Create a multimodal message
    message = HumanMessage(
        content=[
            {
                "type": "text",
                "text": "Extract all text from this image. Return only the extracted text."
            },
            {
                "type": "image_url",
                "image_url": {
                    "url": (
                        f"data:{mime_type};base64,{encoded}"
                    )
                }
            }
        ]
    )

    # Call the vision model
    response = vision_llm.invoke([message])

    return str(response.content)