import streamlit as st
from langchain_core.messages import AIMessage, HumanMessage
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from dotenv import load_dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate


load_dotenv()

#laoding embeddings
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vectordb = FAISS.load_local("faiss_index_pymupdf", embeddings)

# app config
st.set_page_config(page_title="MANOchatBot", page_icon="🤖")
st.title("MANO ChatBot")


llm = ChatGoogleGenerativeAI(model="gemini-pro",temperature=0.7,convert_system_message_to_human=True)

def augment_prompt(user_query):
    # get top  results from knowledge base
    results = vectordb.similarity_search(user_query, k=10)
    # get the text from the results
    source_knowledge = "\n".join([x.page_content for x in results])
    
    # feed into an augmented prompt
    augmented_prompt = f"""Based on the context  provided, provide an answer to the best of your knowledge.If answer is not found in the context then web search.
Use your skills to determine what kind of context is provided and tailor your response accordingly.
Also, use html bullet list format when needed.
    Contexts:
    {source_knowledge}
      
    Query: {user_query}"""
    return augmented_prompt 

def get_response(user_query):
    messages=[]
    #messages.append(res)
    prompt = HumanMessage(content=augment_prompt(user_query))
    # add to messages
    messages.append(prompt)

    res = llm(messages)

    return res.content

# session state
if "chat_history" not in st.session_state:
    st.session_state.chat_history = [
        AIMessage(content="Hello, I am a ChatBot. How can I help you?"),
    ]

    
# conversation
for message in st.session_state.chat_history:
    if isinstance(message, AIMessage):
        with st.chat_message("AI"):
            st.write(message.content)
    elif isinstance(message, HumanMessage):
        with st.chat_message("Human"):
            st.write(message.content)

# user input
user_query = st.chat_input("Type your query here...")
if user_query is not None and user_query != "":
    st.session_state.chat_history.append(HumanMessage(content=user_query))

    with st.chat_message("Human"):
        st.markdown(user_query)

    with st.chat_message("AI"):
        response = st.write(get_response(user_query))

    #st.session_state.chat_history.append(AIMessage(content=response))