pages/3_Multimodal.py

import streamlit as st
import sqlite3
import time
import datetime
from PIL import Image
import google.generativeai as genai
import os
from reportlab.pdfgen import canvas
from reportlab.lib.pagesizes import A4, letter 
from io import BytesIO
import tempfile
import json
import re
from reportlab.platypus import Paragraph, Frame, Spacer
from reportlab.lib.styles import getSampleStyleSheet
import shutil

MODEL_ID = "gemini-2.0-flash-exp" 
api_key = os.getenv("GEMINI_API_KEY")
model_id = MODEL_ID
genai.configure(api_key=api_key)
enable_stream = False

if "model" not in st.session_state:
    st.session_state.model = genai.GenerativeModel(MODEL_ID)

if "chat" not in st.session_state:
    st.session_state.chat = st.session_state.model.start_chat()

if "is_new_file" not in st.session_state:
    st.session_state.is_new_file = True

def get_system_instruction(username):
    """ Retrieves the system instruction for the user from the database. """
    conn = sqlite3.connect('users.db')
    c = conn.cursor()
    c.execute('SELECT instruction FROM system_instructions WHERE username=?', (username,))
    instruction = c.fetchone()
    conn.close()
    if instruction:
        return instruction[0]
    else:
        return "Default system instruction."

def save_user_prompt(username, prompt_time, prompt_type):
    """ Saves the user prompt to the database for monitoring purposes. """

    conn = sqlite3.connect('users.db')
    c = conn.cursor()
    c.execute('INSERT INTO user_prompts(username, prompt_time, prompt_type) VALUES (?,?,?)', (username, prompt_time, prompt_type))
    conn.commit()
    conn.close()

def merge_json_strings(json_str1, json_str2):
    """
    Merges two JSON strings into one, handling potential markdown tags.

    Args:
        json_str1: The first JSON string, potentially with markdown tags.
        json_str2: The second JSON string, potentially with markdown tags.

    Returns:
        A cleaned JSON string representing the merged JSON objects.
    """

    # Clean the JSON strings by removing markdown tags
    cleaned_json_str1 = _clean_markdown(json_str1)
    cleaned_json_str2 = _clean_markdown(json_str2)

    try:
        # Parse the cleaned JSON strings into Python dictionaries
        data1 = json.loads(cleaned_json_str1)
        data2 = json.loads(cleaned_json_str2)

        # Merge the dictionaries
        merged_data = _merge_dicts(data1, data2)

        # Convert the merged dictionary back into a JSON string
        return json.dumps(merged_data, indent=2)
    except json.JSONDecodeError as e:
        return f"Error decoding JSON: {e}"


def _clean_markdown(text):
    """
    Removes markdown tags from a string if they exist. 
    Otherwise, returns the original string unchanged.

    Args:
        text: The input string.

    Returns:
        The string with markdown tags removed, or the original string 
        if no markdown tags were found.
    """
    try:
        # Check if the string contains markdown 
        if re.match(r"^```json\s*", text) and re.search(r"\s*```$", text):
            # Remove leading ```json
            text = re.sub(r"^```json\s*", "", text) 
            # Remove trailing ```
            text = re.sub(r"\s*```$", "", text) 
        return text
    except Exception as e:
        # Log the error 
        st.error(f"Error cleaning markdown: {e}") 
        return None

def _merge_dicts(data1, data2):
    """
    Recursively merges two data structures.

    Handles merging of dictionaries and lists. 
    For dictionaries, if a key exists in both and both values are dictionaries 
    or lists, they are merged recursively. Otherwise, the value from data2 is used.
    For lists, the lists are concatenated.

    Args:
        data1: The first data structure (dictionary or list).
        data2: The second data structure (dictionary or list).

    Returns:
        The merged data structure.

    Raises:
        ValueError: If the data types are not supported for merging.
    """
    if isinstance(data1, dict) and isinstance(data2, dict):
        for key, value in data2.items():
            if key in data1 and isinstance(data1[key], (dict, list)) and isinstance(value, type(data1[key])):
                _merge_dicts(data1[key], value)
            else:
                data1[key] = value
        return data1
    elif isinstance(data1, list) and isinstance(data2, list):
        return data1 + data2
    else:
        raise ValueError("Unsupported data types for merging")

