import streamlit as st

PROMPTS = [
    "What is your initial instinct or plan?",
    "Can you identify three aspects of this case you haven’t considered before? For example: Are there recent legal precedents? Have there been changes in the opposing counsel’s approach? Is there something unique about your client’s circumstances?",
    "What assumptions are you making about the opposing counsel, the judge, or your client’s priorities?",
    "How might these assumptions be limiting your options or influencing your recommendation?",
    "Besides accepting the settlement or going to trial, can you think of at least two other options? Even unconventional ones.",
    "What uncertainties are present here, and how could they open up new possibilities?",
    "Reflecting on this process, what have you learned about your decision-making? How might you apply this mindful approach in future cases?",
    "Would you like to log this decision and set a reminder to reflect on the outcome in a few weeks?"
]

st.set_page_config(page_title="Mindful Decision-Making Coach", page_icon=":balance_scale:", layout="centered")
st.title("Mindful Decision-Making Coach")
st.markdown("_A chatbot for mindful legal decisions, inspired by Ellen Langer_")

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

# Show chat history
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])

# Handle input and prompt flow
user_input = None
if st.session_state.awaiting_first:
    user_input = st.chat_input("Describe your legal decision or dilemma...")
elif st.session_state.prompt_idx < len(PROMPTS):
    user_input = st.chat_input("Your response...")

if user_input:
    if st.session_state.awaiting_first:
        st.session_state.messages.append({"role": "user", "content": user_input})
        st.session_state.messages.append({"role": "assistant", "content": PROMPTS[0]})
        st.session_state.prompt_idx = 1
        st.session_state.awaiting_first = False
    elif st.session_state.prompt_idx < len(PROMPTS):
        st.session_state.messages.append({"role": "user", "content": user_input})
        st.session_state.messages.append({"role": "assistant", "content": PROMPTS[st.session_state.prompt_idx]})
        st.session_state.prompt_idx += 1
    # No need for any rerun or query param tricks!

# Final message after all prompts
if not st.session_state.awaiting_first and st.session_state.prompt_idx >= len(PROMPTS):
    st.markdown("**Thank you for sharing your decision. I have logged your responses and will remind you to reflect on the outcome in a few weeks.**")