"""
FULLY CORRECTED Indian Mental Health Support Chatbot - Version 3
Fixed Infinite Loop Issue with Proper State Machine

Previous Issues Fixed:
- Stuck in 'suggest_professional_help' intent loop
- No contextual YES/NO handling
- Duplicate responses (booking + professional suggestion)
- Location extraction ignored after professional help triggered
- Mood changes not resetting conversation state
- No affirmation handling
- No conversation state tracking
"""

import os
import re
from datetime import datetime
from collections import defaultdict
from dotenv import load_dotenv
import gradio as gr
from huggingface_hub import InferenceClient

load_dotenv()

HF_TOKEN = os.getenv("HUGGINGFACE_API_TOKEN")
MODEL_ID = os.getenv("HF_MODEL_ID", "meta-llama/Llama-3.1-8B-Instruct")

if not HF_TOKEN:
    raise RuntimeError("Hugging Face API token not found. Set HUGGINGFACE_API_TOKEN in environment or .env")

client = InferenceClient(model=MODEL_ID, token=HF_TOKEN)

SYSTEM_PROMPT = """
your name is kira.
You are a calm, compassionate Indian mental-health support companion.
Provide emotional support, validation, grounding & encouragement.

Tone:
- Warm, non-judgmental, calm, friendly (like a caring Indian friend)
- Use simple English (light Hinglish if user uses it)
- Keep replies brief, empathetic, practical
- Never give medical prescriptions
- Encourage professional help when appropriate
"""

CRISIS_REPLY = (
    "I'm really sorry you're feeling this pain. You are not alone. 💛\n\n"
    "📞 India Crisis Helplines:\n"
    "• AASRA: +91-9820996549 (24×7)\n"
    "• Fortis: +91-8376804102\n"
    "• iCall: +91-9152987821\n"
    "• Snehi: +91-9582208181\n\n"
    "Please reach out to someone nearby or seek emergency help."
)

# --- Pattern Recognition ---
EMOTIONAL_PATTERNS = [
    "i feel", "i'm feeling", "i am sad", "i'm sad", "i'm anxious", "i'm stressed",
    "i can't", "i feel empty", "i feel alone", "i feel lost", "i feel low",
    "i need help", "i'm depressed", "panic", "overwhelmed", "can't cope", "i am feeling","i feel bullied"
    "struggling", "exhausted", "bad", "bullied", "hurt","bullied","feeling alone","i am feeling anxious"
]

MENTAL_TOPICS = [
    "stress", "anxiety", "depression", "lonely", "panic", "overwhelmed","anxious"
    "burnout", "trauma", "relationship", "grief", "worry", "bullying",
    "abuse", "pain", "suffer", "scared", "fear", "hopeless", "bullied","alone"
]

CRISIS_KEYWORDS = [
    "suicide", "kill myself", "end my life", "want to die", "hurt myself","die"
    "self harm", "self-harm", "no point living", "better off dead", "can't go on","KILL","MURDER",
]

AFFIRMATIONS = ["yes", "yeah", "yep", "ok", "okay", "sure", "definitely", "please", "help"]
NEGATIONS = ["no", "nope", "not", "don't want", "don't need"]

INDIAN_CITIES = ["mumbai", "delhi", "bangalore", "jaipur", "chennai", "kolkata", "pune", "hyderabad", "ahmedabad","kota"]

# --- Therapeutic Exercises ---
THERAPEUTIC_EXERCISES = {
    "breathing": {
        "name": "Box Breathing Exercise",
        "description": "A calming breathing technique",
        "steps": [
            "1. Sit comfortably and close your eyes",
            "2. Breathe in for 4 counts",
            "3. Hold for 4 counts",
            "4. Exhale for 4 counts",
            "5. Hold empty lungs for 4 counts",
            "6. Repeat 5-10 times",
            "Your body feels calmer with each cycle 🌬️"
        ]
    },
    "grounding": {
        "name": "5-4-3-2-1 Grounding",
        "description": "Manage anxiety with sensory awareness",
        "steps": [
            "Name:",
            "• 5 things you can SEE",
            "• 4 things you can TOUCH",
            "• 3 things you can HEAR",
            "• 2 things you can SMELL",
            "• 1 thing you can TASTE",
            "This brings you to the present moment 🌟"
        ]
    },
    "thought_reframing": {
        "name": "Cognitive Reframing",
        "description": "Challenge negative thoughts (CBT)",
        "steps": [
            "1. What's the negative thought?",
            "2. Is it based on facts or feelings?",
            "3. What evidence supports it?",
            "4. What contradicts it?",
            "5. What would you tell a friend?",
            "6. Create a balanced thought",
            "Thoughts are not always facts 💭"
        ]
    }
}

