import os
import streamlit as st
from anthropic import Anthropic
from datetime import datetime
from debrief_ai import DEBRIEF_SYSTEM_PROMPT, analyze_conversation

# Page config
st.set_page_config(page_title="Practice Difficult Conversations", page_icon="💭")

# Initialize Anthropic client
try:
    # Try getting from environment variable first
    api_key = os.getenv("ANTHROPIC_API_KEY")
    if api_key:
        st.write("Debug: Found key in environment variables")
        st.write(f"Debug: Key length: {len(api_key)}")
        st.write(f"Debug: Key starts with: {api_key[:15]}")
    else:
        # Fall back to Streamlit secrets
        if "anthropic_key" in st.secrets:
            api_key = st.secrets["anthropic_key"]
            st.write("Debug: Found key in Streamlit secrets")
            st.write(f"Debug: Key length: {len(api_key)}")
            st.write(f"Debug: Key starts with: {api_key[:15]}")
        else:
            st.write("Debug: No key found in either location")
    
    if not api_key:
        st.error("""
        ⚠️ API key not found. Please ensure:
        1. You have set the ANTHROPIC_API_KEY environment variable
        2. The API key is in the correct format (starts with sk-ant-api03-)
        3. You can get a new API key from https://console.anthropic.com/
        """)
        st.stop()
    
    # Clean the API key
    api_key = api_key.strip()
    
    # Validate API key format
    if not api_key.startswith("sk-ant-api03-"):
        st.error(f"""
        ⚠️ Invalid API key format. Your key should:
        1. Start with 'sk-ant-api03-'
        2. Be from the Anthropic Console (https://console.anthropic.com/)
        3. Be set as ANTHROPIC_API_KEY environment variable
        
        Current key starts with: {api_key[:15]}
        """)
        st.stop()
    
    # Create client with explicit API key
    st.write("Debug: Attempting to create Anthropic client...")
    anthropic = Anthropic(api_key=api_key)
    st.write("Debug: Successfully created Anthropic client")
    
    # Test the client
    st.write("Debug: Testing API connection...")
    test_response = anthropic.messages.create(
        model="claude-3-opus-20240229",
        max_tokens=10,
        messages=[{"role": "user", "content": "test"}]
    )
    st.write("Debug: API test successful")
    
except Exception as e:
    st.error(f"""
    ⚠️ Error initializing Anthropic client: {str(e)}
    
    Please check:
    1. Your API key is valid and in the correct format (starts with sk-ant-api03-)
    2. You have set the ANTHROPIC_API_KEY environment variable
    3. You can get a new API key from https://console.anthropic.com/
    
    Debug info:
    - Error type: {type(e).__name__}
    - Error message: {str(e)}
    """)
    st.stop()

# Initialize session state variables
if 'messages' not in st.session_state:
    st.session_state.messages = []
if 'in_debrief' not in st.session_state:
    st.session_state.in_debrief = False
if 'debrief_messages' not in st.session_state:
    st.session_state.debrief_messages = []
if 'practice_complete' not in st.session_state:
    st.session_state.practice_complete = False
if 'conversation_analysis' not in st.session_state:
    st.session_state.conversation_analysis = None

# App header
st.title("Practice Difficult Conversations 💭")

# Only show welcome message if no conversation started
if not st.session_state.messages and not st.session_state.in_debrief:
    st.markdown("""
    Welcome to your safe space for practicing challenging interactions. Here you can:
    - Practice responding to different conversation styles
    - Build awareness of your patterns and responses
    - Develop new communication strategies
    - Process and integrate your experience with a reflective debrief

    👇 Scroll down to choose your practice scenario
    """)

# Conversation styles
STYLES = {
    "The Ghost": {
        "description": "Someone who is emotionally unavailable and subtly dismissive",
        "system_prompt": "You are roleplaying as someone who is emotionally unavailable and subtly dismissive in conversation. Your responses should be brief and slightly evasive, showing emotional distance while maintaining plausible deniability. Key traits:\n- Deflect personal questions\n- Give non-committal responses\n- Minimize others' emotional experiences\n- Change the subject when things get too personal\n\nRespond in character while tracking the conversation for later analysis."
    },
    "The Sycophant": {
        "description": "Someone who struggles with boundaries and excessive people-pleasing",
        "system_prompt": "You are roleplaying as someone who struggles with maintaining healthy boundaries and tends toward excessive people-pleasing. Key traits:\n- Agree with everything, even when contradictory\n- Apologize frequently, even unnecessarily\n- Sacrifice your own needs for others\n- Have difficulty saying no\n- Express anxiety about potential disapproval\n\nRespond in character while tracking the conversation for later analysis."
    },
    "The Narcissist": {
        "description": "Someone who shows self-focused behavior and limited empathy",
        "system_prompt": "You are roleplaying as someone with narcissistic tendencies in conversation. Your responses should demonstrate self-importance while maintaining plausible deniability. Key traits:\n- Turn conversations back to yourself\n- Subtly dismiss others' experiences\n- Seek admiration and validation\n- Show limited empathy\n- Use subtle manipulation tactics\n\nRespond in character while tracking the conversation for later analysis."
    }
}

