import os
import streamlit as st
from anthropic import Anthropic
from dotenv import load_dotenv

# Initialize page configuration first
st.set_page_config(
    page_title="Attachment Style Roleplay Simulator",
    page_icon="🗣️",
    layout="wide"
)

# Handle API key
try:
    # Load from .env file for local development
    if os.path.exists(".env"):
        load_dotenv()
    
    # Try multiple ways to get the API key
    api_key = None
    
    # Try direct environment variable
    if 'ANTHROPIC_API_KEY' in os.environ:
        api_key = os.environ['ANTHROPIC_API_KEY']
    
    # Try getenv
    if not api_key:
        api_key = os.getenv('ANTHROPIC_API_KEY')
    
    # Try Streamlit secrets
    if not api_key and hasattr(st, 'secrets') and 'ANTHROPIC_API_KEY' in st.secrets:
        api_key = st.secrets['ANTHROPIC_API_KEY']
    
    if not api_key:
        st.error("""
        ⚠️ No API key found. Please make sure:
        1. You've added the secret in Hugging Face Space settings
        2. The secret is named exactly 'ANTHROPIC_API_KEY'
        3. The Space has been rebuilt after adding the secret
        """)
        st.stop()

    # Initialize Anthropic client
    c = Anthropic(api_key=api_key)
    
except Exception as e:
    st.error(f"""
    ⚠️ Error initializing Anthropic client: {str(e)}
    
    Please check:
    1. Your ANTHROPIC_API_KEY is set correctly
    2. You're using a valid API key from Anthropic
    """)
    st.stop()

# Initialize session state variables
if 'messages' not in st.session_state:
    st.session_state.messages = []
if 'setup_complete' not in st.session_state:
    st.session_state.setup_complete = False
if 'system_message' not in st.session_state:
    st.session_state.system_message = ""
if 'simulation_params' not in st.session_state:
    st.session_state.simulation_params = {}
if 'in_debrief' not in st.session_state:
    st.session_state.in_debrief = False

# Main page header
st.title("Attachment Style Roleplay Simulator")

# Welcome text in main content area
st.markdown("""
Welcome!

Created by [Jocelyn Skillman LMHC](http://www.jocelynskillman.com), this tool is designed to help you practice and understand relational dynamics through the lens of attachment theory.

To learn more about attachment theory and therapeutic approaches, check out my newsletter: [@jocelynskillmanlmhc](https://jocelynskillmanlmhc.substack.com/)
""")

st.markdown("""
This specific project — the Attachment Roleplay Simulator — invites users to practice emotionally difficult conversations, like setting boundaries, through roleplay tailored to their attachment style. It's designed to help people feel into relational patterns, develop language for self-advocacy, and build nervous system capacity for connection and repair.

💡 Not sure about your attachment style?
You can take the free quiz by Sarah Peyton — or simply choose the one that resonates when you read it:

Anxious – "I often worry I've upset people or need to explain myself."

Avoidant – "I'd rather deal with things alone and not depend on anyone."

Disorganized – "I want closeness, but also feel unsafe or mistrusting."

Secure – "I can express needs and handle conflict without losing connection."

Choose what vibes — this is a practice space, not a test.
""")

# Sidebar welcome and introduction
with st.sidebar:
    st.markdown("""
Hi, I'm Jocelyn Skillman, LMHC — a clinical therapist, relational design ethicist, and creator of experimental tools that explore how AI can support (not replace) human care.

Each tool in this collection is thoughtfully designed to:

Extend therapeutic support between sessions

Model emotional safety and relational depth

Help clients and clinicians rehearse courage, regulation, and repair

Stay grounded in trauma-informed, developmentally sensitive frameworks
    """)
    
    st.divider()
    
    # Add the moved content to the bottom of sidebar
    st.markdown("""
I use Claude (by Anthropic) as the primary language model for these tools, chosen for its relational tone, contextual nuance, and responsiveness to emotionally complex prompts.

As a practicing therapist, I imagine these resources being especially helpful to clinicians like myself — companions in the work of tending to others with insight, warmth, and care.

To learn more about my work, visit:
🌐 [jocelynskillman.com](http://www.jocelynskillman.com)
📬 [Substack: Relational Code](https://jocelynskillmanlmhc.substack.com/)
    """)

