import streamlit as st
from modules.data_class import DataState
from modules.tools import data_node
from modules.nodes import chatbot_with_tools, human_node, maybe_exit_human_node, maybe_route_to_tools

from langgraph.graph import StateGraph, START, END

from IPython.display import Image, display
from pprint import pprint
from typing import Literal

from langgraph.prebuilt import ToolNode

from collections.abc import Iterable
from IPython.display import display, clear_output
import sys


# Define the LangGraph chatbot
graph_builder = StateGraph(DataState)

# Add nodes
graph_builder.add_node("chatbot_healthassistant", chatbot_with_tools)
graph_builder.add_node("patient", human_node)
graph_builder.add_node("documenting", data_node)

# Define edges
graph_builder.add_conditional_edges("chatbot_healthassistant", maybe_route_to_tools)
graph_builder.add_conditional_edges("patient", maybe_exit_human_node)
graph_builder.add_edge("documenting", "chatbot_healthassistant")
graph_builder.add_edge(START, "chatbot_healthassistant")

# Compile the graph
graph_with_order_tools = graph_builder.compile()

# Streamlit UI
st.title("LangGraph Chatbot")
st.markdown("Chat with an AI-powered health assistant.")

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

user_input = st.text_input("You:", key="input")

if st.button("Send"):
    if user_input:
        # Add user input to history
        st.session_state.messages.append(("User", user_input))

        # Run LangGraph chatbot
        state = DataState(messages=st.session_state.messages, data={}, finished=False)
        for output in graph_with_order_tools.stream(state):
            response = output["messages"][-1]  # Get the last chatbot response
            st.session_state.messages.append(("Bot", response))

# Display chat history
for sender, message in st.session_state.messages:
    st.write(f"**{sender}:** {message}")