import gradio as gr
import json
import datetime
from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, tool
from typing import Dict, List, Any
import re

# Import your existing tools (assuming they're in the same file or properly imported)
# For this example, I'll include the key functions inline

@tool
def speech_emotion_analyzer(text: str, context: str = "therapy") -> str:
    """
    Analyzes patient's emotional state from speech text.
    
    Args:
        text: Patient's speech text
        context: Context (therapy, assessment, crisis)
    
    Returns:
        Analysis of emotional state and recommendations
    """
    if not text or len(text.strip()) < 10:
        return "Error: Text must be at least 10 characters long for analysis."
    
    # Emotion indicators for analysis
    emotion_indicators = {
        "depression": {
            "keywords": ["sad", "hopeless", "tired", "meaningless", "empty", "lonely", 
                        "depressed", "despair", "miserable", "worthless", "numb"],
            "patterns": [r"can't.*sleep", r"no.*energy", r"everything.*bad", r"don't.*want", 
                        r"life.*no.*meaning", r"better.*dead"]
        },
        "anxiety": {
            "keywords": ["anxious", "worried", "nervous", "scared", "panic", "tense",
                        "anxiety", "worry", "fear", "stress", "phobia"],
            "patterns": [r"what.*if", r"afraid.*that", r"can't.*cope", r"heart.*racing",
                        r"hands.*shaking", r"can't.*breathe", r"panic.*attack"]
        },
        "anger": {
            "keywords": ["angry", "furious", "irritated", "annoyed", "hate", "mad",
                        "rage", "aggression", "hatred", "frustration"],
            "patterns": [r"fed up", r"sick of", r"can't.*stand", r"want.*hit",
                        r"everything.*irritates", r"pissed off"]
        },
        "stress": {
            "keywords": ["stress", "pressure", "overwhelmed", "exhausted", "burned out",
                        "burnout", "fatigue","job interview", "drained"],
            "patterns": [r"too.*much.*work", r"can't.*keep up", r"headache", r"no.*time",
                        r"swamped", r"working.*nonstop",r"job.*interview"]
        },
        "trauma": {
            "keywords": ["trauma", "flashbacks", "nightmares", "memories", "avoiding",
                        "ptsd", "triggered", "shock"],
            "patterns": [r"keep.*remembering", r"bad.*dreams", r"can't.*forget",
                        r"avoid.*places", r"triggers"]
        }
    }
    
    # Analyze text
    text_lower = text.lower()
    detected_emotions = {}
    
    for emotion, indicators in emotion_indicators.items():
        score = 0
        matches = []
        
        # Check keywords
        for keyword in indicators["keywords"]:
            if keyword in text_lower:
                score += 1
                matches.append(keyword)
        
        # Check patterns
        for pattern in indicators["patterns"]:
            if re.search(pattern, text_lower):
                score += 2
                matches.append(f"pattern: {pattern}")
        
        if score > 0:
            detected_emotions[emotion] = {"score": score, "matches": matches}
    
    # Check for crisis indicators
    crisis_indicators = ["suicide", "kill myself", "end it all", "not worth living", "better off dead"]
    crisis_risk = any(indicator in text_lower for indicator in crisis_indicators)
    
    # Check positive indicators
    positive_indicators = ["happy", "joy", "good", "great", "wonderful", "excited", "grateful", "love"]
    positive_count = sum(1 for indicator in positive_indicators if indicator in text_lower)
    
    # Create report
    report = "## 🧠 Emotional State Analysis\n\n"
    
    if detected_emotions:
        report += "### 📊 Detected Emotional States:\n"
        for emotion, data in sorted(detected_emotions.items(), key=lambda x: x[1]["score"], reverse=True):
            emotion_names = {
                "depression": "😔 Depression signs",
                "anxiety": "😰 Anxiety", 
                "anger": "😠 Anger/Irritation",
                "stress": "😵 Stress",
                "trauma": "💔 Trauma-related"
            }
            intensity = "High" if data['score'] >= 4 else "Medium" if data['score'] >= 2 else "Low"
            report += f"- {emotion_names[emotion]} (intensity: {intensity})\n"
            report += f"  Found indicators: {', '.join(data['matches'][:3])}\n\n"
    else:
        report += "✅ No clear signs of emotional distress detected.\n\n"
    
    # Positive aspects
    if positive_count > 0:
        report += f"😊 Positive indicators: Found {positive_count} positive marker(s)\n\n"
    
    # Crisis assessment
    if crisis_risk:
        report += "🚨 CRITICAL: Suicide risk indicators detected!\n"
        report += "📞 Immediate consultation and safety assessment recommended.\n"
        report += "Emergency services: 101, Crisis hotline: \n\n"
    
    # Treatment recommendations
    report += "### 🎯 Treatment Recommendations:\n"
    
    if "depression" in detected_emotions:
        report += "- 🧘 CBT (Cognitive Behavioral Therapy) for negative thought patterns\n"
        report += "- 💊 Consider psychiatric consultation for medication support\n"
        report += "- 🏃 Behavioral activation (increase pleasant activities)\n"
    
    if "anxiety" in detected_emotions:
        report += "- 🫁 Relaxation techniques: breathing exercises, progressive muscle relaxation\n"
        report += "- 🎭 Exposure therapy for phobic avoidance\n"
        report += "- 🧠 Mindfulness therapy for anxious thoughts\n"
    
    if "anger" in detected_emotions:
        report += "- 🎯 Anger management training and self-regulation techniques\n"
        report += "- 🔍 Trigger identification and behavior patterns\n"
        report += "- 🤝 Communication skills training\n"
    
    if "stress" in detected_emotions:
        report += "- ⏰ Stress management and time management techniques\n"
        report += "- ⚖️ Work-life balance assessment and adjustment\n"
        report += "- 🛡️ Coping strategies for stressful situations\n"
    
    if "trauma" in detected_emotions:
        report += "- 👁️ EMDR therapy for traumatic memories\n"
        report += "- 🏠 Trauma-informed approach in therapy\n"
        report += "- 🧘 Somatic techniques for body-based trauma symptoms\n"
    
    if not detected_emotions or positive_count > 2:
        report += "- 💪 Supportive therapy to maintain resources\n"
        report += "- 🎯 Personal growth and life goals development\n"
    
    report += f"\n📅 Analysis date: {datetime.datetime.now().strftime('%m/%d/%Y %H:%M')}"
    report += f"\n🌍 Context: {context}"
    
    return report

