from langchain_google_genai import ChatGoogleGenerativeAI
import os
import base64
from langchain_core.messages import HumanMessage
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI

load_dotenv()

class ImageQuestionAnswer:
    def __init__(self):
        '''self.vision_llm = ChatGoogleGenerativeAI(
            model="gemini-2.0-flash", 
            temperature=0.1,
            api_key=os.getenv("GEMINI_API_KEY")
        )'''
        self.vision_llm = ChatOpenAI(model="gpt-4o", openai_api_key=os.getenv("OPENAI_API_KEY"))

    def answer(self, image_path: str, question: str) -> str:
        print(f"Sending image to OpenAI: {image_path}")
        print(f"Question: {question}")
        with open(image_path, "rb") as image_file:
            image_bytes = image_file.read()

        image_base64 = base64.b64encode(image_bytes).decode("utf-8")

        # Prepare the prompt including the base64 image data
        message = [
            HumanMessage(
                content=[
                    {
                        "type": "text",
                        "text": (
                            "Answer the question based on the image. "
                            f"The question is: {question}"
                        ),
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/png;base64,{image_base64}"
                        },
                    },
                ]
            )
        ]

        # Call the vision-capable model
        response = self.vision_llm.invoke(message)
        print(f"Image question answer: {response.content}")

        return response.content
    
'''if __name__ == "__main__":
    image_question_answer = ImageQuestionAnswer()
    image_question_answer.answer("cca530fc-4052-43b2-b130-b30968d8aa44.png", "Review the chess position provided in the image. It is black's turn. Provide the correct next move for black which guarantees a win. Please provide your response in algebraic notation.")'''