# import os
# import base64
# import zipfile
# from pathlib import Path
# import streamlit as st
# from byaldi import RAGMultiModalModel
# from openai import OpenAI

# import os
# st.write("Current Working Directory:", os.getcwd())

# # Function to unzip a folder if it does not exist
# # def unzip_folder_if_not_exist(zip_path, extract_to):
# #     if not os.path.exists(extract_to):
# #         with zipfile.ZipFile(zip_path, 'r') as zip_ref:
# #             zip_ref.extractall(extract_to)

# # # Example usage
# # zip_path = 'medical_index.zip'
# # extract_to = 'medical_index'
# # unzip_folder_if_not_exist(zip_path, extract_to)

# # Preload the RAGMultiModalModel
# @st.cache_resource
# def load_model():
#     return RAGMultiModalModel.from_index("./medical_index")

# RAG = load_model()

# # OpenAI API key from environment
# api_key = os.getenv("OPENAI_API_KEY")
# client = OpenAI(api_key=api_key)

# # Streamlit UI
# st.title("Medical Diagnostic Assistant")
# st.write("Enter a medical query and get diagnostic recommendations along with visual references.")

# # User input
# query = st.text_input("Query", "What should be the appropriate diagnostic test for peptic ulcer?")

# if st.button("Submit"):
#     if query:
#         # Search using RAG model
#         with st.spinner('Retrieving information...'):
#             try:
#                 returned_page = RAG.search(query, k=1)[0].base64

#                 # Decode and display the retrieved image
#                 image_bytes = base64.b64decode(returned_page)
#                 filename = 'retrieved_image.jpg'
#                 with open(filename, 'wb') as f:
#                     f.write(image_bytes)

#                 # Display image in Streamlit
#                 st.image(filename, caption="Reference Image", use_column_width=True)

#                 # Get model response
#                 response = client.chat.completions.create(
#                     model="gpt-4o-mini-2024-07-18",
#                     messages=[
#                         {"role": "system", "content": "You are a helpful assistant. You only answer the question based on the provided image"},
#                         {
#                             "role": "user",
#                             "content": [
#                                 {"type": "text", "text": query},
#                                 {
#                                     "type": "image_url",
#                                     "image_url": {"url": f"data:image/jpeg;base64,{returned_page}"},
#                                 },
#                             ],
#                         },
#                     ],
#                     max_tokens=300,
#                 )
                
#                 # Display the response
#                 st.success("Model Response:")
#                 st.write(response.choices[0].message.content)
#             except Exception as e:
#                 st.error(f"An error occurred: {e}")
#     else:
#         st.warning("Please enter a query.")


import os
import base64
import zipfile
from pathlib import Path
import streamlit as st
from byaldi import RAGMultiModalModel
from openai import OpenAI

# Preload the RAGMultiModalModel
@st.cache_resource
def load_model():
    return RAGMultiModalModel.from_index("/home/user/app/medical_index1")

RAG = load_model()

# OpenAI API key from environment
api_key = os.getenv("OPENAI_API_KEY")
client = OpenAI(api_key=api_key)

# Streamlit UI
st.title("Medical Diagnostic Assistant")
st.write("Enter a medical query and get diagnostic recommendations along with visual references.")

# User input for selecting the model
model_options = ["gpt-4o", "gpt-4o-mini", "o1-preview", "o1-mini"]
selected_model = st.selectbox("Choose a GPT model", model_options)

# User input for query
query = st.text_input("Query", "What should be the appropriate diagnostic test for peptic ulcer?")

if st.button("Submit"):
    if query:
        # Search using RAG model
        with st.spinner('Retrieving information...'):
            try:
                # Get top 10 images
                returned_pages = RAG.search(query, k=10)
                image_urls = []
                for i, page in enumerate(returned_pages):
                    image_bytes = base64.b64decode(page.base64)
                    filename = f'retrieved_image_{i}.jpg'
                    with open(filename, 'wb') as f:
                        f.write(image_bytes)
                    image_urls.append(f"data:image/jpeg;base64,{page.base64}")

                # Get model response
                response = client.chat.completions.create(
                    model=selected_model,
                    messages=[
                        {"role": "system", "content": "You are a helpful assistant. You only answer the question based on the provided images."},
                        {
                            "role": "user",
                            "content": [
                                {"type": "text", "text": query},
                                *[{"type": "image_url", "image_url": {"url": url}} for url in image_urls],
                            ],
                        },
                    ],
                    max_tokens=300,
                )

                # Display the response
                st.success("Model Response:")
                st.write(response.choices[0].message.content)

                # Option to see all references
                # # Option to see all references
                # if st.button("Show References"):
                #     st.subheader("References")
                #     for i, page in enumerate(returned_pages):
                #         st.image(f'retrieved_image_{i}.jpg', caption=f"Reference Image {i+1}", use_column_width=True)
            except Exception as e:
                st.error(f"An error occurred: {e}")
    else:
        st.warning("Please enter a query.")