def create_json(metadata, content):
    """
    Creates a JSON string combining metadata and content.

    Args:
        metadata: A dictionary containing metadata information.
        content: A dictionary containing the quiz content.

    Returns:
        A string representing the combined JSON data.
    """

    # Create metadata with timestamp
    metadata = {
        "subject": metadata.get("subject", ""),
        "topic": metadata.get("topic", ""),
        "num_questions": metadata.get("num_questions", 0),
        "exam_type": metadata.get("exam_type", ""),
        "timestamp": datetime.datetime.now().isoformat()
    }

    # Combine metadata and content
    combined_data = {"metadata": metadata, "content": content}

    # Convert to JSON string
    json_string = json.dumps(combined_data, indent=4)

    return json_string

def create_pdf(data):
    """Creates a PDF file with text wrapping for quiz content."""
    try:
        # Load the JSON data
        data = json.loads(data)

        if 'metadata' not in data or 'content' not in data:
            st.error("Error: Invalid data format. Missing 'metadata' or 'content' keys.")
            return None

        metadata = data['metadata']
        content = data['content']

        # Validate metadata
        required_metadata_keys = ['subject', 'topic', 'exam_type', 'num_questions']
        if not all(key in metadata for key in required_metadata_keys):
            st.error("Error: Invalid metadata format. Missing required keys.")
            return None

        # Create a unique filename with timestamp
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        pdf_filename = f"quiz_output_{timestamp}.pdf"

        # Get the temporary directory
        temp_dir = tempfile.gettempdir()
        pdf_path = os.path.join(temp_dir, pdf_filename)

        c = canvas.Canvas(pdf_path, pagesize=A4)
        c.setFont("Helvetica", 10)

        exam_type = metadata['exam_type']

        styles = getSampleStyleSheet()
        style_normal = styles["Normal"]

        y_position = 750
        line_height = 15
        frame_width = 500
        first_page = True

        for idx, q in enumerate(content):
            if not isinstance(q, dict):
                st.error(f"Error: Invalid question format at index {idx}. Skipping...")
                continue

            if first_page:
                # Print metadata once
                for key, label in [("subject", "Subject"), ("topic", "Topic"),
                                                ("exam_type", "Type"), ("num_questions", "Number of Questions")]:
                    c.drawString(50, y_position, f"{label}: {metadata[key]}")
                    y_position -= line_height

                y_position -= line_height  # Extra space before questions
                first_page = False

            # Print question number
            question_text = f"{idx+1}. "
            c.drawString(50, y_position, question_text)
            x_position = 70  # Adjust starting position for question text

            # --- Changes for better text flow ---
            # Split the question into words
            words = q.get('question', q.get('statement', '')).split() 
            current_line = ""
            for word in words:
                temp_line = current_line + " " + word
                text_width = c.stringWidth(temp_line, "Helvetica", 10)
                if text_width <= frame_width:
                    current_line = temp_line
                else:
                    c.drawString(x_position, y_position, current_line)
                    y_position -= line_height
                    current_line = word
                    if y_position < 50:
                        c.showPage()
                        c.setFont("Helvetica", 10)
                        y_position = 750
            # Print the last line of the question
            c.drawString(x_position, y_position, current_line)
            y_position -= line_height
            # --- End of changes ---

            if exam_type == "Multiple Choice":
                # Validate question structure
                required_question_keys = ['question', 'options', 'correct_answer']
                if not all(key in q for key in required_question_keys):
                    st.error(f"Error: Invalid question format at index {idx}. Skipping...")
                    continue

                # Print options
                for option_idx, option in enumerate(q['options'], ord('a')):
                    c.drawString(70, y_position, f"{chr(option_idx)}) {option}")
                    y_position -= line_height
                    if y_position < 50:
                        c.showPage()
                        c.setFont("Helvetica", 10)
                        y_position = 750

                # Print correct answer
                c.drawString(70, y_position, f"Correct Answer: {q['correct_answer']}")
                y_position -= line_height * 2

            elif exam_type == "True or False":
                # Validate question structure
                required_question_keys = ['statement', 'options', 'correct_answer']
                if not all(key in q for key in required_question_keys):
                    st.error(f"Error: Invalid question format at index {idx}. Skipping...")
                    continue

                # Print options
                for option in q['options']:
                    c.drawString(70, y_position, f"{option}")
                    y_position -= line_height
                    if y_position < 50:
                        c.showPage()
                        c.setFont("Helvetica", 10)
                        y_position = 750

                # Print correct answer
                c.drawString(70, y_position, f"Correct Answer: {q['correct_answer']}")
                y_position -= line_height * 2

            elif exam_type in ["Short Response", "Essay Type"]:
                # Validate question structure
                required_question_keys = ['question', 'correct_answer']
                if not all(key in q for key in required_question_keys):
                    st.error(f"Error: Invalid question format at index {idx}. Skipping...")
                    continue

                # Print correct answer 
                answer_text = f"Correct Answer: {q['correct_answer']}"
                
                # --- Changes for better text flow ---
                # Split the answer into words
                words = answer_text.split()
                current_line = ""
                for word in words:
                    temp_line = current_line + " " + word
                    text_width = c.stringWidth(temp_line, "Helvetica", 10)
                    if text_width <= frame_width:
                        current_line = temp_line
                    else:
                        c.drawString(x_position, y_position, current_line)
                        y_position -= line_height
                        current_line = word
                        if y_position < 50:
                            c.showPage()
                            c.setFont("Helvetica", 10)
                            y_position = 750
                # Print the last line of the answer
                c.drawString(x_position, y_position, current_line)
                y_position -= line_height * 2
                # --- End of changes ---

            if y_position < 50:
                c.showPage()
                c.setFont("Helvetica", 10)
                y_position = 750

        # Add the notice at the end
        notice = "This exam was generated by the WVSU Exam Maker (c) 2025 West Visayas State University"
        c.drawString(50, y_position, notice)

        c.save()
        return pdf_path

    except Exception as e:
        st.error(f"Error creating PDF: {e}")
        return None
    