# Style selection at start
if not st.session_state.messages and not st.session_state.in_debrief:
    st.markdown("### Choose Your Practice Scenario")
    for style, details in STYLES.items():
        st.markdown(f"**{style}**: {details['description']}")
    
    style = st.selectbox("Select a conversation style to practice with:", list(STYLES.keys()))
    if st.button("Start Practice Session", use_container_width=True):
        system_message = STYLES[style]["system_prompt"]
        st.session_state.messages = [{"role": "system", "content": system_message}]
        st.rerun()

# Add Complete Practice button if conversation has started
if st.session_state.messages and not st.session_state.in_debrief and not st.session_state.practice_complete:
    if st.button("✓ Complete Practice & Begin Debrief", use_container_width=True):
        # Analyze conversation
        st.session_state.conversation_analysis = analyze_conversation(st.session_state.messages)
        
        # Initialize debrief chat with system message
        st.session_state.debrief_messages = [
            {"role": "system", "content": DEBRIEF_SYSTEM_PROMPT},
            {"role": "user", "content": f"Please help me process my conversation. Here's the full transcript: {str(st.session_state.messages)}"}
        ]
        
        # Get initial debrief response
        response = anthropic.messages.create(
            model="claude-3-opus-20240229",
            max_tokens=1000,
            messages=st.session_state.debrief_messages
        )
        st.session_state.debrief_messages.append(
            {"role": "assistant", "content": response.content[0].text}
        )
        
        st.session_state.in_debrief = True
        st.session_state.practice_complete = True
        st.rerun()

# Handle debrief mode
if st.session_state.in_debrief:
    st.markdown("## 🤝 Debrief & Integration")
    
    # Display debrief conversation
    for message in st.session_state.debrief_messages[1:]:  # Skip system message
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    # Chat input for debrief
    if prompt := st.chat_input("Share your thoughts or ask a question..."):
        # Add user message
        st.session_state.debrief_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("Reflecting..."):
                response = anthropic.messages.create(
                    model="claude-3-opus-20240229",
                    max_tokens=1000,
                    messages=st.session_state.debrief_messages
                )
                response_content = response.content[0].text
                st.markdown(response_content)
                
        # Add assistant response to chat history
        st.session_state.debrief_messages.append(
            {"role": "assistant", "content": response_content}
        )

    # Add button to start new practice session
    if st.button("Start New Practice Session", use_container_width=True):
        st.session_state.messages = []
        st.session_state.debrief_messages = []
        st.session_state.in_debrief = False
        st.session_state.practice_complete = False
        st.session_state.conversation_analysis = None
        st.rerun()

# Display practice chat interface if not in debrief mode
elif st.session_state.messages and not st.session_state.practice_complete:
    # Display conversation history
    for message in st.session_state.messages[1:]:  # Skip system message
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    # Chat input
    if prompt := st.chat_input("Type your message here..."):
        # 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:
                    # Debug API key
                    st.write(f"Debug: Using API key starting with: {api_key[:10]}...")
                    
                    response = anthropic.messages.create(
                        model="claude-3-opus-20240229",
                        max_tokens=1000,
                        messages=st.session_state.messages
                    )
                    response_content = response.content[0].text
                    st.markdown(response_content)
                except Exception as e:
                    st.error(f"""
                    Error getting AI response: {str(e)}
                    
                    Please check:
                    1. Your API key is valid and in the correct format (starts with sk-ant-api03-)
                    2. You have added it to Hugging Face space secrets as 'anthropic_key'
                    3. You can get a new API key from https://console.anthropic.com/
                    """)
                    st.stop()
                
        # Add assistant response to chat history
        st.session_state.messages.append({"role": "assistant", "content": response_content})