@tool
def create_therapy_worksheet(therapy_type: str, client_age: int, issue: str) -> str:
    """
    Creates a therapeutic worksheet/exercise for client.
    
    Args:
        therapy_type: Type of therapy (cbt, dbt, art, mindfulness)
        client_age: Client's age
        issue: Main problem
    
    Returns:
        Structured therapeutic worksheet
    """
    if client_age < 5 or client_age > 100:
        return "Error: Age must be between 5 and 100."
    
    therapy_type = therapy_type.lower()
    valid_types = ["cbt", "dbt", "art", "mindfulness"]
    
    if therapy_type not in valid_types:
        return f"Error: Therapy type must be one of: {', '.join(valid_types)}"
    
    # Determine level by age
    if client_age < 12:
        level = "child"
    elif client_age < 18:
        level = "teen" 
    else:
        level = "adult"
    
    worksheets = {
        "cbt": {
            "child": {
                "title": "🌈 My Thoughts and Feelings",
                "exercises": [
                    "🌤️ Mood Diary: Draw or mark your mood with an emoji each day",
                    "🤔 Thought Catcher: Write down thoughts when sad or scared",
                    "🎯 Helper Thoughts: Find 3 good thoughts for each bad thought",
                    "⭐ My Successes: Write one good thing you did each day"
                ]
            },
            "teen": {
                "title": "🧠 Cognitive Restructuring",
                "exercises": [
                    "📊 ABC Model: Situation (A) → Thought (B) → Feeling/Behavior (C)",
                    "🔍 Thought Detective: Find evidence FOR and AGAINST your beliefs",
                    "⚖️ Alternative Thoughts: Find 2-3 realistic alternatives for negative thoughts",
                    "📝 Thought Diary: Daily record of automatic thoughts"
                ]
            },
            "adult": {
                "title": "🔄 Cognitive Behavioral Therapy",
                "exercises": [
                    "🔄 Thought-Feeling-Behavior Cycle: Analyze connections in situations",
                    "⚡ Cognitive Distortions: Identify and correct thinking errors",
                    "🎯 Behavioral Experiments: Test catastrophic beliefs",
                    "📈 Graded Tasks: Step-by-step overcoming avoidance"
                ]
            }
        },
        "dbt": {
            "teen": {
                "title": "🌊 Emotion Regulation Skills",
                "exercises": [
                    "🌊 TIPP: Temperature, Intense exercise, Paced breathing, Paired muscle relaxation",
                    "🚦 STOP: Stop → Take a breath → Observe → Proceed mindfully",
                    "🎭 Emotion Diary: Track emotion intensity (1-10) with triggers",
                    "🧘 Mindfulness: 5-10 minutes daily breathing meditation"
                ]
            },
            "adult": {
                "title": "⚖️ Dialectical Behavior Therapy",
                "exercises": [
                    "⚖️ Wise Mind: Balance emotional and rational mind",
                    "🛡️ Crisis Skills: TIPP, distraction, self-soothing",
                    "🤝 DEAR MAN: Technique for healthy communication",
                    "🎯 Radical Acceptance: Practice accepting unchangeable reality"
                ]
            }
        },
        "mindfulness": {
            "child": {
                "title": "🐸 Mindfulness for Kids",
                "exercises": [
                    "🐸 Frog Pose: Sit still like a frog and listen to sounds",
                    "🌬️ Belly Breathing: Watch a toy rise and fall on your belly",
                    "👀 5-4-3-2-1: Find 5 things you see, 4 you hear, 3 you can touch...",
                    "🍯 Slow Eating: Eat one raisin very slowly and mindfully"
                ]
            },
            "teen": {
                "title": "🧘 Teen Mindfulness",
                "exercises": [
                    "📱 Digital Detox: 30-60 minutes without devices",
                    "🚶 Mindful Walking: Walk slowly noticing each step",
                    "💭 Thought Watching: Watch thoughts like clouds passing",
                    "❤️ Body Scan: Relax and notice each body part"
                ]
            },
            "adult": {
                "title": "🧘‍♀️ Mindfulness Practice",
                "exercises": [
                    "🧘 Formal Meditation: 20-30 minutes daily sitting practice",
                    "🍽️ Mindful Eating: Full attention to taste, smell, texture",
                    "💼 Work Mindfulness: Regular pauses to return to present",
                    "😴 Sleep Meditation: Body scan and breathing for better sleep"
                ]
            }
        }
    }
    
    # Get appropriate exercises
    therapy_data = worksheets.get(therapy_type, {})
    level_data = therapy_data.get(level, therapy_data.get("adult", {}))
    
    if not level_data:
        return f"Exercises for {therapy_type} therapy and age {client_age} not yet available."
    
    # Create worksheet
    worksheet = f"# {level_data['title']}\n\n"
    worksheet += f"**👤 Client**: {client_age} years old\n"
    worksheet += f"**🎯 Main issue**: {issue}\n"
    worksheet += f"**🔬 Therapy type**: {therapy_type.upper()}\n"
    worksheet += f"**📅 Date created**: {datetime.datetime.now().strftime('%m/%d/%Y')}\n\n"
    
    worksheet += "## 📋 Exercises:\n\n"
    
    for i, exercise in enumerate(level_data['exercises'], 1):
        worksheet += f"### {i}. {exercise}\n\n"
        
        # Add space for notes
        worksheet += "**📝 Notes:**\n"
        worksheet += "```\n"
        worksheet += "Date completed: ___________\n"
        worksheet += "Observations: \n\n\n"
        worksheet += "Emotions during exercise: \n\n"
        worksheet += "Insights: \n\n"
        worksheet += "```\n\n"
    
    worksheet += "---\n\n"
    worksheet += "## 📖 Instructions:\n"
    worksheet += "✅ Complete exercises at your own pace\n"
    worksheet += "✅ Record all observations - this is important for progress\n"
    worksheet += "✅ Discuss results with your therapist\n"
    worksheet += "✅ Stop if you feel too uncomfortable and seek support\n\n"
    
    worksheet += "## 🆘 When to get help:\n"
    worksheet += "⚠️ If exercises cause severe anxiety or panic\n"
    worksheet += "⚠️ If you have thoughts of self-harm\n"
    worksheet += "⚠️ If symptoms worsen\n\n"
    
    worksheet += "📞 Emergency contacts:\n"
    worksheet += "- Crisis hotline: 988\n"
    worksheet += "- Emergency: 101\n"
    worksheet += "- Your therapist: ________________\n\n"
    
    return worksheet