def generate_quiz_content(data):
    """
    Separates the metadata and content from a JSON string containing exam data.
    Creates a markdown formatted text that contains the exam metadata and 
    enumerates the questions, options and answers nicely formatted for readability.

    Args:
      data: A JSON string containing the exam data.

    Returns:
      A markdown formatted string.
    """
    data = json.loads(data)
    metadata = data["metadata"]
    content = data["content"]
    exam_type = metadata["exam_type"]
    if exam_type == "Multiple Choice":
        md_text = f"""# {metadata['subject']} - {metadata['topic']}

**Exam Type:** {metadata['exam_type']}  
**Number of Questions:** {metadata['num_questions']}  
**Timestamp:** {metadata['timestamp']}

---

"""
        for i, q in enumerate(content):
            md_text += f"""Question {i+1}:
            {q['question']}

"""
            for j, option in enumerate(q['options'], ord('a')):
                md_text += f"""{chr(j)}. {option}  

"""
            md_text += f"""**Correct Answer:** {q['correct_answer']}

---

"""
        md_text += """This exam was generated by the WVSU Exam Maker
            (c) 2025 West Visayas State University
"""
            
    elif exam_type == "True or False":
        md_text = f"""# {metadata['subject']} - {metadata['topic']}

**Exam Type:** {metadata['exam_type']}  
**Number of Questions:** {metadata['num_questions']}  
**Timestamp:** {metadata['timestamp']}

---

"""

        for i, q in enumerate(content):
            md_text += f"""Statement {i+1}:

{q['statement']}

"""
            for j, option in enumerate(q['options'], ord('a')):
                md_text += f"""{option}  
"""

            md_text += f"""**Correct Answer:** {q['correct_answer']}

---
"""
        md_text += """This exam was generated by the WVSU Exam Maker
(c) 2025 West Visayas State University"""

    elif exam_type == "Short Response" or exam_type == "Essay Type":
        md_text = f"""# {metadata['subject']} - {metadata['topic']}

**Exam Type:** {metadata['exam_type']}  
**Number of Questions:** {metadata['num_questions']}  
**Timestamp:** {metadata['timestamp']}

---

"""

        for i, q in enumerate(content):
            md_text += f"""Question {i+1}:

{q['question']}

"""
            md_text += f"""**Correct Answer:** {q['correct_answer']}

---
"""
        md_text += """This exam was generated by the WVSU Exam Maker
(c) 2025 West Visayas State University"""

    return md_text
    
