from langchain.agents import (
    load_tools,
    create_react_agent,
    AgentExecutor,
    tool,
)
from langchain_openai import ChatOpenAI
from langchain.memory import ConversationBufferMemory
from langchain_community.tools import DuckDuckGoSearchRun
from langchain import hub
import streamlit as st

search = DuckDuckGoSearchRun()


@tool
def duckduckgo_webmd_search(text: str) -> str:
    """Uses the web to gather more medical information about medical query from webmd.com.
    Use this tool when you need more medical information from webmd.com to
    provide more accurate results. Don't use this if you want to search from other websites
    or if you want information regarding non-medical query. Receives a medical query as input,
    searches the web for results and gives relevant information as output string"""
    result = search.run(f"site:webmd.com {text}")
    output = str(result)
    return output


memory = ConversationBufferMemory(memory_key="chat_history")
llm = ChatOpenAI(temperature=0)
tools = load_tools([], llm=llm)
tools = tools + [duckduckgo_webmd_search]

# Get the prompt to use - you can modify this!
prompt = hub.pull("hwchase17/react")
agent = create_react_agent(llm=llm, tools=tools, prompt=prompt)

agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    memory=memory,
    verbose=True,
    handle_parsing_errors=True,
)


def handle_userinput(user_question):
    agent_executor.invoke({"input": user_question})
    st.session_state.chat_history = memory.chat_memory.messages

    for i, message in enumerate(st.session_state.chat_history):
        if i % 2 == 0:
            st.markdown(("User: " + message.content))
        else:
            st.markdown(("AI: " + message.content))


def main():
    if "conversation" not in st.session_state:
        st.session_state.conversation = None
    if "chat_history" not in st.session_state:
        st.session_state.chat_history = None

    st.header("Ask any medical query that you want answers from webmd.com")
    user_question = st.chat_input("Ask any medical query")
    if user_question:
        handle_userinput(user_question)


if __name__ == "__main__":
    main()