def mood_tracker_analyzer(mood_data_str: str) -> str:
    """
    Analyzes mood tracking data to identify patterns.
    """
    try:
        # Try to parse as JSON
        data = json.loads(mood_data_str)
    except:
        # If not JSON, try to parse simple format like "5,7,3,8,6"
        try:
            moods = [float(x.strip()) for x in mood_data_str.split(',')]
            data = [{"mood": mood, "date": f"Day {i+1}"} for i, mood in enumerate(moods)]
        except:
            return "Error: Please provide mood data as JSON or comma-separated numbers (e.g., '5,7,3,8,6')"
    
    if not isinstance(data, list) or len(data) == 0:
        return "Error: Mood data should be a non-empty list of entries."
    
    # Analyze mood patterns
    moods = []
    dates = []
    
    for entry in data:
        if isinstance(entry, dict) and 'mood' in entry:
            moods.append(float(entry['mood']))
            dates.append(entry.get('date', f'Entry {len(dates)+1}'))
        elif isinstance(entry, (int, float)):
            moods.append(float(entry))
            dates.append(f'Entry {len(dates)+1}')
    
    if not moods:
        return "Error: No valid mood entries found."
    
    # Calculate statistics
    avg_mood = sum(moods) / len(moods)
    mood_trend = "stable"
    
    if len(moods) > 3:
        recent_avg = sum(moods[-3:]) / 3
        older_avg = sum(moods[:-3]) / (len(moods) - 3)
        
        if recent_avg > older_avg + 0.5:
            mood_trend = "improving"
        elif recent_avg < older_avg - 0.5:
            mood_trend = "declining"
    
    # Create report
    report = "# 📊 Mood Tracking Analysis\n\n"
    report += f"**Period**: {dates[0]} to {dates[-1]}\n"
    report += f"**Entries analyzed**: {len(moods)}\n"
    report += f"**Average mood**: {avg_mood:.1f}/10\n"
    report += f"**Trend**: {mood_trend}\n\n"
    
    report += "## 📈 Pattern Analysis:\n"
    
    # Check for low mood periods
    low_mood_days = sum(1 for m in moods if m < 5)
    if low_mood_days > len(moods) * 0.3:
        report += "⚠️ Concern: Low mood detected in >30% of entries\n"
        report += "   → Consider professional evaluation for depression\n\n"
    
    # Check volatility
    if len(moods) > 2:
        volatility = sum(abs(moods[i] - moods[i-1]) for i in range(1, len(moods))) / (len(moods) - 1)
        if volatility > 2:
            report += "🎢 High mood volatility detected\n"
            report += "   → May benefit from emotion regulation skills\n\n"
    
    report += "## 💡 Recommendations:\n"
    
    if mood_trend == "declining":
        report += "- 🚨 Schedule appointment with mental health professional\n"
        report += "- 🧘 Increase self-care activities\n"
        report += "- 👥 Reach out to support network\n"
    elif mood_trend == "improving":
        report += "- ✅ Continue current strategies\n"
        report += "- 📝 Document what's working well\n"
        report += "- 🎯 Set new wellness goals\n"
    else:
        report += "- 📊 Continue tracking for more data\n"
        report += "- 🔍 Look for mood triggers\n"
        report += "- 💪 Maintain healthy routines\n"
    
    report += f"\n*Analysis completed: {datetime.datetime.now().strftime('%m/%d/%Y %H:%M')}*"
    
    return report

