from langchain.document_loaders import PyPDFLoader
from utils import *
import os
import google.generativeai as genai
from langchain_google_genai import ChatGoogleGenerativeAI
from dotenv import load_dotenv
import streamlit as st
st.set_page_config(layout="wide", page_title="QA Pair Generation from Documents",page_icon='deep-learning.png')

temperature = 0.3
pages = []
numPairs = 2
option = ''
optionCategory = ("Long QA Pairs", "MCQs", "Short QA Pairs")

load_dotenv()
genai.configure(api_key=os.environ.get("GOOGLE_API_KEY"))
model = ChatGoogleGenerativeAI(model="gemini-pro", temperature=temperature)

def LongQAPairs():
    response = []
    with st.spinner('Generating Long Question Answer Pairs...'):
        response =getLongQAPairs(pages[0:len(pages) - 1], numPairs, model)

    for qaPair in response:
        with st.chat_message("user"):
            st.write("Question : {}".format(qaPair['question']))
            st.write("Answer : {}".format(qaPair['answer']))

def ShortQAPairs():
    response = []
    with st.spinner('Generating Short Question Answer Pairs...'):
        response = getShortQAPairs(pages[0:len(pages) - 1], numPairs, model)

    for qaPair in response:
        with st.chat_message("user"):
            st.write("Question : {}".format(qaPair['question']))
            st.write("Answer : {}".format(qaPair['answer']))


def McqQAPairs():
    response = []
    with st.spinner('Generating MCQ Question Answer Pairs...'):
        response = getMcqQAPairs(pages[0:len(pages) - 1], numPairs, model)

    for qaPair in response:
        with st.chat_message("user"):
            st.radio(label=qaPair['question'],options=qaPair["options"],disabled=True,index=qaPair['correct_option_index'])


with st.sidebar:
    st.image('Pic.png')
    st.title("Final Year Project")
    st.divider()
    with st.container(border=True):
        st.text('Model: Gemini Pro', help='Developed by Google \n')
        temperature = st.slider('Temperature:', 0.0, 1.0, 0.3, 0.1)

    code = '''Team Members CSE(20-37):
    \nAmbuj Raj BT20CSE054 \nSrishti Pandey BT20CSE068 \nPrateek Niket BT20CSE211 \nSmriti Singh BT20CSE156'''
    st.code(code, language='JAVA')
    code = '''Mentored By: \nDr. Amol Bhopale'''
    st.code(code, language='JAVA')

st.title('Question Answer Pair Generation From Documents')

with st.container(border=True):
    col1, col2 = st.columns(2)
    with col1:
        st.write("Please Upload Your File")
        uploaded_file = st.file_uploader("Choose a file", type='.pdf', accept_multiple_files=False)
        if uploaded_file is not None:
            with open("temp.pdf", "wb") as f:
                f.write(uploaded_file.getbuffer())

            # Get the path of the uploaded file
            file_path = "temp.pdf"

            pdf_loader = PyPDFLoader(file_path)
            pages = pdf_loader.load_and_split()
            print(len(pages))

    with col2:
        st.write('Please Choose your Configuration')
        option = st.selectbox(
            "In Which Category would you like to Generate Question Answer Pairs?",
            optionCategory,
            index=None,
            placeholder="Select Category of Question Answer Pairs",
        )
        numPairs = st.number_input('Number of QA Pairs', min_value=1, max_value=20, step=2,value=2)

if st.button("Generate", type="primary"):
    if option == "Long QA Pairs" and len(pages) and option in optionCategory:
        LongQAPairs()
    elif option == "MCQs" and len(pages) and option in optionCategory:
        McqQAPairs()
    elif option == "Short QA Pairs" and len(pages) and option in optionCategory:
        ShortQAPairs()
    elif len(pages) or option not in optionCategory or uploaded_file is None:
        st.error('Required Fields are Missing!', icon="🚨")