import io
import os
import streamlit as st
from dotenv import load_dotenv
from PyPDF2 import PdfReader
from langchain_community.embeddings import OpenAIEmbeddings
from langchain_community.llms import OpenAI
from langchain.prompts import PromptTemplate
from langchain_community.vectorstores import FAISS
import openai

# Load environment variables
load_dotenv()
openai_api_key = os.getenv('OPENAI_API_KEY')

# Initialize Streamlit session states
if 'vectorDB' not in st.session_state:
    st.session_state.vectorDB = None

# Function to extract text from a PDF file
def get_pdf_text(pdf):
    text = ""
    pdf_reader = PdfReader(pdf)
    for page in pdf_reader.pages:
        text += page.extract_text()
    return text

# Function to create a vector database
def get_vectorstore(text_chunks):
    embeddings = OpenAIEmbeddings(api_key=openai_api_key)
    vectorstore = FAISS.from_texts(texts=text_chunks, embedding=embeddings)
    return vectorstore

# Function to split text into chunks
def get_text_chunks(text):
    text_chunks = text.split('\n\n')  # Modify this based on your text splitting requirements
    return text_chunks

# Function to process PDF and create vector database
def processing(pdf):
    raw_text = get_pdf_text(pdf)
    text_chunks = get_text_chunks(raw_text)
    vectorDB = get_vectorstore(text_chunks)
    return vectorDB

# Function to generate quiz questions, options, and correct answers
def generate_quiz_content(text, num_questions):
    messages = [{'role': 'system', 'content': 'You are a helpful assistant.'}]
    for i in range(num_questions):
        user_message = {'role': 'user', 'content': f'Generate question {i + 1} from the given text:\n{text}'}
        messages.append(user_message)

    response = openai.ChatCompletion.create(
        model='gpt-3.5-turbo-16k',
        messages=messages,
        max_tokens=200,
        temperature=0.7
    )
    quiz_data = []

    for choice in response.get('choices', []):
        content = choice.get('content', '').strip()
        options_start = content.find("Options:") + len("Options:")
        correct_answer_start = content.find("Correct Answer:") + len("Correct Answer:")

        question = content[:options_start].strip()
        options_str = content[options_start:correct_answer_start].strip()
        correct_answer = content[correct_answer_start:].strip()

        options = [opt.strip() for opt in options_str.split(',')]
        quiz_data.append((question, options, correct_answer))

    return quiz_data

# Function to retrieve quiz data from the vector database
def retrieve_quiz_data(vectorDB, num_questions):
    # Replace this with your actual logic to query the vector database and retrieve quiz data
    return [(f"Question {i + 1}", [f"Option {j}" for j in range(1, 5)], f"Option {i % 4 + 1}") for i in range(num_questions)]

# Function to generate quiz questions
def generate_quiz(quiz_name, quiz_topic, num_questions, pdf_content):
    st.header(f"Quiz Generator: {quiz_name}")
    st.subheader(f"Topic: {quiz_topic}")

    # Process PDF and create vector database
    if st.button('Generate Quiz'):
        st.session_state['vectorDB'] = processing(pdf_content)
        st.success('PDF Processed and Vector Database Created')

    # Generate Quiz Questions
    if st.session_state.vectorDB:
        # Retrieve quiz data from the vector database
        generated_quiz_data = retrieve_quiz_data(st.session_state.vectorDB, num_questions)

        # Display retrieved questions, options, and correct answers
        for i, (question, options, correct_answer) in enumerate(generated_quiz_data):
            st.subheader(f"Question {i + 1}")
            st.write(f"Retrieved Question: {question}")
            st.write(f"Options: {', '.join(options)}")
            st.write(f"Correct Answer: {correct_answer}")

    # Save button to store vector database
    if st.session_state.vectorDB:
        if st.button('Save Vector Database'):
            st.success('Vector Database Saved')

if __name__ =='__main__':
    st.set_page_config(page_title="CB Quiz Generator", page_icon="📝")
    st.title('CB Quiz Generator')

    # User inputs
    quiz_name = st.text_input('Enter Quiz Name:')
    quiz_topic = st.text_input('Enter Quiz Topic:')
    num_questions = st.number_input('Enter Number of Questions:', min_value=1, value=1, step=1)
    pdf_content = st.file_uploader("Upload PDF Content for Questions:", type='pdf')

    # Generate quiz if all inputs are provided
    if quiz_name and quiz_topic and num_questions and pdf_content:
        generate_quiz(quiz_name, quiz_topic, num_questions, pdf_content)