import os
import streamlit as st
from dotenv import load_dotenv
from langchain.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain.schema import StrOutputParser
from docx import Document
import fitz  # PyMuPDF

def extract_text_from_pdf_or_docx(file):
    """Extract text from PDF or Word document."""
    filename = file.name
    text = ""
    if filename.endswith('.pdf'):
        # Extract text from PDF
        with fitz.open(file) as doc:
            for page in doc:
                text += page.get_text()
    elif filename.endswith('.docx'):
        # Extract text from Word document
        doc = Document(file)
        for paragraph in doc.paragraphs:
            text += paragraph.text + "\n"
    else:
        text = "Unsupported file format. Please upload a PDF or Word document."
    return text

def create_multiple_choice_prompt(num_questions, quiz_context, expertise):
    """Create the prompt template for multiple-choice quiz."""
    template = f"""
You are an expert in {expertise}. Generate a quiz with {num_questions} multiple-choice questions that are relevant to {expertise} based on the following content: {quiz_context}.

The questions should be at the level of {expertise} and should challenge the knowledge of someone proficient in this field.

The format of the quiz is as follows:
- Multiple-choice:
- Questions:
    1. <Question1>:
        a. Answer 1
        b. Answer 2
        c. Answer 3
        d. Answer 4

    2. <Question2>:
        a. Answer 1
        b. Answer 2
        c. Answer 3
        d. Answer 4
    ....
- Answers:
    1. <a|b|c|d>
    2. <a|b|c|d>
    ....
Example:
- Questions:
    1. What is the time complexity of a binary search tree?
        a. O(n)
        b. O(log n)
        c. O(n^2)
        d. O(1)
- Answers:
    1. b
"""
    return template

def create_true_false_prompt(num_questions, quiz_context, expertise):
    """Create the prompt template for true-false quiz."""
    template = f"""
You are an expert in {expertise}. Generate a quiz with {num_questions} true-false questions that are relevant to {expertise} based on the following content: {quiz_context}.

The questions should be at the level of {expertise} and should challenge the knowledge of someone proficient in this field.

The format of the quiz is as follows:
- True-false:
- Questions:
    1. <Question1>: <True|False>
    2. <Question2>: <True|False>
    .....
- Answers:
    1. <True|False>
    2. <True|False>
    .....
Example:
- Questions:
    1. A binary search tree is a type of data structure.
    2. Binary search trees are typically used for sorting and searching operations.
- Answers:
    1. True
    2. True
"""
    return template

def create_open_ended_prompt(num_questions, quiz_context, expertise):
    """Create the prompt template for open-ended quiz."""
    template = f"""
You are an expert in {expertise}. Generate a quiz with {num_questions} open-ended questions that are relevant to {expertise} based on the following content: {quiz_context}.

The questions should be at the level of {expertise} and should challenge the knowledge of someone proficient in this field.

The format of the quiz is as follows:
- Open-ended:
- Questions:
    1. <Question1>
    2. <Question2>
    ....
Example:
- Questions:
    1. What is a binary search tree?
    2. How are binary search trees implemented?
"""
    return template

def create_fill_in_the_blank_prompt(num_questions, quiz_context, expertise):
    """Create the prompt template for fill-in-the-blank quiz."""
    template = f"""
You are an expert in {expertise}. Generate a quiz with {num_questions} fill-in-the-blank questions that are relevant to {expertise} based on the following content: {quiz_context}.

The questions should be at the level of {expertise} and should challenge the knowledge of someone proficient in this field.

The format of the quiz is as follows:
- Fill-in-the-blank:
- Questions:
    1. <Question1>: <Fill-in-the-blank>
    2. <Question2>: <Fill-in-the-blank>
    ....
Example:
- Questions:
    1. A binary search tree is a ________ data structure.
    2. Binary search trees are implemented using ________.
- Answers:
    1. hierarchical
    2. linked lists
"""
    return template

