import streamlit as st

from dotenv import load_dotenv
import os
import random
from datetime import datetime

import warnings
warnings.filterwarnings("ignore", category=FutureWarning)

# ────────────────────────────────────────────────
# xAI SDK import
# ────────────────────────────────────────────────
from xai_sdk import Client
from xai_sdk.chat import system, user

# Load environment variables
load_dotenv(encoding="utf-8")
XAI_API_KEY = os.getenv("XAI_API_KEY")

# Page configuration
st.set_page_config(
    page_title="Mental Health Companion",
    page_icon="🌱",
    layout="centered"
)

# Custom CSS for better UI
st.markdown("""
    <style>
    .stChatMessage {
        padding: 1rem;
        border-radius: 0.5rem;
    }
    .crisis-warning {
        background-color: #fff3cd;
        border: 2px solid #ffc107;
        padding: 1rem;
        border-radius: 0.5rem;
        margin: 1rem 0;
    }
    </style>
""", unsafe_allow_html=True)

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

# Simple default mood (handled by Grok instead of transformers)
if "mood" not in st.session_state:
    st.session_state.mood = "neutral"
    st.session_state.confidence = 0.5

# Initialize xAI client
if XAI_API_KEY:
    try:
        client = Client(api_key=XAI_API_KEY)
    except Exception as e:
        st.error(f"Error initializing xAI client: {e}")
        st.stop()
else:
    st.error("⚠️ XAI_API_KEY not found. Please add it to your .env file.")
    st.stop()

# Enhanced relaxation tips with categories
relaxation_tips = {
    "breathing": [
        "Box breathing: Inhale for 4 seconds, hold for 4, exhale for 4, hold for 4. Repeat 3-5 times.",
        "4-7-8 breathing: Breathe in for 4 counts, hold for 7, exhale slowly for 8.",
    ],
    "physical": [
        "Try progressive muscle relaxation: Tense and relax each muscle group from toes to head.",
        "Do some gentle stretches or yoga poses for 5-10 minutes.",
        "Go for a short walk, even if it's just around your room or building.",
    ],
    "mindfulness": [
        "Practice the 5-4-3-2-1 grounding technique: Name 5 things you see, 4 you can touch, 3 you hear, 2 you smell, 1 you taste.",
        "Spend 5 minutes on mindful breathing, focusing only on your breath.",
        "Try a brief guided meditation using a free app or YouTube.",
    ],
    "creative": [
        "Journal your thoughts: Write freely for 5-10 minutes without editing.",
        "Listen to calming music or nature sounds.",
        "Color, draw, or engage in any creative activity you enjoy.",
    ],
    "social": [
        "Reach out to a friend or family member for a quick chat.",
        "Join a study group or student organization to connect with peers.",
        "Consider talking to a counselor at your school's wellness center.",
    ]
}

# Crisis keywords detection
crisis_keywords = [
    "suicide", "kill myself", "end it all", "want to die", 
    "no reason to live", "better off dead", "hurt myself"
]

def detect_crisis(text):
    """Detect potential crisis situations"""
    text_lower = text.lower()
    return any(keyword in text_lower for keyword in crisis_keywords)

def get_random_tips(num=2):
    """Get random relaxation tips from different categories"""
    tips = []
    categories = random.sample(list(relaxation_tips.keys()), min(num, len(relaxation_tips)))
    for category in categories:
        tips.append(random.choice(relaxation_tips[category]))
    return tips

# Header
st.title("🌱 Mental Health Companion")
st.markdown("*A supportive space for students navigating stress, anxiety, and life's challenges*")

# Sidebar with resources
with st.sidebar:
    st.header("📚 Resources")
    st.markdown("""
    ### Crisis Support (24/7)
    - 🇮🇳 **Vandrevala Foundation**: 1860-2662-345
    - 🇮🇳 **iCall**: 9152987821
    - 🌍 **International**: Find local helplines
    
    ### Student Wellness Tips
    - Maintain a regular sleep schedule
    - Take study breaks every 45-60 minutes
    - Stay connected with friends and family
    - Exercise regularly, even if just 15 minutes
    - Seek help early from campus counseling
    
    ### About This Bot
    This chatbot uses AI to provide emotional support 
    and coping strategies. It's not a replacement for 
    professional mental health care.
    """)
    
    if st.button("🔄 Clear Chat History"):
        st.session_state.messages = []
        st.rerun()

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

# User input
user_input = st.chat_input("Share what's on your mind...")

if user_input:
    # Check for crisis situation
    if detect_crisis(user_input):
        st.markdown("""
        <div class="crisis-warning">
        <h3>🆘 Immediate Support Available</h3>
        <p>I'm concerned about what you've shared. Please reach out to a crisis counselor immediately:</p>
        <ul>
            <li><strong>India - Vandrevala Foundation</strong>: 1860-2662-345 (24/7)</li>
            <li><strong>India - iCall</strong>: 9152987821 (Mon-Sat, 8am-10pm)</li>
            <li><strong>Your campus counseling center</strong></li>
        </ul>
        <p>You don't have to face this alone. Professional help is available and things can get better.</p>
        </div>
        """, unsafe_allow_html=True)
        st.stop()
    
    # Add user message to history & display
    st.session_state.messages.append({"role": "user", "content": user_input})
    with st.chat_message("user"):
        st.markdown(user_input)
    

    # Default mood (sentiment handled implicitly)
    mood = st.session_state.get("mood", "neutral")
    confidence = st.session_state.get("confidence", 0.5)


    
    # Prepare prompt (now as messages list)
    system_prompt = f"""
You are a compassionate mental health support chatbot for students.

Guidelines:
1. Show empathy and validate their feelings
2. Be warm, supportive, and non-judgmental
3. Use encouraging language appropriate for students
4. If the mood seems negative, acknowledge their struggle
5. Offer hope and perspective when appropriate
6. Keep responses conversational and under 120 words
7. Do NOT provide medical diagnosis or treatment
8. If issues seem serious, gently suggest professional support
9. Use "I" statements to make it personal (e.g., "I hear you")

Detected emotional tone: {mood} (confidence: {confidence:.2f})
    """.strip()

    messages = [
        system(system_prompt),
        user(user_input)
    ]

    # Generate response
    with st.spinner("💭 Thinking..."):
        try:
            chat = client.chat.create(model="grok-4")          # or "grok-beta", "grok-3", etc.
            for msg in messages:
                chat.append(msg)

            response = chat.sample()
            bot_response = response.content.strip()

            # Add practical tips for negative sentiment
            if mood == 'negative' and confidence > 0.6:
                tips = get_random_tips(2)
                bot_response += "\n\n**Quick coping strategies to try:**\n"
                for i, tip in enumerate(tips, 1):
                    bot_response += f"{i}. {tip}\n"

            # Add timestamp (optional)
            timestamp = datetime.now().strftime("%I:%M %p")
            # bot_response += f"\n\n<small>{timestamp}</small>"

        except Exception:
              bot_response = (
                   "⚠️ I'm temporarily unable to connect to the AI service right now.\n\n"
                   "That said, I still want to support you. If you're feeling overwhelmed, "
                   "consider reaching out to someone you trust or your campus counseling services.\n\n"
                   "You’re not alone, and help is available 💙"
        )

    
    # Add bot response to history & display
    st.session_state.messages.append({"role": "assistant", "content": bot_response})
    with st.chat_message("assistant"):
        st.markdown(bot_response)

# Footer
st.markdown("---")
st.caption(
    "💙 Remember: This chatbot is a support tool, not a replacement for "
    "professional mental health care. If you're in crisis, please contact "
    "emergency services or a crisis hotline immediately."
)