import streamlit as st
import os
from together import Together
import PyPDF2

# Initialize the client
client = Together(api_key=os.environ.get("TOGETHER_API_KEY"))

def get_response(question, pdf_files):
    texts = [get_pdf_content(pdf) for pdf in pdf_files]
    combined_text = " ".join(texts)  # Combine all PDF texts into a single string

    # Embedding the prompt template
    prompt_template = f"""
    Answer the question concisely, focusing on the most relevant and important details from the PDF context. 
    Refrain from mentioning any mathematical equations, even if they are present in provided context. 
    Focus on the textual information available. Please provide direct quotations or references from PDF
    to back up your response. If the answer is not found within the PDF, 
    please state "answer is not available in the context."

    Context:\n{combined_text}\n
    Question:\n{question}\n
    Example response format:
    Overview: 
    (brief summary or introduction)
    Key points: 
    (point 1: paragraph for key details)
    (point 2: paragraph for key details)
    ...
    Use a mix of paragraphs and points to effectively convey the information.
    """
    
    response = client.chat.completions.create(
        model="meta-llama/Llama-3-8b-chat-hf",
        messages=[{"role": "user", "content": prompt_template}],
    )
    return response.choices[0].message.content

def get_pdf_content(pdf_file):
    reader = PyPDF2.PdfReader(pdf_file)
    text = ""
    for page in reader.pages:
        text += page.extract_text()
    return text

# Setup session state for maintaining conversation history
if 'history' not in st.session_state:
    st.session_state.history = []

# Set up the layout
st.set_page_config(page_title="Llama-3 Chatbot", layout="wide")
st.title('Interactive Llama-3 Chatbot with PDF Context')

# Sidebar for PDF upload
with st.sidebar:
    uploaded_files = st.file_uploader("Upload PDF files:", accept_multiple_files=True, type=['pdf'])
    st.sidebar.markdown("""
    **Instructions:**
    - Upload PDF files to include their text as context in the conversation.
    - Type your questions in the input box below the chat history.
    """)

# Main chat area
st.markdown("## Conversation")
chat_container = st.container()

# Text input for questions
user_input = st.text_input("Ask your question:", key="user_input")

# Submit button
# Submit button
if st.button("Send"):
    if uploaded_files and user_input:
        # Append user question to history
        st.session_state.history.append(f"You: {user_input}")

        # Get the response from Llama-3
        with st.spinner("Fetching your answer..."):
            answer = get_response(user_input, uploaded_files)
            st.session_state.history.append(f"Llama-3: {answer}")

        # Display conversation history
        with chat_container:
            for message in st.session_state.history:
                st.text_area("", value=message, height=100, key=message[:15])

        # Clearing handled by re-running the script with cleared input.


# Display errors if needed
if not uploaded_files:
    st.error("Please upload at least one PDF file to proceed.")
elif not user_input:
    st.error("Please enter a question to get an answer.")