def generate_metadata(subject, topic, num_questions, exam_type):
    """Generates quiz metadata as a dictionary combining num_questions, 
    exam_type, and timestamp.

    Args:
        num_questions: The number of questions in the exam (int).
        exam_type: The type of exam (str).

    Returns:
        A dictionary containing the quiz metadata.
    """

    # Format the timestamp
    timestamp = datetime.datetime.now()
    formatted_timestamp = timestamp.strftime("%Y-%m-%d %H:%M:%S") 
    
    metadata = {
        "subject": subject,
        "topic": topic,
        "num_questions": num_questions,
        "exam_type": exam_type,
        "timestamp": formatted_timestamp
    }

    return metadata

def generate_text(uploaded_file, mime_type, prompt):
    """Generates text based on the uploaded file and prompt."""
    try:
        if st.session_state.is_new_file:
            # Upload the file with the correct MIME type
            file_data = genai.upload_file(uploaded_file, mime_type=mime_type)

            # Send file and prompt to Gemini API
            chat = st.session_state.chat
            response = chat.send_message(
                [
                    prompt, 
                    file_data
                ],
                stream=enable_stream  
            )
            st.session_state.is_new_file = False
        else:
            # continue chat without sending the file again
            # Send a text prompt to Gemini API
            chat = st.session_state.chat
            response = chat.send_message(
                [
                    prompt
                ],
                stream=enable_stream  
            )

        return response.text
    
    except Exception as e:
        st.error(f"An error occurred while generating text: {e}")
        return None

