import streamlit as st
from langchain_community.chat_models import ChatOpenAI
from langchain_community.callbacks import get_openai_callback
from langchain.chains.question_answering import load_qa_chain

from utils.process_data import process_text, pdf_to_text

MODEL = st.secrets["MODEL4"]

st.set_page_config(page_title="Summarizer with LLM QA", page_icon="⚖️")
st.title("Summarize Text")
st.subheader("🚗 🔗 LLM/Question Answering")

maxw = st.slider('MAX words', 50, 1000, step=10, value=200)
minw = st.slider('MIN words', 10, 500, step=10, value=50)

sentence = st.text_area('Please paste your article:', height=50)
button = st.button("Summarize")
query = f"Summarize the content of the uploaded PDF file in more that {minw} words and less than {maxw} words. Focus on capturing the main ideas and key points discussed in the document. Use your own words and ensure clarity and coherence in the summary."

with st.spinner("Generating Summary.."):
    if button and sentence:
        knowledgeBase = process_text(sentence)
        docs = knowledgeBase.similarity_search(query)
        llm = ChatOpenAI(model=MODEL, temperature=0.1, openai_api_key=st.secrets["OPENAI_API_KEY"])
        chain = load_qa_chain(llm, chain_type='stuff')
        with get_openai_callback() as cost:
            response = chain.run(input_documents=docs, question=query)
            print(cost)
        st.subheader('Summary Results:')
        st.write(response)

st.divider()

st.subheader('🚙🔗 Summarize PDF')
pdf_path = st.file_uploader('Upload your PDF Document', type='pdf')
button2 = st.button("Summarize PDF")

if pdf_path is not None and button2:
    text = pdf_to_text(pdf_path)
    knowledgeBase = process_text(text)
    with st.spinner("Generating PDF Summary.."):
        docs = knowledgeBase.similarity_search(query)
        llm = ChatOpenAI(model=MODEL, temperature=0.1, openai_api_key=st.secrets["OPENAI_API_KEY"])
        chain = load_qa_chain(llm, chain_type='stuff')
        with get_openai_callback() as cost:
            response2 = chain.run(input_documents=docs, question=query)
            print(cost)
        st.subheader('Summary Results:')
        st.write(response2)