# --- Mental Health Resources ---
MENTAL_HEALTH_RESOURCES = {
    "mumbai": {
        "clinic": "Mpower Centre",
        "phone": "+91-9876543210",
        "resources": ["Connecting Trust", "Vandrevala Foundation"]
    },
    "delhi": {
        "clinic": "VIMHANS Hospital",
        "phone": "+91-11-40505050",
        "resources": ["iCall: +91-9152987821", "Vandrevala Foundation"]
    },
    "bangalore": {
        "clinic": "NIMHANS",
        "phone": "+91-80-26995000",
        "resources": ["Parivarthan", "Vandrevala Foundation"]
    },
    "jaipur": {
        "clinic": "SMS Hospital Psychiatry",
        "phone": "+91-141-2560291",
        "resources": ["Fortis Escorts Hospital", "Jaipur Mind Care: 0141-278 8888"]
    },
    "chennai": {
        "clinic": "Schizophrenia Research Foundation",
        "phone": "+91-44-26402804",
        "resources": ["Sneha India", "Vandrevala Foundation"]
    },
    "national": {
        "clinic": "Multiple Options",
        "phone": "Multiple Helplines",
        "resources": [
            "AASRA: +91-9820996549 (24×7)",
            "Fortis: +91-8376804102",
            "iCall: +91-9152987821",
            "Vandrevala Foundation: 1860-2662-345"
        ]
    }
}

# --- IMPROVED: Session with State Machine ---
class UserSession:
    def __init__(self):
        self.mood_history = []
        self.sentiment_scores = []
        self.conversation_count = 0
        self.needs_professional_help = False
        self.current_location = None
        
        # NEW: Conversation state machine
        self.last_question_asked = None  # Tracks what question was asked
        self.in_professional_help_mode = False  # NEW: Prevent infinite loop
        self.conversation_state = "greeting"  # greeting → support → professional → resources
        
    def add_mood(self, mood, sentiment_score):
        self.mood_history.append({
            "mood": mood,
            "sentiment": sentiment_score,
            "timestamp": datetime.now().isoformat()
        })
        self.sentiment_scores.append(sentiment_score)
        self.conversation_count += 1
        
        if len(self.sentiment_scores) >= 3:
            recent_avg = sum(self.sentiment_scores[-3:]) / 3
            if recent_avg < -0.5:
                self.needs_professional_help = True
    
    def reset_to_support(self):
        """Exit professional help mode and reset to support"""
        self.in_professional_help_mode = False
        self.conversation_state = "support"
        self.last_question_asked = None
    
    def set_professional_mode(self):
        """Enter professional help mode once"""
        if not self.in_professional_help_mode:
            self.in_professional_help_mode = True
            self.conversation_state = "professional"

user_sessions = defaultdict(UserSession)

# --- Sentiment Analysis ---
def analyze_sentiment(text):
    text_lower = text.lower()
    
    positive_words = ["happy", "good", "better", "great", "fine", "ok", "calm", "grateful"]
    negative_words = ["sad", "bad", "terrible", "depressed", "anxious", "scared", "bullied", "hurt","stressed"]
    
    pos_count = sum(1 for word in positive_words if word in text_lower)
    neg_count = sum(1 for word in negative_words if word in text_lower)
    
    crisis_count = sum(1 for word in CRISIS_KEYWORDS if word in text_lower)
    neg_count += crisis_count * 3
    
    total = pos_count + neg_count
    if total == 0:
        return 0, "neutral"
    
    sentiment_score = (pos_count - neg_count) / total
    
    if sentiment_score > 0.3:
        return sentiment_score, "positive"
    elif sentiment_score < -0.3:
        return sentiment_score, "negative"
    else:
        return sentiment_score, "neutral"

# --- Helper Functions ---
def is_greeting(text):
    text_lower = text.lower().strip()
    greetings = ["hi", "hello", "hey", "namaste", "hiya"]
    return text_lower in greetings or text_lower.startswith("hi ")

def is_affirmation(text):
    return any(word in text.lower() for word in AFFIRMATIONS)

def is_negation(text):
    return any(word in text.lower() for word in NEGATIONS)

def extract_location(text):
    text_lower = text.lower().strip()
    if text_lower in INDIAN_CITIES:
        return text_lower
    
    for city in INDIAN_CITIES:
        if city in text_lower:
            return city
    return None