def show_multimodal():
    st.subheader("Multimodal")
    username = st.session_state["username"]
    st.write(f"Welcome, {username}! This page allows you to generate questions based on an image or PDF file.")

    # Display username and logout button on every page
    st.sidebar.write(f"Current user: {st.session_state['username']}")
    
    # we dont use the system instruction for now
    #system_instruction = get_system_instruction(username)

    # File uploader with allowed types
    uploaded_file = st.file_uploader("Choose an image or PDF...", type=["jpg", "jpeg", "png", "pdf"])

    if uploaded_file is not None:
        # Determine file type
        file_type = uploaded_file.type
        if file_type.startswith('image'):
            # Display the uploaded image
            image = Image.open(uploaded_file)
            st.image(image, caption="Uploaded Image.", use_container_width=True)
            mime_type = "image/jpeg"  # Use a consistent MIME type for images
            # Display a message for PDF upload
            st.write("Image file was uploaded. Questions will be generated based on its contents.")
        elif file_type == 'application/pdf':
            # Display a message for PDF upload
            st.write("PDF file uploaded. Questions will be generated based on its contents.")
            mime_type = "application/pdf"
        else:
            st.error("Unsupported file type. Please upload an image or PDF.")
            st.stop()

    # User inputs
    # Course selection
    course = st.selectbox("Select Course", 
                        ["Diploma in Teaching", 
                        "Post Baccalaureate Diploma in Early Childhood Education", 
                        "Master of Arts in Education - Language Teaching (English)", 
                        "Master in Education major in Early Childhood Education"])

    # Year level selection
    year_level = st.selectbox("Select Year Level", 
                            ["1st Year", 
                            "2nd Year", 
                            "3rd Year", 
                            "4th Year"])

    # Subject selection
    subject = st.text_input("Enter Subject", 
                        "e.g.,The Teaching Profession, Facilitating Learner-Centered Teaching")

    # Topic selection
    topic = st.text_input("Enter Topic", 
                        "e.g., Teacher as a professional, Introduction to Learner-Centered Teaching")

    # Question type selection
    question_type = st.selectbox("Select Question Type", 
                                ["Multiple Choice", 
                                "True or False", 
                                "Short Response", 
                                "Essay Type"])

    difficulty = st.selectbox("Select Difficulty",["easy","average","hard"])

    #number of questions to generate
    if question_type != "Essay Type":
        num_questions = st.selectbox("Number of Questions to Generate",
                                [10, 20, 30, 40, 50])  
    else:
        num_questions = st.selectbox("Number of Questions to Generate",
                                [1, 2, 3, 4, 5]) 

    # Combine user inputs into a prompt
    prompt = f"""Refer to the uploaded document. Generate a {question_type} question for a {year_level} {course} student
                    in {subject} on the topic of {topic} with a {difficulty} difficulty level.
                    The questions should require higher order thinking skills. 
                    """

    if question_type == "Multiple Choice":    
        prompt += """Provide 4 choices. Provide the correct answer in the format 'Answer: A'.
        Use the following JSON format for each question:
        [{
        "question": "Your question here?",
        "options": ["Option A", "Option B", "Option C", "Option D"],
        "correct_answer": "full text of the correct answer"
        }, ... more questions]
        Ensure that the response only contains the JSON array of questions and nothing else.
        """
    elif question_type == "True or False":
        prompt += """Indicate whether the statement is true or false. Keep the statement brief and concise.
                    Use the following JSON format for each question:
                    [{
                    "statement": "Your statement here",
                    "options": ["True", "False"],
                    "correct_answer": True"
                    }, ... more questions]
                    Ensure that the response only contains the JSON array of questions and nothing else.    
                    """   
    elif question_type == "Short Response":
        prompt += """Create question that require a word or short phrase as answer.  Use the following JSON format for each question:
                    [{
                    "question": "Your question here?",
                    "correct_answer": A word or phrase"
                    },  ... more questions]
                    Ensure that the response only contains the JSON array of questions and nothing else.
                    """
    elif question_type == "Essay Type":
        prompt += """Create questions that require a short essay between 300 to 500 words. 
                    Provide a detailed answer. Use the following JSON format for each question:
                    [{
                    "question": "Your question here?",
                    "correct_answer": The essay answer goes here."
                    }, ... more questions]
                    Ensure that the response only contains the JSON array of questions and nothing else.
                    """

    if not question_type == "Essay Type":
        prompt += f"Generate 10 questions. Do not repeat questions you have already given in previous prompts. Exclude markdown tags in the response."
    else:
        prompt += f" Generate {num_questions} questions. Do not repeat questions you have already given in previous prompts. Exclude markdown tags in the response"  

    full_quiz = ""    
    
    # Send button
    if st.button("Generate Questions"):
        

        if not uploaded_file:
            
            st.warning("Please upload an image or PDF and enter a prompt.")
            st.stop()               
        else:
            if question_type == "Essay Type":
                #prompt once
                with st.spinner('Generating questions...'):
                    full_quiz = _clean_markdown(generate_text(uploaded_file, mime_type, prompt))  

            else:
                if num_questions == 10:

                    #prompt once
                    with st.spinner('Generating questions...'):
                        full_quiz = _clean_markdown(generate_text(uploaded_file, mime_type, prompt))
                else:
                    #prompt multiple times
                    times = num_questions//10
                    for i in range(times):
                        with st.spinner('Generating questions...'):
                            response = generate_text(uploaded_file, mime_type, prompt)  

                            if i==0:
                                full_quiz = _clean_markdown(response)
                            else:                        
                                full_quiz = merge_json_strings(full_quiz, response) 

        metadata = generate_metadata(subject, topic, num_questions, question_type)

        try:
            # Attempt to load the string as JSON to validate it            
            content = json.loads(full_quiz)
        except json.JSONDecodeError:
            st.error("Error: Invalid JSON string for quiz content.")
            st.stop()

        json_string = create_json(metadata, content)

        quiz_markdown = generate_quiz_content(json_string)
        st.markdown(quiz_markdown)

        pdf_path = create_pdf(json_string)
        
        if pdf_path:
            """Click the button to download the generated PDF."""
            try:
                with open(pdf_path, "rb") as f:
                    st.download_button("Download PDF", f, file_name=os.path.basename(pdf_path))
            except Exception as e:
                st.error(f"Error handling file download: {e}")
        else:
            st.error("Failed to generate the PDF. Please try again.")

        #record the prompt for monitoring
        save_user_prompt(username, datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "Multimodal")

if st.session_state["authenticated"]:
    show_multimodal()   
else:
    if not st.session_state["is_starting"]:
        st.write("You are not authenticated. Please log in to access this page.")