import os 
import streamlit as st
import re
import time 
from PIL import Image
import ast
import numpy as np

def reset_conversation():
    st.session_state.messages = []

def display_mask_image(image_path):
    if os.path.isfile(image_path):
        image = Image.open(image_path)
        st.image(image, caption='Final Mask', use_column_width=True)
    

def tyre_synap_bot(filter_agent,image_file_path):
    if "messages" not in st.session_state:
        st.session_state.messages = []
    
    print("Found image file path: ",image_file_path)
    # Display chat messages from history on app rerun
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])
    
    # React to user input
    if prompt := st.chat_input("What is up?"):
        # Display user message in chat message container
        st.chat_message("user").markdown(prompt)
        # Add user message to chat history
        st.session_state.messages.append({"role": "user", "content": prompt})
        
        ai_response = filter_agent.invoke(
            {
                "input": f'{prompt}, provided image path: {image_file_path}'
            }
            )
        
        # ai_response = filter_agent.run(f'{prompt} provided image path :{image_file_path}')
        
        response = f"Echo: {ai_response['output']}"
        with st.chat_message("assistant"):
            message_placeholder = st.empty()
            full_response = ""
            if 'mask' in ai_response['output']:
                display_mask_image('final_mask.png')
            
            for chunk in re.split(r'(\s+)', response):
                full_response += chunk + " "
                time.sleep(0.01)
                # Add a blinking cursor to simulate typing
                message_placeholder.markdown(full_response + "▌")
        # Add assistant response to chat history
        st.session_state.messages.append({"role": "assistant", "content": full_response})
        st.button('Reset Chat', on_click=reset_conversation)