import streamlit as st
from langchain.agents import initialize_agent, AgentType, Tool
from langchain.chains import LLMMathChain, RetrievalQA
from langchain.chat_models import ChatOpenAI
from langchain.embeddings import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.utilities import SQLDatabase
from langchain.vectorstores import FAISS
from langchain_experimental.sql import SQLDatabaseChain
from langchain.document_loaders import PyPDFLoader
from langchain.schema import ChatMessage
from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler


# Dummy users and passwords
USERS = {
    "johndoe": {"password": "Streamlit@12", "username": "John Doe"},
    "janesmith": {"password": "Streamlit@12", "username": "Jane Smith"}
}

# Initialize the system
llm = ChatOpenAI(temperature=0, model="gpt-3.5-turbo-0613")
llm_math_chain = LLMMathChain.from_llm(llm=llm, verbose=True)
db = SQLDatabase.from_uri("sqlite:///utilities_monthly.db")
db_chain = SQLDatabaseChain.from_llm(llm, db, verbose=True)
pdf_loader = PyPDFLoader("sewa.pdf")
pages = pdf_loader.load_and_split()
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(pages)
embeddings = OpenAIEmbeddings()
retriever = FAISS.from_documents(docs, embeddings).as_retriever()

tools = [
    Tool(
        name="FAQ",
        func=RetrievalQA.from_chain_type(llm=llm, retriever=retriever),
        description="Useful for answering generic questions related to SEWA, bill, usage, etc."
    ),
    Tool(
        name="Electricity-DB",
        func=db_chain.run,
        description="Useful for answering questions about Customer Energy Usage.Plase use this username given in the prompt Input should be in the form of a question containing full context"
    )
]
callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
agent_executor = initialize_agent(tools, llm, agent=AgentType.OPENAI_FUNCTIONS, verbose=True,
                                  callback_manager=callback_manager)
# Styling for sidebar
# Set the secondary text color
st.markdown("""
    <style>
        .stApp > div > div > div > div:first-child {
            background-color: blue;
        }
    </style>
    """, unsafe_allow_html=True)



with st.sidebar:
    if "logged_in_user" not in st.session_state:
        st.session_state["logged_in_user"] = None

    if st.session_state["logged_in_user"]:
        user_data = st.session_state["logged_in_user"]
        st.title(f"Welcome {user_data['username']}!")

        st.write("You can ask about energy usage, bill, due date, etc..")
        # col1, col2, col3 = st.columns([1,18,1])

        if st.button("Log Out"):
            st.session_state["logged_in_user"] = None
            st.session_state["messages"] = [ChatMessage(role="assistant", content="How can I help you?")]  # Reset chat history
            st.experimental_rerun()  # Refresh the page to show the login options again
    else:
        st.title("SEWA - At your service")
        user_id = st.text_input("User ID")
        password = st.text_input("Password", type="password")
        if st.button("Login"):
            if user_id in USERS and USERS[user_id]['password'] == password:
                st.session_state["logged_in_user"] = USERS[user_id]
                st.success(f"Logged in as {USERS[user_id]['username']}!")
                st.experimental_rerun()  # Refresh the page immediately after a successful login
            else:
                st.error("Invalid user ID or password")
# Chat Interface
if "messages" not in st.session_state:
    st.session_state["messages"] = [ChatMessage(role="assistant", content="How can I help you?")]

for msg in st.session_state.messages:
    st.chat_message(msg.role).write(msg.content)

if prompt := st.chat_input():
    # If the user is logged in, replace "my" with the user's name.


    st.session_state.messages.append(ChatMessage(role="user", content=prompt))
    st.chat_message("user").write(prompt)
    if st.session_state["logged_in_user"]:
        user_name = st.session_state["logged_in_user"]["username"]
        prompt = prompt.replace("my", user_name)
    response = agent_executor.invoke({"input": prompt})
    st.session_state.messages.append(ChatMessage(role="assistant", content=response['output']))
    st.chat_message("assistant").write(response['output'])