import streamlit as st
import pandas as pd

def display_chat(history):
    """Renders the chat history with custom bubbles for each message."""
    chat_container = st.container()
    with chat_container:
        for idx, chat in enumerate(history):
            # --- User message ---
            st.markdown(
                f"""
                <div class="chat-bubble user-bubble">
                    <strong>You:</strong> {chat['user']}
                </div>
                """,
                unsafe_allow_html=True
            )

            # --- Bot bubble: use the 'type' key to decide how to render ---
            st.markdown(
                """
                <div class="chat-bubble bot-bubble">
                <strong>Bot:</strong>
                """,
                unsafe_allow_html=True,
            )

            response_type = chat.get('type', 'string')  # default to 'string'
            bot_response = chat['bot']

            if response_type == 'dataframe' and isinstance(bot_response, pd.DataFrame):
                # Show top 5 rows
                df_to_display = bot_response
                if len(df_to_display) > 5:
                    st.info("Showing the first 5 rows of the DataFrame.")
                st.dataframe(df_to_display.head(5))

                # Provide a CSV download
                csv_data = df_to_display.to_csv(index=False).encode('utf-8')
                st.download_button(
                    label="Download data as CSV",
                    data=csv_data,
                    file_name=f'result_{idx+1}.csv',
                    mime='text/csv',
                    key=f'download_{idx}'
                )

            elif response_type == 'plot':
                # If it's an image object (e.g., PIL Image), show it
                st.image(bot_response, use_container_width=True)

            else:  # "string" or any other text
                st.markdown(f"{bot_response}", unsafe_allow_html=True)

            st.markdown("</div>", unsafe_allow_html=True)