import streamlit as st
import os
import google.generativeai as genai
import json
from PIL import Image


MODEL_ID = "gemini-2.0-flash-exp"  # Keep the model ID as is
try:
    api_key = os.getenv("GEMINI_API_KEY")
    model_id = MODEL_ID
    genai.configure(api_key=api_key)
except Exception as e:
    st.error(f"Error: {e}")
    st.stop()

model = genai.GenerativeModel(MODEL_ID)
chat = model.start_chat()

def get_local_pdf_path():
    """
    Returns the path to the local PDF file.
    """
    try:
        pdf_path = os.path.join("problems", "problems.pdf")
        if not os.path.exists(pdf_path):
            raise FileNotFoundError(f"{pdf_path} does not exist.")
        return pdf_path
    except Exception as e:
        st.error(f"Failed to find the local PDF: {e}")
        st.stop()  # Stop if the file is not found

# Initialize conversation history in Streamlit session state
if "conversation_history" not in st.session_state:
    st.session_state.conversation_history = []
if "uploaded_file_part" not in st.session_state:  # Store the file *part*
    st.session_state.uploaded_file_part = None
if "uploaded_pdf_path" not in st.session_state:
    st.session_state.uploaded_pdf_path = get_local_pdf_path()

def multimodal_prompt(pdf_path, text_prompt, file_type="PDF"):
    """
    Sends a multimodal prompt to Gemini, handling file uploads efficiently.
    Args:
        pdf_path: The path to the file (PDF or image).
        text_prompt: The text prompt for the model.
        file_type: "PDF" or "image" to specify the file type.
    Returns:
        The model's response as a string, or an error message.
    """
    try:
        if file_type == "PDF":
            mime_type = "application/pdf"
        elif file_type == "image":
            import mimetypes
            mime_type, _ = mimetypes.guess_type(pdf_path)
            if mime_type is None:
              return "Could not determine MIME type for image. Please check the file path or type."
        else:
            return "Invalid file_type. Must be 'PDF' or 'image'."

        if st.session_state.get("uploaded_file_part") is None:  # First time, upload
            pdf_part = genai.upload_file(pdf_path, mime_type=mime_type)
            st.session_state.uploaded_file_part = pdf_part
            prompt = [text_prompt, pdf_part]  # First turn includes the actual file
        else:  # Subsequent turns, reference the file
            prompt = [text_prompt, st.session_state.uploaded_file_part]  # Subsequent turns include the file reference

        response = chat.send_message(prompt)

        # Update conversation history
        st.session_state.conversation_history.append({"role": "user", "content": text_prompt, "has_file": True})
        st.session_state.conversation_history.append({"role": "assistant", "content": response.text})
        return response.text

    except Exception as e:
        return f"An error occurred: {e}"


# --- Main Page ---
st.title("📚❓Problem Solving Tutor")
about = """
**How to use this App**
Replace this placeholder with the actual text.
"""


import re
import json

# Define constants
TEXT_PROMPT = """Use the provided document. Read the list of 5 quadratic equations. 
Return your response as a JSON list. Do not include any extra text, explanations, or backslashes.

Example JSON output:
[
  "x^2 - 5x + 6 = 0",
  "2x^2 + 3x - 1 = 0",
  "x^2 - 9 = 0",
  "3x^2 - 2x + 4 = 0",
  "x^2 + 8x + 15 = 0"
]
"""

# Main code
with st.spinner("AI is thinking..."):
    if st.session_state.get("uploaded_pdf_path") is None:
        st.session_state.uploaded_pdf_path = get_local_pdf_path()

    filepath = st.session_state.uploaded_pdf_path
    response = multimodal_prompt(filepath, TEXT_PROMPT, file_type="PDF")

# Remove the ```json and ``` and extra spaces.
try:
    json_string = response.replace('```json', '').replace('```', '').strip()

    # Parse the JSON string into a Python list.
    problems_list = json.loads(json_string)

    # Iterate over the list and print each item using st.write().
    for item in problems_list:
        st.write(item)

except json.JSONDecodeError:
    st.write("Error: Invalid JSON format in the response.")
except Exception as e:
    st.write(f"An unexpected error occurred: {e}")


st.markdown("Visit our Hugging Face Space!")
st.markdown("© 2025 WVSU AI Dev Team 🤖 ✨")