import streamlit as st
from llama_index.core import StorageContext, load_index_from_storage, VectorStoreIndex, SimpleDirectoryReader, ChatPromptTemplate
from llama_index.llms.huggingface import HuggingFaceInferenceAPI
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core import Settings
import os
import base64

# Configure the Llama index settings
Settings.llm = HuggingFaceInferenceAPI(
    model_name="google/gemma-1.1-7b-it",
    tokenizer_name="google/gemma-1.1-7b-it",
    context_window=3000,
    token=os.environ.get("HF_TOKEN"),
    max_new_tokens=512,
    generate_kwargs={"temperature": 0.1},
)
Settings.embed_model = HuggingFaceEmbedding(
    model_name="BAAI/bge-small-en-v1.5"
)

# Define the directory for persistent storage and data
PERSIST_DIR = "./db"
DATA_DIR = "data"

# Ensure data directory exists
os.makedirs(DATA_DIR, exist_ok=True)
os.makedirs(PERSIST_DIR, exist_ok=True)

def display_pdf(file_path):
    """Display a PDF file in the Streamlit app."""
    with open(file_path, "rb") as f:
        base64_pdf = base64.b64encode(f.read()).decode('utf-8')
    pdf_display = f'<iframe src="data:application/pdf;base64,{base64_pdf}" width="100%" height="600" type="application/pdf"></iframe>'
    st.markdown(pdf_display, unsafe_allow_html=True)

def data_ingestion():
    """Ingest and process PDF data."""
    documents = SimpleDirectoryReader(DATA_DIR).load_data()
    storage_context = StorageContext.from_defaults()
    index = VectorStoreIndex.from_documents(documents)
    index.storage_context.persist(persist_dir=PERSIST_DIR)

def handle_query(query):
    """Handle user query and return response."""
    storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
    index = load_index_from_storage(storage_context)
    chat_text_qa_msgs = [
        (
            "user",
            """You are a Q & A assistant named PDF CHATTY, created by Huzaifa.
             You have a specific response programmed for when users specifically ask about your creator, Huzaifa.
             The response is: "I was created by Huzaifa, an enthusiast in Artificial Intelligence.
             He is dedicated to solving complex problems and delivering innovative solutions. 
             With a strong focus on machine learning, deep learning, Python, generative AI, NLP,
             and computer vision, Huzaifa is passionate about pushing the boundaries of AI to explore new possibilities." 
             For all other inquiries, your main goal is to provide answers as accurately as possible, based on the instructions and context you have been given. 
             If a question does not match the provided context or is outside the scope of the document, kindly advise the user to ask questions within the context
             of the document
              
            
            Context:
            {context_str}
            Question:
            {query_str}
            """
        )
    ]
    text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)
    
    query_engine = index.as_query_engine(text_qa_template=text_qa_template)
    answer = query_engine.query(query)
    
    if hasattr(answer, 'response'):
        return answer.response
    elif isinstance(answer, dict) and 'response' in answer:
        return answer['response']
    else:
        return "Sorry, I couldn't find an answer."

# Streamlit app initialization
st.set_page_config(page_title="PDF Information and Inference", page_icon=":book:", layout="wide")

st.title("📚 PDF Information and Inference")
st.markdown("Welcome to the Retrieval-Augmented Generation tool! You can either upload a PDF or chat with the language model directly.")

# Sidebar for PDF upload and processing
with st.sidebar:
    st.title("Menu")
    st.subheader("Upload and Process PDF")
    uploaded_file = st.file_uploader("Upload your PDF file here:", type="pdf")
    if st.button("Submit & Process PDF"):
        if uploaded_file is not None:
            with st.spinner("Processing..."):
                filepath = os.path.join(DATA_DIR, "saved_pdf.pdf")
                with open(filepath, "wb") as f:
                    f.write(uploaded_file.getbuffer())
                data_ingestion()  # Process PDF every time a new file is uploaded
                st.success("PDF processed successfully!")
                st.sidebar.write("### Uploaded PDF:")
                display_pdf(filepath)  # Display the uploaded PDF
        else:
            st.sidebar.error("Please upload a PDF file first.")

# Main section for user interaction
st.header("Chat with the PDF or the LLM")

if uploaded_file:
    st.write("You can now ask questions about the uploaded PDF.")
else:
    st.write("No PDF uploaded. You can still chat with the LLM.")

user_prompt = st.text_input("Ask me anything:")

if user_prompt:
    with st.spinner("Fetching answer..."):
        if uploaded_file:
            # Handle query with PDF context
            response = handle_query(user_prompt)
        else:
            # Handle query without PDF, just using the LLM
            response = Settings.llm.generate([user_prompt])
        st.write(response[0]["generated_text"] if isinstance(response, list) else response)