import streamlit as st
from crewai.tasks.task_output import TaskOutput

def streamlit_callback(step_output:TaskOutput):
    st.markdown("---")
    for step in step_output:
        if isinstance(step, tuple) and len(step) == 2:
            action, observation = step

            # Display action information
            if isinstance(action, dict) and all(key in action for key in ["tool", "tool_input", "log"]):
                st.markdown(f"**Tool:** {action['tool']}")
                st.markdown(f"**Tool Input:** {action['tool_input']}")
                st.markdown(f"**Log:** {action['log']}")
                if "Action" in action:
                    st.markdown(f"# 📝 Processing Action...")
                    st.markdown(f"**Action:** {action['Action']}")
                st.markdown(f"**Action Input:** ```json\n{action['tool_input']}\n```")
            elif isinstance(action, str):
                if action != 'log':
                    continue
                st.markdown(f"# Action Result...")
                st.markdown(f"**Action:** {action}")
            else:
                st.markdown(f"**Action:** {str(action)}")

            # Display observation information
            st.markdown(f"**Thought**")
            if isinstance(observation, str):
                observation_lines = observation.split('\n')
                for line in observation_lines:
                    if line.startswith('Title: '):
                        st.markdown(f"**Title:** {line[7:]}")
                    elif line.startswith('Link: '):
                        st.markdown(f"**Link:** {line[6:]}")
                    elif line.startswith('Snippet: '):
                        st.markdown(f"**Snippet:** {line[9:]}")
                    elif line.startswith('-'):
                        st.markdown(line)
                    else:
                        st.markdown(line)
            else:
                st.markdown(str(observation))
        else:
            st.markdown(step)