import streamlit as st
import os
from llama_index import (
    ServiceContext,
    SimpleDirectoryReader,
    VectorStoreIndex,
)
from llama_index.llms import OpenAI
from openai import OpenAI

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

# Define Streamlit layout and interaction
st.title("Grounded Generations")

# Upload PDF
uploaded_file = st.file_uploader("Choose a PDF file", type="pdf")

@st.cache_resource(show_spinner=False)
def load_data(uploaded_file):
    with st.spinner('Indexing document...'):
        # Save the uploaded file temporarily
        with open("temp.pdf", "wb") as f:
            f.write(uploaded_file.read())
        # Read and index documents using SimpleDirectoryReader
        reader = SimpleDirectoryReader(input_dir="./", recursive=False)
        docs = reader.load_data()
        service_context = ServiceContext.from_defaults(
            llm=OpenAI(
                model="gpt-3.5-turbo-16k",
                temperature=0.1,
            ),
            system_prompt="You are an AI assistant that uses context from a PDF to assist the user in generating text."
        )
        index = VectorStoreIndex.from_documents(docs, service_context=service_context)
    return index

# Placeholder for document indexing
if uploaded_file is not None:
    index = load_data(uploaded_file)

# Take user query input
user_query = st.text_input("Search for the products/info you want to use to ground your generated text content:")

# Initialize session_state for retrieved_text if not already present
if 'retrieved_text' not in st.session_state:
    st.session_state['retrieved_text'] = ''

# Search and display retrieved text
if st.button("Retrieve"):
    with st.spinner('Retrieving text...'):
        # Use VectorStoreIndex to search
        query_engine = index.as_query_engine(similarity_top_k=3) 
        st.session_state['retrieved_text'] = query_engine.query(user_query)
        st.write(f"Retrieved Text: {st.session_state['retrieved_text']}") # store the retrieved text as a streamlit state variable 

# Select content type
content_type = st.selectbox("Select content type:", ["Blog", "Tweet"]) # make some default nonsense 

# Generate text based on retrieved text and selected content type
if st.button("Generate") and content_type:
    with st.spinner('Generating text...'):
        # Generate text using OpenAI API
        try:
            if content_type == "Blog":
                prompt = f"Write a blog about 500 words in length using the {st.session_state['retrieved_text']}" #uses content from the stored state
            elif content_type == "Tweet":
                prompt = f"Compose a tweet using the {st.session_state['retrieved_text']}"
            response = client.chat.completions.create(model="gpt-3.5-turbo-16k",
            messages=[
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": prompt}
            ])
            generated_text = response.choices[0].message.content
            st.write(f"Generated Text: {generated_text}")
        except Exception as e:
            st.write(f"An error occurred: {e}")