# Simulation setup form
with st.form("setup_form"):
    st.header("Set Up Your Simulation")
    
   tone_archetype = st.selectbox(
    "Select the voice you'd like to engage with:",
    ["Ghost", "Sycophant", "Critic"]
)
    
    scenario = st.text_area(
        "Describe the scenario you'd like to practice:",
        placeholder="Example: Having a difficult conversation with my partner about feeling disconnected"
    )
    
    tone = st.text_input(
        "How would you like the AI to respond?",
        placeholder="Example: Direct but understanding"
    )
    
    goals = st.text_area(
        "What are your practice goals?",
        placeholder="Example: Practice expressing needs without becoming defensive"
    )
    
    submitted = st.form_submit_button("Start Simulation")
    
    if submitted:
        # Store simulation parameters for debrief
        st.session_state.simulation_params = {
            "tone_archetype": tone_archetype,
            "scenario": scenario,
            "tone": tone,
            "goals": goals
        }
        
        # Prepare system message with simulation parameters
        st.session_state.system_message = f"""
        You are a conversational partner helping someone practice difficult conversations. 
        The user has a {tone_archetype} attachment style and wants to practice: {scenario}
        
        Respond in a {tone} manner. Help them achieve their goals: {goals}
        
        Maintain a realistic conversation while providing gentle guidance when needed.
        
        IMPORTANT: If the user types "debrief" or "end roleplay", switch to debrief mode using this format:

        📝 DEBRIEF SUMMARY

        **Emotional Arc**: [Analyze how the user shifted emotionally during the interaction, noting moments of courage, freeze, protest, backtracking, or people-pleasing]

        **Goal Alignment**: [Reflect on how well they stayed aligned with their stated goal. Note small wins and struggles, affirm effort and awareness]

        **Attachment Insight**: [Offer one insight based on their {tone_archetype} attachment style. Normalize the response as protective. Gently suggest a next growth step without shaming]

        **Practical Tool**: [Offer a brief NVC or DBT tool they could try next time (e.g., needs check-in, opposite action, self-validation, distress tolerance)]

        **Bold Reframe**: [Offer one bold, loving sentence they might say if they trusted their relational worth]

        **Journaling Prompt**: [Suggest one reflective question or embodied inquiry for integration]

        Speak warmly and with psychological precision, like an emotionally attuned therapist.
        """
        
        # Reset message history and state
        st.session_state.messages = []
        st.session_state.setup_complete = True
        st.session_state.in_debrief = False
        st.rerun()

# Display status or chat interface
if not st.session_state.setup_complete:
    st.info("Please complete the simulation setup to begin.")
else:
    # Display chat history
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    # Optional somatic input logging
st.subheader("🧘 Somatic Reflection (Optional)")
somatic_input = st.text_input("What’s happening in your body right now?", key="soma_input")

if st.button("Log Somatic Entry"):
    if 'soma_log' not in st.session_state:
        st.session_state.soma_log = []
    st.session_state.soma_log.append({
        "timestamp": datetime.now().isoformat(),
        "tone": st.session_state.simulation_params.get('tone_archetype', st.session_state.simulation_params.get('tone_archetype', 'Unknown')),
        "somatic_input": somatic_input
    })
    st.success("Somatic reflection logged.")

    
    # Chat input
    if prompt := st.chat_input("Type your message here (type 'debrief' or 'end roleplay' when ready to reflect)"):
        # Check for debrief trigger
        if prompt.lower() in ["debrief", "end roleplay"] and not st.session_state.in_debrief:
            st.session_state.in_debrief = True
            # Construct conversation summary
            conversation_summary = "\n".join([
                f"{msg['role']}: {msg['content']}" 
                for msg in st.session_state.messages
            ])
            
            # Add debrief request to messages
            prompt = f"""Please provide a therapeutic debrief for this roleplay session using:

            Attachment Style: {st.session_state.simulation_params['tone_archetype']}
            Scenario: {st.session_state.simulation_params['scenario']}
            Goals: {st.session_state.simulation_params['goals']}
            AI Tone: {st.session_state.simulation_params['tone']}
            
            Conversation Summary:
            {conversation_summary}
            """
        
        # Add user message to chat history
        st.session_state.messages.append({"role": "user", "content": prompt})
        
        # Display user message
        with st.chat_message("user"):
            st.markdown(prompt)
        
        # Get AI response
        with st.chat_message("assistant"):
            with st.spinner("Thinking..."):
                try:
                    message = c.messages.create(
                        model="claude-3-opus-20240229",
                        max_tokens=1000,
                        system=st.session_state.system_message,
                        messages=[
                            {"role": msg["role"], "content": msg["content"]} 
                            for msg in st.session_state.messages
                        ]
                    )
                    
                    ai_response = message.content[0].text
                    st.markdown(ai_response)
                    
                    # Add AI response to chat history
                    st.session_state.messages.append(
                        {"role": "assistant", "content": ai_response}
                    )
                except Exception as e:
                    st.error(f"Error getting AI response: {str(e)}")
                    st.error("Please try again or contact support if the issue persists.")

# Add restart button after debrief
if st.session_state.in_debrief:
    if st.button("Start New Roleplay"):
        st.session_state.setup_complete = False
        st.session_state.in_debrief = False
        st.session_state.messages = []
        st.session_state.system_message = ""
        st.session_state.simulation_params = {}
        st.rerun()

# Footer
st.markdown("---")
st.markdown(
    "Created by [Jocelyn Skillman LMHC](http://www.jocelynskillman.com) | "
    "Learn more: [@jocelynskillmanlmhc](https://jocelynskillmanlmhc.substack.com/)"
)