# Create the Gradio interface
def create_interface():
    """Create and return the Gradio interface"""
    
    with gr.Blocks(title="🧠 Mental Health Therapy Assistant", theme=gr.themes.Soft()) as app:
        gr.Markdown("""
        # 🧠 Mental Health Therapy Assistant
        
        **⚠️ Important Disclaimer**: This tool is for educational and supportive purposes only. 
        It is not a substitute for professional mental health care. If you're in crisis, 
        please contact emergency services (911) or the crisis hotline (988).
        """)
        
        with gr.Tabs():
            # Emotion Analysis Tab
            with gr.TabItem("🎭 Emotion Analysis"):
                gr.Markdown("### Analyze emotional content from speech or text")
                
                with gr.Row():
                    with gr.Column():
                        emotion_text = gr.Textbox(
                            label="Patient Speech/Text",
                            placeholder="Enter the text you want to analyze for emotional content...",
                            lines=5
                        )
                        emotion_context = gr.Dropdown(
                            choices=["therapy", "assessment", "crisis"],
                            value="therapy",
                            label="Context"
                        )
                        emotion_btn = gr.Button("🔍 Analyze Emotions", variant="primary")
                    
                    with gr.Column():
                        emotion_output = gr.Markdown(label="Analysis Results")
                
                emotion_btn.click(
                    fn=speech_emotion_analyzer,
                    inputs=[emotion_text, emotion_context],
                    outputs=emotion_output
                )
            
            # Worksheet Generator Tab
            with gr.TabItem("📝 Therapy Worksheets"):
                gr.Markdown("### Generate age-appropriate therapeutic exercises")
                
                with gr.Row():
                    with gr.Column():
                        worksheet_therapy = gr.Dropdown(
                            choices=["CBT", "DBT", "Art", "Mindfulness"],
                            value="CBT",
                            label="Therapy Type"
                        )
                        worksheet_age = gr.Slider(
                            minimum=5,
                            maximum=100,
                            value=25,
                            step=1,
                            label="Client Age"
                        )
                        worksheet_issue = gr.Textbox(
                            label="Main Issue/Concern",
                            placeholder="e.g., anxiety, depression, anger management..."
                        )
                        worksheet_btn = gr.Button("📋 Generate Worksheet", variant="primary")
                    
                    with gr.Column():
                        worksheet_output = gr.Markdown(label="Generated Worksheet")
                
                worksheet_btn.click(
                    fn=create_therapy_worksheet,
                    inputs=[worksheet_therapy, worksheet_age, worksheet_issue],
                    outputs=worksheet_output
                )
            
            # Mood Tracker Tab
            with gr.TabItem("📊 Mood Analysis"):
                gr.Markdown("### Analyze mood tracking patterns")
                
                with gr.Row():
                    with gr.Column():
                        mood_data = gr.Textbox(
                            label="Mood Data",
                            placeholder="Enter mood scores (1-10) separated by commas: 5,7,3,8,6\nOr JSON format: [{\"mood\": 5, \"date\": \"2024-01-01\"}, ...]",
                            lines=5
                        )
                        mood_btn = gr.Button("📈 Analyze Mood Patterns", variant="primary")
                        
                        gr.Markdown("""
                        **Example formats:**
                        - Simple: `5,7,3,8,6,4,9`
                        - JSON: `[{"mood": 5, "date": "2024-01-01"}, {"mood": 7, "date": "2024-01-02"}]`
                        """)
                    
                    with gr.Column():
                        mood_output = gr.Markdown(label="Mood Analysis Results")
                
                mood_btn.click(
                    fn=mood_tracker_analyzer,
                    inputs=mood_data,
                    outputs=mood_output
                )
            
            # Resources Tab
            with gr.TabItem("📚 Resources"):
                gr.Markdown("""
                ### 🆘 Crisis Resources
                
                **Immediate Help:**
                - 🚨 Emergency: **911**
                - 📞 Crisis Hotline: **988** (Suicide & Crisis Lifeline)
                - 💬 Crisis Text Line: Text HOME to **741741**
                
                **Online Resources:**
                - [National Alliance on Mental Illness (NAMI)](https://nami.org)
                - [Mental Health America](https://www.mhanational.org)
                - [Psychology Today Therapist Finder](https://www.psychologytoday.com)
                
                ### 🧘 Self-Care Tips
                
                **Daily Practices:**
                - 🌅 Maintain regular sleep schedule
                - 🏃 Regular physical activity
                - 🥗 Balanced nutrition
                - 🧘 Mindfulness/meditation practice
                - 👥 Social connections
                - 📚 Limit news/social media if overwhelming
                
                **Remember:** This tool provides educational support but cannot replace professional mental health care.
                """)
    
    return app

# Main execution
if __name__ == "__main__":
    # Create and launch the interface
    app = create_interface()
    
    # Launch with appropriate settings for deployment
    app.launch(
        server_name="0.0.0.0",  # Allow external access
        server_port=7860,       # Standard port for Hugging Face Spaces
        share=False,            # Don't create a public link
        debug=False             # Disable debug mode for production
    )