import streamlit as st
import llm_part
import os
from langchain_groq import ChatGroq

# Sidebar to select the LLM model
st.sidebar.title("LLM Model Selector")
llm_model = st.sidebar.selectbox("Select LLM Model", ("Google Gemini", "Llama"))


# Define Llama-specific configurations
if llm_model == "Google Gemini":  # Check if "Google Gemini" is selected
    llm = llm_part.llm_1  # Assign Google Gemini to llm
else:
    llm = llm_part.llm_2  # Use the Llama model




# Main app
st.title("Jony's Custom Research Notes Extracted from PDFs Using " + llm_model)
option = st.selectbox("Select PDF Source:", ("Enter URL", "Upload Local File"))

document_text = ""

if option == "Enter URL":
    pdf_url = st.text_input("Enter the PDF URL:")
    
    if pdf_url:
        try:
            with st.spinner("Processing PDF from URL..."):
                local_pdf_path = "downloaded_paper.pdf"
                llm_part.download_pdf_from_url(pdf_url, local_pdf_path)
                document_text = llm_part.extract_text_from_pdf(local_pdf_path)
                os.remove(local_pdf_path)

        except Exception as e:
            st.error(f"Error processing PDF from URL: {e}")

elif option == "Upload Local File":
    uploaded_file = st.file_uploader("Choose a PDF file", type="pdf")
    
    if uploaded_file is not None:
        try:
            with st.spinner("Processing uploaded PDF..."):
                local_pdf_path = "uploaded_paper.pdf"
                with open(local_pdf_path, "wb") as f:
                    f.write(uploaded_file.read())
                document_text = llm_part.extract_text_from_pdf(local_pdf_path)
                os.remove(local_pdf_path)

        except Exception as e:
            st.error(f"Error processing uploaded PDF: {e}")

if document_text:
    with st.spinner("Generating the summary..."):
        query = llm_part.prompt.format(document_text=document_text[:20000])
        result = llm.invoke(query)
        st.write("### Summary in Table Format:")
        st.write(result.content) 
        lines = result.content.split('\n')
        paragraph_output = []

        for line in lines[2:]:
            if "|" not in line or not line.strip():
                continue

            parts = [part.strip() for part in line.split("|") if part.strip()]
            if len(parts) == 2:
                _, details = parts
                if "Not specified" in details or "Not mentioned" in details:
                    continue
                details_clean = llm_part.clean_html_tags(details)
                paragraph_output.append(details_clean) 

        paragraph_output = ". ".join(paragraph_output) + "."
        paragraph_output = paragraph_output.replace(" ,", ",").replace(" .", ".")
        paragraph_output = paragraph_output.replace(". CNN", ". In this approach, CNN").replace("Federated learning (FL)", "The use of Federated Learning (FL)")

        paragraph_output = paragraph_output.replace("The use of Federated Learning", "The study explores the use of Federated Learning")
        paragraph_output = paragraph_output.replace("In this approach, CNN", "In this approach, a combination of CNN models was used to enhance performance")
        paragraph_output = paragraph_output.replace("achieved", "yielded results indicating")
        paragraph_output = paragraph_output.replace("slightly lower", "only marginally lower")
    
        query2 = llm_part.prompt2.format(paragraph=paragraph_output)
        result2 = llm.invoke(query2)
        st.write("### Answer in Paragraph Style:")
        st.markdown("""
            <style>
            .justified-text {
                text-align: justify;
            }
            </style>
            """, unsafe_allow_html=True)

        st.markdown(f"<div class='justified-text'>{result2.content}</div>", unsafe_allow_html=True)


        #pip install -r requirements.txt