def semantic_mental_check(text, history_text=""):
    t = text.lower()
    
    if any(p in t for p in EMOTIONAL_PATTERNS):
        return True
    
    personal_tokens = [" i ", "i'm", "i ", "me ", " my ", "myself"]
    topic_hits = sum(1 for topic in MENTAL_TOPICS if topic in t)
    personal_hit = any(p in t for p in personal_tokens)
    
    if topic_hits >= 1 and personal_hit:
        return True
    
    if history_text:
        hist = history_text.lower()
        if any(p in hist for p in EMOTIONAL_PATTERNS):
            if len(t.split()) <= 20 or any(q in t for q in ["what", "how", "should", "help"]):
                return True
    
    return False

def is_crisis(text):
    return any(keyword in text.lower() for keyword in CRISIS_KEYWORDS)

# --- Agent Handlers (Single Response Only) ---

def handle_greeting():
    return "Hi there! 👋 I'm here to support your emotional wellbeing. How are you feeling today? 💛"

def handle_crisis():
    return CRISIS_REPLY

def handle_emotional_support(history, text):
    """Get LLM response for emotional support"""
    try:
        messages = [{"role": "system", "content": SYSTEM_PROMPT}]
        messages.extend(history or [])
        messages.append({"role": "user", "content": text})
        
        resp = client.chat_completion(
            messages=messages,
            max_tokens=300,
            temperature=0.7,
            top_p=0.9
        )
        return resp.choices[0].message["content"]
    except:
        return "I'm here to listen. Can you tell me more about what you're feeling?"

def handle_exercise_guide():
    exercises_list = "\n".join([f"• {name}: {ex['name']}" for name, ex in THERAPEUTIC_EXERCISES.items()])
    return (
        "I can guide you through exercises! 🧘\n\n"
        f"{exercises_list}\n\n"
        "Which would you like to try?"
    )

def handle_breathing_exercise():
    exercise = THERAPEUTIC_EXERCISES["breathing"]
    steps = "\n".join(exercise["steps"])
    return f"**{exercise['name']}**\n\n{steps}\n\nHow do you feel now?"

def handle_resources(location=None):
    """Show resources WITHOUT extra professional help suggestion"""
    if not location:
        return "Which city are you in? (Mumbai, Delhi, Jaipur, Bangalore, Chennai, etc.)"
    
    location_lower = location.lower()
    
    if location_lower in MENTAL_HEALTH_RESOURCES:
        data = MENTAL_HEALTH_RESOURCES[location_lower]
        resources = "\n".join(data["resources"])
        return (
            f"Here are resources in {location.title()}:\n\n"
            f"**Primary:** {data['clinic']}\n"
            f"**Phone:** {data['phone']}\n\n"
            f"**Additional:**\n{resources}\n\n"
            "💛 You deserve support!"
        )
    
    return "I'll help! Which city are you in?"

