import streamlit as st
import base64
# from openai import OpenAI
from langchain_core.messages import HumanMessage
from langchain_google_genai import ChatGoogleGenerativeAI

# nFunction to convert HTML to Markdown
def html_to_markdown(data):
    # Replace header tags
    data = data.replace('<h2><b>', '## ').replace('</b></h2>', '')
    # Replace bold tags
    data = data.replace('<b>', '**').replace('</b>', '**')
    # Remove center tags (Markdown doesn't support center alignment)
    data = data.replace('<center>', '').replace('</center>', '')
    
    return data

# Function to encode the image to base64
def encode_image(image_file):
    return base64.b64encode(image_file.getvalue()).decode("utf-8")

st.set_page_config(page_title="Scientific and Engineering Image Analyst", layout="centered", initial_sidebar_state="collapsed")

# Streamlit page setup
st.title("Scientific and Engineering Image Analyst X Omnisys")

# Text input for the user to enter their OpenAI API Key
api_key = st.text_input("Enter your OpenAI API Key:", type="password")

# Initialize the OpenAI client with the API key
if api_key:
    client = ChatGoogleGenerativeAI(model="gemini-pro-vision",google_api_key = api_key)

# File uploader allows user to add their own image
uploaded_file = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])

# Checkbox to add an additional prompt
add_prompt = st.checkbox("Add additional prompt instructions")

# Initialize a variable for additional prompt text
additional_prompt_text = ""

# Conditional text input for additional prompt
if add_prompt:
    additional_prompt_text = st.text_area("Enter additional prompt instructions:")

if uploaded_file:
    # Display the uploaded image
    with st.expander("Image", expanded=True):
        st.image(uploaded_file, caption=uploaded_file.name, use_column_width=True)


# Toggle for showing additional details input
show_details = st.checkbox("Add details about the image", value=False)

if show_details:
    # Text input for additional details about the image, shown only if toggle is True
    additional_details = st.text_area(
        "Add any additional details or context about the image here:",
        disabled=not show_details
    )

# Button to trigger the analysis
analyze_button = st.button("Analyse the Image")

# Check if an image has been uploaded, if the API key is available, and if the button has been pressed
if uploaded_file is not None and api_key and analyze_button:

    with st.spinner("Analysing the image ..."):
        # Encode the image
        base64_image = encode_image(uploaded_file)

        # Standard prompt for image analysis
        prompt_text = (
            "As an expert in scientific and engineering diagram analysis, your keen eye for detail is crucial. "
            "Your primary task is to conduct a meticulous examination of the provided image. "
            "Focus on identifying every numerical value visible in the diagram, such as dimensions, tolerances, and material properties. "
            "Offer a detailed, fact-based, and technically precise explanation of the diagram, with an emphasis on the scientific or engineering principles it illustrates. "
            "Highlight the significance of each numerical value, explaining how they affect the diagram's functionality and design. "
            "Structure your analysis in a clear, markdown format, targeting an audience with a background in science or engineering. "
            "Incorporate appropriate scientific or engineering terminology to provide a thorough understanding of the numerical details. "
            "Conclude with a bold, concise caption summarizing the key aspects and numerical details of the image, and their relevance in the diagram's context."
        )



        # Append additional prompt text if provided
        if additional_prompt_text:
            prompt_text += f"\n\nAdditional Prompt Instructions:\n{additional_prompt_text}"

        # Append additional details if provided
        if show_details and additional_details:
            prompt_text += f"\n\nAdditional Context Provided by the User:\n{additional_details}"

        # Create the payload for the completion request
        messages = HumanMessage(
            content = [
                {"type": "text", "text": prompt_text},
                {
                      "type": "image_url",
                    "image_url": f"data:image/jpeg;base64,{base64_image}",
                },
            ]            
        )

        # Make the request to the OpenAI API
        try:
            # Stream the response
            full_response = html_to_markdown(client.invoke([messages]).content)
            message_placeholder = st.empty()

            message_placeholder.markdown(full_response)
            
        except Exception as e:
            st.error(f"An error occurred: {e}")
else:
    # Warnings for user action required
    if not uploaded_file and analyze_button:
        st.warning("Please upload an image.")
    if not api_key:
        st.warning("Please enter your API key.")