import os
import streamlit as st
from anthropic import Anthropic
from datetime import datetime
from debrief_sequence import DEBRIEF_SEQUENCE

# Initialize session state for debrief
if 'in_reflection' not in st.session_state:
    st.session_state.in_reflection = False
if 'debrief_stage' not in st.session_state:
    st.session_state.debrief_stage = 0
if 'final_notes' not in st.session_state:
    st.session_state.final_notes = ""

# ... existing code ...

# Add Reflect button at the top of chat
if not st.session_state.in_reflection:
    if st.button("🤔 I'm ready to debrief", use_container_width=True):
        st.session_state.in_reflection = True
        st.session_state.debrief_stage = 0
        st.rerun()

# Handle reflection mode
if st.session_state.in_reflection:
    st.markdown("## 🪞 Final Debrief Sequence")
    
    # Display current stage of debrief
    current_debrief = DEBRIEF_SEQUENCE[st.session_state.debrief_stage]
    
    # Display the debrief content in a styled container
    st.markdown(f"""
    <div style="background-color: #f8f9fa; padding: 20px; border-radius: 10px; margin: 20px 0;">
        <div style="color: #666; font-size: 0.9em; margin-bottom: 10px;">
            {current_debrief["type"].replace("_", " ").title()}
        </div>
        <div style="white-space: pre-line;">
            {current_debrief["content"]}
        </div>
    </div>
    """, unsafe_allow_html=True)
    
    # Navigation buttons
    col1, col2, col3 = st.columns([1, 2, 1])
    
    with col1:
        if st.session_state.debrief_stage > 0:
            if st.button("← Previous", use_container_width=True):
                st.session_state.debrief_stage -= 1
                st.rerun()
    
    with col3:
        if st.session_state.debrief_stage < len(DEBRIEF_SEQUENCE) - 1:
            if st.button("Next →", use_container_width=True):
                st.session_state.debrief_stage += 1
                st.rerun()
        elif st.session_state.debrief_stage == len(DEBRIEF_SEQUENCE) - 1:
            if st.button("Complete Reflection", use_container_width=True):
                st.session_state.in_reflection = False
                st.session_state.debrief_stage = 0
                st.rerun()
    
    # Optional note-taking area
    st.markdown("### Your Notes")
    st.text_area(
        "Use this space to write any thoughts, feelings, or insights that arise...",
        value=st.session_state.final_notes,
        key="reflection_notes",
        height=100,
        help="Your notes are private and will be cleared when you start a new session."
    )