import os
import streamlit as st
import time
from openai import OpenAI
from langchain.prompts import PromptTemplate 
from langchain.chat_models import ChatOpenAI
from langchain.memory import ConversationBufferWindowMemory
from langchain.chains import ConversationalRetrievalChain
from langchain.embeddings import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS

# Set up OpenAI API key
api_key = os.getenv("OPENAI_API_KEY")
client = OpenAI(api_key=api_key)

# Custom template to guide the LLM model
custom_template = """<s>[INST]You will start the conversation by greeting the user and introducing yourself as qanoon-bot,
stating your availability for legal assistance. Your next step will depend on the user's response.
If the user expresses a need for legal assistance in Pakistan, you will ask them to describe their case or problem.
After receiving the case or problem details from the user, you will provide the solutions and procedures according to the knowledge base and also give related penal codes and procedures.
However, if the user does not require legal assistance in Pakistan, you will immediately thank them and
say goodbye, ending the conversation. Remember to base your responses on the user's needs, providing accurate and
concise information regarding the Pakistan legal law and rights where applicable. Your interactions should be professional and
focused, ensuring the user's queries are addressed efficiently without deviating from the set flows.
CONTEXT: {context}
CHAT HISTORY: {chat_history}
QUESTION: {question}
ANSWER:
</s>[INST]
"""

embeddings = OpenAIEmbeddings()

# Load vector database
db = FAISS.load_local("vectordb", embeddings, allow_dangerous_deserialization=True)
db_retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 4})

# Streamlit page configuration
st.set_page_config(page_title="Qanoon-Bot")
col1, col2, col3 = st.columns([1, 4, 1])
with col2:
    st.image("https://s3.ap-south-1.amazonaws.com/makerobosfastcdn/cms-assets/Legal_AI_Chatbot.png")

st.markdown(
    """
    <style>
div.stButton > button:first-child {
    background-color: #ffd0d0;
}
div.stButton > button:active {
    background-color: #ff6262;
}
div[data-testid="stStatusWidget"] div button {
    display: none;
}
.reportview-container {
    margin-top: -2em;
}
#MainMenu {visibility: hidden;}
.stDeployButton {display:none;}
footer {visibility: hidden;}
#stDecoration {display:none;}
button[title="View fullscreen"] {
    visibility: hidden;
}
    </style>
""",
    unsafe_allow_html=True,
)

# Function to reset conversation
def reset_conversation():
    st.session_state.messages = []
    st.session_state.memory.clear()

# Initialize session state
if "messages" not in st.session_state:
    st.session_state.messages = []

if "memory" not in st.session_state:
    st.session_state.memory = ConversationBufferWindowMemory(k=5, memory_key="chat_history", return_messages=True, output_key='answer')

# Initialize the prompt
prompt = PromptTemplate(template=custom_template, input_variables=['context', 'chat_history', 'question'])

# Initialize the LLM
llm = ChatOpenAI(temperature=0.2, model_name='gpt-3.5-turbo-0125')

# Create the ConversationalRetrievalChain
qa = ConversationalRetrievalChain.from_llm(
    llm=llm,
    memory=st.session_state.memory,
    retriever=db_retriever,
    combine_docs_chain_kwargs={'prompt': prompt}
)

# Display chat history
for message in st.session_state.messages:
    with st.chat_message(message.get("role")):
        st.write(message.get("content"))

# Handle user input
input_prompt = st.chat_input("Say something")

if input_prompt:
    with st.chat_message("user"):
        st.write(input_prompt)

    st.session_state.messages.append({"role": "user", "content": input_prompt})

    with st.chat_message("assistant"):
        with st.status("Thinking 💡...", expanded=True):
            result = qa.invoke(input=input_prompt)

            message_placeholder = st.empty()

            full_response = "**_Note: Information provided by Qanoon-Bot may be inaccurate._** \n\n\n"
            for chunk in result["answer"]:
                full_response += chunk
                time.sleep(0.02)
                message_placeholder.markdown(full_response + " ▌")
                
    st.session_state.messages.append({"role": "assistant", "content": result["answer"]})
    st.button('Reset All Chat 🗑️', on_click=reset_conversation)