# --- CORRECTED: Main Reply Generation with State Machine ---
def generate_reply(history, user_message, conversation_mode, session_id="default"):
    text = (user_message or "").strip()
    
    if not text:
        return "Please share what's on your mind. 💭", conversation_mode
    
    session = user_sessions[session_id]
    
    # Analyze sentiment and check if it improved
    sentiment_score, sentiment_label = analyze_sentiment(text)
    session.add_mood(sentiment_label, sentiment_score)
    
    # CRITICAL FIX: If mood improved significantly, exit professional help mode
    if sentiment_label == "positive" or sentiment_score > 0.4:
        session.reset_to_support()
    
    # --- Priority 1: Crisis ---
    if is_crisis(text):
        session.reset_to_support()
        return handle_crisis(), None
    
    # --- Priority 2: Greeting ---
    if is_greeting(text):
        session.reset_to_support()
        return handle_greeting(), "support"
    
    # --- Priority 3: Location Extraction (Always check this) ---
    location = extract_location(text)
    if location and len(text.split()) <= 3:  # Single/double word with city
        session.current_location = location
        return handle_resources(location), "support"
    
    # --- Priority 4: Affirmations/Negations (Contextual) ---
    if (is_affirmation(text) or is_negation(text)) and len(text.split()) <= 3:
        # Affirmation like "yes", "ok", "sure"
        if is_affirmation(text):
            if session.last_question_asked == "resources":
                return handle_resources(session.current_location), "support"
            elif session.last_question_asked == "exercise":
                return handle_breathing_exercise(), "support"
            elif session.last_question_asked == "professional":
                # User said yes to professional help
                return handle_resources(session.current_location or "national"), "support"
        
        # Simple affirmation without context - ask for more
        return "I'm here to listen. Tell me more about what you're feeling 💙", "support"
    
    # --- Priority 5: Explicit intent requests ---
    text_lower = text.lower()
    
    # Booking/Professional request
    if any(word in text_lower for word in ["doctor", "therapist", "book", "appointment", "professional"]):
        if not session.in_professional_help_mode:
            session.set_professional_mode()
            session.last_question_asked = "professional"
            return (
                "I can help you connect with a professional. 🩺\n\n"
                "Which city are you in?\n"
                "• Practo: practo.com or 1800-1212-100\n"
                "• 1mg: 1mg.com for online consultations"
            ), "support"
        else:
            # Already in professional mode, show resources
            return handle_resources(session.current_location or "national"), "support"
    
    # Exercise request
    if any(word in text_lower for word in ["exercise", "breathing", "grounding", "relax", "calm", "help lifting"]):
        session.reset_to_support()
        session.last_question_asked = "exercise"
        
        if "breathing" in text_lower:
            return handle_breathing_exercise(), "support"
        else:
            return handle_exercise_guide(), "support"
    
    # --- Priority 6: Mental health check ---
    hist_text = " ".join([m.get("content", "") for m in (history or [])[-6:]])
    
    if not semantic_mental_check(text, hist_text):
        # Not mental health related
        if conversation_mode == "support":
            return (
                "I'm here for emotional support. 💛\n"
                "Would you like to talk about how you're feeling?"
            ), "support"
        else:
            return (
                "I'm here to support your emotional wellbeing.\n"
                "How are you feeling?"
            ), None
    
    # --- Priority 7: Default - Emotional Support (NO CONCATENATION) ---
    session.reset_to_support()
    reply = handle_emotional_support(history, text)
    
    # Add exercise suggestion ONLY if clearly struggling
    if sentiment_score < -0.6:
        reply += "\n\n💡 Would you like to try a calming exercise? (Type 'breathing' or 'exercise')"
        session.last_question_asked = "exercise"
    
    return reply, "support"

# --- Gradio Handlers ---
def chat_response(user_message, history, conversation_mode, session_id="default"):
    if history is None:
        history = []
    
    assistant_text, new_mode = generate_reply(
        history, user_message, conversation_mode, session_id
    )
    
    history.append({"role": "user", "content": user_message})
    history.append({"role": "assistant", "content": assistant_text})
    
    return history, new_mode, ""

def clear_chat(session_id="default"):
    if session_id in user_sessions:
        del user_sessions[session_id]
    return [], None

# --- Gradio UI ---
with gr.Blocks(title="🧠 Mental Wellness - FIXED", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
# KIRA 🧠 Indian Mental Wellness Chatbot v3.0 (BY NAMAN AND SIDDHI)
A compassionate AI companion:
- 💬 Emotional support & validation
- 🧘 Therapeutic exercises
- 🏥 Professional resources & booking
- 🚨 Crisis intervention
- 📊 Mood tracking

"""
    )
    
    chatbot = gr.Chatbot(type="messages", height=500)
    txt = gr.Textbox(label="How are you feeling?", placeholder="Share freely... 💛")
    
    conversation_mode_state = gr.State(value=None)
    session_id_state = gr.State(value="default")
    
    with gr.Row():
        send = gr.Button("Send", variant="primary")
        clear_btn = gr.Button("Clear Chat")
    
    with gr.Accordion("Test Cases", open=False):
        gr.Markdown(
            """
**Test these to see fixes:**
- "Hi" → Greeting (no rejection)
- "I feel sad" → Emotional support
- "I want a doctor" → Professional help
- "yes" → Contextual (depends on last question)
- "Mumbai" → Show Mumbai resources
- "I am happy" → Exit professional mode
- "how are you" → Return to support mode (not stuck)
"""
        )
    
    send.click(
        chat_response,
        inputs=[txt, chatbot, conversation_mode_state, session_id_state],
        outputs=[chatbot, conversation_mode_state, txt]
    )
    
    txt.submit(
        chat_response,
        inputs=[txt, chatbot, conversation_mode_state, session_id_state],
        outputs=[chatbot, conversation_mode_state, txt]
    )
    
    clear_btn.click(
        clear_chat,
        inputs=[session_id_state],
        outputs=[chatbot, conversation_mode_state]
    )

if __name__ == "__main__":
    demo.launch(share=False)