def create_mixed_questions_prompt(num_questions, quiz_context, expertise):
    """Create the prompt template for a mix of all question types."""
    template = f"""
You are an expert in {expertise}. Generate a quiz with {num_questions} questions that include a mix of multiple-choice, true-false, open-ended, and fill-in-the-blank questions relevant to {expertise} based on the following content: {quiz_context}.

The questions should be at the level of {expertise} and should challenge the knowledge of someone proficient in this field.

The format of the quiz is as follows:
- Mixed Questions:
- Questions:
    1. <Question1> (Multiple-choice):
        a. Answer 1
        b. Answer 2
        c. Answer 3
        d. Answer 4

    2. <Question2> (True/False):
        <Question2>: <True|False>

    3. <Question3> (Open-ended):
        <Question3>

    4. <Question4> (Fill-in-the-blank):
        <Question4>: <Fill-in-the-blank>
    ....
- Answers:
    1. <a|b|c|d>
    2. <True|False>
    3. <Open-ended Answer>
    4. <Fill-in-the-blank Answer>
    ....
Example:
- Questions:
    1. What is the time complexity of a binary search tree? (Multiple-choice)
        a. O(n)
        b. O(log n)
        c. O(n^2)
        d. O(1)
    2. A binary search tree is a type of data structure. (True/False)
        True
    3. What is a binary search tree? (Open-ended)
    4. A binary search tree is a ________ data structure. (Fill-in-the-blank)
- Answers:
    1. b
    2. True
    3. A binary search tree is a data structure used to store data in a sorted manner.
    4. hierarchical
"""
    return template

def create_quiz_chain(openai_api_key):
    """Creates the chain for the quiz app."""
    llm = ChatOpenAI(temperature=0.0, openai_api_key=openai_api_key)
    return llm | StrOutputParser()

def split_questions_answers(quiz_response):
    """Function that splits the questions and answers from the quiz response."""
    if "Answers:" in quiz_response:
        questions = quiz_response.split("Answers:")[0]
        answers = quiz_response.split("Answers:")[1]
    else:
        questions = quiz_response
        answers = "Answers section not found in the response."
    return questions, answers

def main():
    st.title("QuesPro")
    st.write("This app generates questions based on the uploaded document.")
    load_dotenv()
    openai_api_key = os.getenv("OPENAI_API_KEY")
    
    uploaded_file = st.file_uploader("Upload a PDF or Word document")
    if uploaded_file is not None:
        text = extract_text_from_pdf_or_docx(uploaded_file)
        num_questions = st.number_input("Enter the number of questions", min_value=1, max_value=10, value=3)
        quiz_type = st.selectbox("Select the quiz type", ["multiple-choice", "true-false", "open-ended", "fill-in-the-blank", "mixed"])
        expertise = st.text_input("Enter the field of expertise for the quiz")

        if st.button("Generate Questions"):
            if quiz_type == "multiple-choice":
                prompt_template = create_multiple_choice_prompt(num_questions, text, expertise)
            elif quiz_type == "true-false":
                prompt_template = create_true_false_prompt(num_questions, text, expertise)
            elif quiz_type == "open-ended":
                prompt_template = create_open_ended_prompt(num_questions, text, expertise)
            elif quiz_type == "fill-in-the-blank":
                prompt_template = create_fill_in_the_blank_prompt(num_questions, text, expertise)
            else:  # mixed
                prompt_template = create_mixed_questions_prompt(num_questions, text, expertise)

            chain = create_quiz_chain(openai_api_key)
            quiz_response = chain.invoke(prompt_template)
            st.write("Quiz Generated!")
            questions, answers = split_questions_answers(quiz_response)
            st.session_state.answers = answers
            st.session_state.questions = questions
            st.write(questions)
    
    if st.button("Show Answers"):
        st.markdown(st.session_state.questions)
        st.write("----")
        st.markdown(st.session_state.answers)

if __name__ == "__main__":
    main()