import cv2
import numpy as np
import google.generativeai as genai
from langchain_core.messages import HumanMessage
from langchain_google_genai import ChatGoogleGenerativeAI
import os
import streamlit as st
from PIL import Image
from streamlit_drawable_canvas import st_canvas

# Set up environment variables and configurations
genai.configure(api_key=os.environ['GOOGLE_API_KEY'])

def main():
    st.title("Virtual Math Calculator")

    col1, col2 = st.columns([3, 1])

    with col1:
        # Create a drawing canvas
        st.write("**Draw your equation below:**")
        canvas_result = st_canvas(
            fill_color="rgb(255, 255, 255)",  # Fill color for drawing
            stroke_width=5,
            stroke_color="rgb(0, 0, 0)",
            background_color="rgb(255, 255, 255)",
            height=300,
            width=600,
            drawing_mode="freedraw",
            key="canvas",
        )

        if st.button("Submit"):
            if canvas_result.image_data is not None:
                # Convert to image
                drawing_image = Image.fromarray(canvas_result.image_data.astype('uint8'), 'RGBA').convert('RGB')
                drawing_image.save('drawing.png')

                # Send the image to Gemini and get the response
                response = send_to_gemini('drawing.png')
                st.session_state.result = response

    with col2:
        st.header("Instructions")
        st.write("1. Use the mouse pointer/finger to draw your equation.")
        st.write("2. Click **Submit** to process the drawing.")
        st.write("3. The result will be displayed below after submission.")
        if 'result' in st.session_state:
            st.write("**Result:**")
            st.text_area("", value=st.session_state.result, height=300)

def send_to_gemini(drawing_path):
    llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash-latest")
    # with open(drawing_path, 'rb') as img_file:
    #     image_content = img_file.read()
    message = HumanMessage(
        content=[
            {
                "type": "text",
                "text": "Give me the answer of any mathematical representation in the image with the complete solution, and does not say the image contains etc.",
            },
            {"type": "image_url", "image_url": drawing_path},
        ]
    )
    response = llm.invoke([message]).content
    return response

if __name__ == "__main__":
    main()