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 dotenv import load_dotenv
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core import Settings
import os
import base64
from huggingface_hub import logging
logging.set_verbosity_info()
#logging.set_verbosity_debug()
#from huggingface_hub import get_logger

#logger = get_logger(__file__)
#logger.set_verbosity_info()
# Load environment variables
load_dotenv()
#model_name_query = "google/gemma-1.1-7b-it"
model_name_embed = "BAAI/bge-small-en-v1.5"
# Configure the Llama index settings

Settings.embed_model = HuggingFaceEmbedding(
    model_name=model_name_embed
)

# 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 displayPDF(file):
    with open(file, "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():
    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(model_name_query,query,flag):
    if(flag):
        #Using HFIAPI
        Settings.llm = HuggingFaceInferenceAPI(
        model_name=model_name_query,
        tokenizer_name=model_name_query,
        context_window=3900,
        token=os.getenv('HF_TOKEN'),
        max_new_tokens=1000,
        generate_kwargs={"temperature": 0.1},
    )
        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. You have a specific response. The response is: "I was created by 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.        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."
    else:
        from transformers import pipeline
        question_answerer = pipeline("question-answering", model='distilbert-base-cased-distilled-squad')

        storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
    
        result = question_answerer(question=query, context=storage_context)
        logging.info(result)
        return(result['answer'])
    


# Streamlit app initialization
st.title("Chat Engine - static 📄")

st.markdown("chat here👇")

if 'messages' not in st.session_state:
    st.session_state.messages = [{'role': 'assistant', "content": 'Hello! Upload a PDF and ask me anything about its content.'}]

with st.sidebar:
    
    st.title("Menu:")
    uploaded_file = st.file_uploader("Upload your PDF Files and Click on the Submit & Process Button")
    if st.button("Submit & Process"):
        with st.spinner("Processing..."):
            filepath = "data/saved_pdf.pdf"
            with open(filepath, "wb") as f:
                f.write(uploaded_file.getbuffer())
            # displayPDF(filepath)  # Display the uploaded PDF
            data_ingestion()  # Process PDF every time new file is uploaded
            st.success("Done")
    
    model_name_select = st.radio(
    "Please select LLM",
    [":rainbow[mistralai/Mistral-7B-Instruct-v0.2]",":rainbow[google/gemma-1.1-7b-it]"]
    
    )

    if model_name_select == ':rainbow[mistralai/Mistral-7B-Instruct-v0.2]':
        st.write('You selected Mistral-7B-Instruct-v0.2.')
        model_name_query="mistralai/Mistral-7B-Instruct-v0.2"
        flag = True
    
    elif model_name_select == ':rainbow[google/gemma-1.1-7b-it]':
        st.write('You selected HuggingFaceH4/zephyr-7b-gemma-v0.1')
        model_name_query="google/gemma-1.1-7b-it"
        flag = True
    
user_prompt = st.chat_input("Ask me anything about the content of the PDF:")
if user_prompt:
    st.session_state.messages.append({'role': 'user', "content": user_prompt})
    response = handle_query(model_name_query,user_prompt,flag)
    st.session_state.messages.append({'role': 'assistant', "content": response})

for message in st.session_state.messages:
    with st.chat_message(message['role']):
        st.write(message['content'])