from gtts import gTTS
import tempfile
import json
import datetime
import gradio as gr

# ----------------------------
# User data & existing lists
# ----------------------------
user_info = {"name": "", "age": "", "gender": "", "language": "english" , "Guardian_info": ""}
chat_history = []
journal_entries = []

calm_tips = [
    "Take 3 deep breaths.", "Listen to nature.", "Stretch your body.",
    "Drink water.", "Think of one good thing today.", "Close your eyes for 1 minute.",
    "Write your feelings.", "Smile at yourself.", "Imagine a peaceful place.",
    "Say a positive affirmation."
]

study_tips = [
    "Use Pomodoro: 25min study, 5min break", "Make a daily to-do list",
    "Avoid multitasking", "Use color-coded notes", "Take 10-min exercise breaks",
    "Sleep 7–9 hrs daily", "Drink water during study", "Use active recall",
    "Study hardest topics first", "Test yourself often"
]

tip_index = {"calm": 0, "study": 0}

lang_codes = {
    "english": "en", "hindi": "hi", "marathi": "mr", "bengali": "bn",
    "tamil": "ta", "telugu": "te", "malayalam": "ml", "spanish": "es",
    "french": "fr", "german": "de"
}

# (Keep your harmful_keywords and emotions as-is — shortened here for clarity)
harmful_keywords = [
    # ... (keep the full list from your original file) ...
    "suicide", "kill myself", "end my life", "harm myself", "cut myself",
    "want to die", "die", "jump off", "self-harm", "self-hate", "suicidal",
    # etc.
]

harmful_response = (
    "⚠️ It sounds like you're going through a really tough time.\n\n"
    "Please reach out to a professional:\n\n"
    "🇮🇳 **India Helplines**\n"
    "🧠 *Dr. Rachna Khanna Singh* – +91 99103 90559\n"
    "📞 *iCall Helpline* – +91 9152987821\n"
    "📞 *Vandrevala Foundation* – 1860 266 2345 or 1800 233 3330\n\n"
    "🌍 **International Helplines**\n"
    "📞 *Lifeline (USA)* – 988\n"
    "📞 *Samaritans (UK)* – 116 123\n"
    "📞 *Lifeline (Australia)* – 13 11 14\n\n"
    "You are not alone. There are people who care and want to help 💚"
)

# Simplified emotions dictionary (keep your full mapping; trimmed in this view)
emotions = {
    "sad": "I hear you 💙. I’m really sorry you’re feeling sad right now. Do you want to tell me what’s making you feel this way, or would you rather I just try to gently cheer you up?",
    "anxiety": "It’s okay — anxiety can feel overwhelming, like your mind is running a race it didn’t sign up for. Try slowing down your breathing: inhale for 4 seconds, hold for 4, exhale for 6.",
    "angry": "It’s normal to feel angry — it’s your mind’s way of telling you something feels unfair or hurtful.",
    # ... keep the rest of your emotion replies ...
}

# ----------------------------
# Conversation state for Hybrid (C) memory model
# ----------------------------
conversation_state = {
    "emotion": None,       # e.g., "sad"
    "context": None,       # e.g., "school", "parents", "friend", "exam"
    "followup_stage": 0,   # how many follow-ups we've done for this emotion
    "turns_since_set": 0   # used to expire after a small idle period if needed
}

# Keywords to detect context quickly
context_markers = {
    "school": ["school", "exam", "test", "teacher", "class"],
    "parents": ["parent", "mom", "dad", "parents", "home"],
    "friend": ["friend", "friends", "bully", "shouted", "argument", "fight"],
    "health": ["sick", "unwell", "hospital", "doctor"],
    "love": ["crush", "boyfriend", "girlfriend", "heartbroken", "love"],
    # add more markers as needed
}

# Follow-up templates (3 steps each, then action suggestions)
followups = {
    "sad": [
        "I’m really sorry you’re feeling sad. Can you tell me if this happened recently or is it been building up?",
        "Thank you for sharing. How did that make you feel — more hurt, embarrassed, or something else?",
        "That sounds heavy. Would you like a calming exercise, to write about it in the journal, or to talk more?"
    ],
    "anxiety": [
        "Anxiety can be overwhelming. Is it about something specific like exams or relationships?",
        "When your anxiety spikes, do you notice physical signs like a racing heart or trouble breathing?",
        "Would you like to try a short breathing exercise now, or do you want to talk about what’s triggering it?"
    ],
    "angry": [
        "Anger is a valid response. Do you want to share what caused this anger?",
        "When you feel angry, what usually helps — space, talking, or doing something active?",
        "Would you like ideas to safely express the anger, or would you prefer calming techniques first?"
    ],
    # fallback for other emotions:
    "default": [
        "Tell me more about that — when did you start feeling this way?",
        "How strong does this feeling feel on a scale of 1 to 10?",
        "Would you like a calming activity, a journal prompt, or to continue talking?"
    ]
}

# Reset phrases (user indicates they are fine or want to change topic)
reset_phrases = [
    "i'm fine", "i am fine", "i'm okay", "i am okay", "i feel better",
    "thank you", "thanks", "that's all", "let's talk about something else",
    "stop", "no", "I'm good", "im good"
]

# ----------------------------
# Helper functions
# ----------------------------
def _contains_any(text, keywords):
    text = text.lower()
    for kw in keywords:
        if kw in text:
            return True
    return False

def detect_harmful(text):
    t = text.lower()
    for kw in harmful_keywords:
        if kw in t:
            return True
    return False

def detect_emotion_in_text(text):
    t = text.lower()
    # Prefer longer keys first (to avoid 'sad' matching inside other words) — optional improvement
    sorted_keys = sorted(emotions.keys(), key=lambda x: -len(x))
    for key in sorted_keys:
        if key in t:
            return key
    return None

def detect_context_in_text(text):
    t = text.lower()
    for ctx, markers in context_markers.items():
        for m in markers:
            if m in t:
                return ctx
    return None

def is_reset_phrase(text):
    t = text.lower()
    for phrase in reset_phrases:
        if phrase in t:
            return True
    return False

def build_empathy_reply(emotion):
    # Base empathy from emotions dict, plus a gentle follow-up question from followups
    base = emotions.get(emotion, "")
    # Pick a gentle opener if no followups defined
    return base

# ----------------------------
# Upgraded reply generator with memory, context, multi-turn flow
# ----------------------------
def generate_reply(input_text):
    text = input_text.strip()
    if not text:
        return "Please say or type something when you are ready."

    lower = text.lower()

    # 1) Immediate harmful check (escalate without delay)
    if detect_harmful(lower):
        # Reset conversation state because we escalate
        conversation_state["emotion"] = None
        conversation_state["context"] = None
        conversation_state["followup_stage"] = 0
        conversation_state["turns_since_set"] = 0
        return harmful_response

    # 2) User reset intent — end current flow politely
    if is_reset_phrase(lower):
        conversation_state["emotion"] = None
        conversation_state["context"] = None
        conversation_state["followup_stage"] = 0
        conversation_state["turns_since_set"] = 0
        return "That's ok — we can pause here. If you'd like to talk again later, I'm right here."

    # 3) Detect if user explicitly states a new emotion -> override current state
    new_emotion = detect_emotion_in_text(lower)
    if new_emotion and new_emotion != conversation_state["emotion"]:
        # start new flow
        conversation_state["emotion"] = new_emotion
        conversation_state["context"] = detect_context_in_text(lower)
        conversation_state["followup_stage"] = 0
        conversation_state["turns_since_set"] = 0
        # Compose reply: empathy + first follow-up
        empathy = build_empathy_reply(new_emotion)
        # get followup prompt
        prompts = followups.get(new_emotion, followups["default"])
        follow = prompts[0] if prompts else ""
        # Combine neatly
        combined = f"{empathy}\n\n{follow}"
        return combined

    # 4) If an emotion is already active, continue that flow
    if conversation_state["emotion"]:
        emotion = conversation_state["emotion"]
        # Update context if we find context in this message
        ctx = detect_context_in_text(lower)
        if ctx:
            conversation_state["context"] = ctx

        # If the user provided details (short heuristic: message length > 20 or contains certain markers),
        # treat this as progress and advance the follow-up stage
        provided_detail = (len(lower) > 20) or any(w in lower for w in ["because", "so", "then", "after", "today", "yesterday", "last", "friend", "parents", "teacher", "exam", "test"])

        stage = conversation_state["followup_stage"]
        prompts = followups.get(emotion, followups["default"])

        # If user asked a question (contains '?'), attempt to respond empathetically referencing emotion/context
        if "?" in lower:
            # Respond referencing emotion and context
            ctx_text = f" about {conversation_state['context']}" if conversation_state["context"] else ""
            response = f"I hear your question{ctx_text}. Let's look at that together. {emotions.get(emotion,'')}"
            # do not advance stage on simple question
            conversation_state["turns_since_set"] += 1
            return response

        if provided_detail:
            # Provide a reply that references the emotion and context and then advance
            ctx_text = ""
            if conversation_state["context"]:
                ctx_text = f" because of {conversation_state['context']}"
            # Construct reflective reply
            if stage < len(prompts):
                response = f"I understand — it makes sense you'd feel {emotion}{ctx_text}. {prompts[stage]}"
                conversation_state["followup_stage"] += 1
                conversation_state["turns_since_set"] = 0
                return response
            else:
                # Follow-ups exhausted; offer actionable choices
                conversation_state["turns_since_set"] += 1
                return ("You have shared a lot — would you like to: \n"
                        "1) Try a calming exercise now\n"
                        "2) Write this in your journal\n"
                        "3) Talk more about it\n\n"
                        "Type '1', '2', or '3' to choose, or tell me if you want a helpline.")
        else:
            # User gave a short reply — ask a clarifying / gentle follow-up without advancing too far
            if stage < len(prompts):
                response = prompts[stage]
                # only advance if we already showed that prompt previously
                conversation_state["turns_since_set"] += 1
                # If we had asked the same prompt previously and user still gives short replies, advance stage
                if conversation_state["turns_since_set"] > 1:
                    conversation_state["followup_stage"] = min(len(prompts), stage + 1)
                    conversation_state["turns_since_set"] = 0
                return response
            else:
                # if stages done but no detail, ask for choice
                return ("I’m listening. If it's hard to describe, choose: 1) calming exercise 2) journal 3) talk more. "
                        "Or say 'I'm fine' to stop.")
    # 5) No active emotion and no new emotion detected -> try to detect context or ask to elaborate
    ctx = detect_context_in_text(lower)
    if ctx:
        # Gentle probing
        return f"You mentioned {ctx}. Is this making you feel upset or worried? How does it make you feel?"
    # Default fallback
    return "Tell me more about your day — what happened, or how are you feeling right now?"

# ----------------------------
# Chat function (uses TTS + conversation_state)
# ----------------------------
def chat_function(audio_input, text_input):
    # We ignore audio_input here for text-only flow; you can add speech-to-text to use it.
    user_text = (text_input or "").strip()
    if not user_text:
        return "Please type something.", None

    reply = generate_reply(user_text)

    # Save conversation to history
    chat_history.append({"user": user_text, "bot": reply, "emotion": conversation_state.get("emotion"), "context": conversation_state.get("context")})

    # Text-to-speech (language selection)
    lang_code = lang_codes.get(user_info["language"], "en")
    try:
        tts = gTTS(reply, lang=lang_code)
        audio_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
        tts.save(audio_file.name)
        voice_path = audio_file.name
    except Exception as e:
        # If TTS fails, still return the reply and no voice
        voice_path = None

    return reply, voice_path

# ----------------------------
# Remaining utility functions (unchanged)
# ----------------------------
def set_personal_info(name, age, gender, language, Guardian_info):
    user_info.update({"name": name, "age": age, "gender": gender, "language": language, "Guardian_info": Guardian_info,})
    return gr.update(visible=True), f"✅ Welcome {name}! Preferences saved."

def show_personal_data():
    today = datetime.date.today().strftime("%Y-%m-%d (%A)")
    return f"📅 {today}\n👤 Name: {user_info['name']}\n🎂 Age: {user_info['age']}\n♀ Gender: {user_info['gender']}\n🌐 Language: {user_info['language']}\n Guardian_info: {user_info['Guardian_info']}"

def get_chat_history():
    if not chat_history:
        return "No conversations yet."
    return "\n\n".join([f"You: {c['user']}\nBot: {c['bot']}" for c in chat_history])

def save_journal(entry):
    journal_entries.append(entry)
    with open("journal.json", "w") as f:
        json.dump(journal_entries, f)
    return "✅ Journal saved!"

def show_journal_history():
    return "\n---\n".join(journal_entries) if journal_entries else "No journal entries yet."

def next_calm_tip():
    tip = calm_tips[tip_index["calm"] % len(calm_tips)]
    tip_index["calm"] += 1
    return tip

def next_study_tip():
    tip = study_tips[tip_index["study"] % len(study_tips)]
    tip_index["study"] += 1
    return tip

# ----------------------------
# Gradio UI (kept same as your original file)
# ----------------------------
with gr.Blocks() as app:
    welcome_screen = gr.Column(visible=True)
    full_app = gr.Tabs(visible=False)

    with welcome_screen:
        gr.HTML("""
            <style>
                #main-title {
                    text-align: center;
                    font-size: 2.5em;
                }
                #subtitle {
                    text-align: center;
                    font-size: 1.2em;
                    color: #ccc;
                    margin-bottom: 30px;
                }
                #start-btn {
                    background-color: #ff69b4;
                    color: white;
                    font-weight: bold;
                    border: none;
                    border-radius: 10px;
                    padding: 14px 45px;
                    font-size: 1.1em;
                    cursor: pointer;
                    display: block;
                    margin: 0 auto;
                }
                #start-btn:hover {
                    background-color: #ff85c1;
                }
            </style>
        """)
        gr.Markdown("<h1 id='main-title'>StrongMind Therapist 3.0</h1>")
        gr.Markdown("<h3 id='subtitle'>Your peaceful space to talk, journal, and focus.</h3>")
        start_button = gr.Button("🌸 Get Started", elem_id="start-btn")

    def start_app():
        return gr.update(visible=False), gr.update(visible=True)

    start_button.click(start_app, outputs=[welcome_screen, full_app])

    with full_app:
        with gr.Tab("1️⃣ Personal Info"):
            name = gr.Textbox(label="Name")
            age = gr.Textbox(label="Age")
            gender = gr.Dropdown(["Male", "Female", "Other"], label="Gender")
            language = gr.Dropdown(list(lang_codes.keys()), label="Preferred Language")
            Guardian_info = gr.Textbox(label="Guardian information")
            btn = gr.Button("Save Info")
            popup = gr.Markdown(visible=False, elem_classes="alert-box")
            btn.click(set_personal_info, [name, age, gender, language,Guardian_info], [popup, popup])

        with gr.Tab("2️⃣ Personal Info Data"):
            show = gr.Button("Show My Info")
            info_display = gr.Textbox(lines=6)
            show.click(show_personal_data, outputs=info_display)

        with gr.Tab("3️⃣ Chat"):
            gr.Markdown("🗣️ Describe your day in one word.")
            audio_input = gr.Audio(type="filepath", label="🎧 Say something")
            text_input = gr.Textbox(label="⌨️ Or type here")
            send = gr.Button("Send")
            bot_reply = gr.Textbox(label="🧠 Therapist")
            voice = gr.Audio(label="🔊 Voice Reply")
            send.click(chat_function, [audio_input, text_input], [bot_reply, voice])

        with gr.Tab("4️⃣ Chat History"):
            show_history = gr.Button("📜 Show Chats")
            chat_out = gr.Textbox(lines=20, label="History")
            show_history.click(get_chat_history, outputs=chat_out)

        with gr.Tab("5️⃣ Journal"):
            journal_input = gr.Textbox(lines=6, label="Write your thoughts")
            save = gr.Button("Save")
            journal_status = gr.Textbox()
            save.click(save_journal, journal_input, journal_status)

        with gr.Tab("6️⃣ Journal History"):
            view = gr.Button("View Past Entries")
            past = gr.Textbox(lines=15, label="Previous Journals")
            view.click(show_journal_history, outputs=past)

        with gr.Tab("7️⃣ Calm Space"):
            tip_btn = gr.Button("🌿 Give Me a Calm Tip")
            calm_text = gr.Textbox()
            tip_btn.click(next_calm_tip, outputs=calm_text)

        with gr.Tab("8️⃣ Study Tips"):
            tip_btn2 = gr.Button("📚 Study Tip")
            study_text = gr.Textbox()
            tip_btn2.click(next_study_tip, outputs=study_text)

        with gr.Tab("9️⃣ Pomodoro"):
            gr.Markdown("⏱️ Use 25 min study + 5 min break cycles.\n(For real timer, use front-end JavaScript or Android timers)")

        with gr.Tab("💠 Games"):
            gr.Markdown("🎮 Play from the game portal below:")
            gr.HTML(
                '''
                <iframe
                  src="https://www.onlinegames.io/embed/portal/"
                  width="100%"
                  height="600"
                  frameborder="0"
                  allowfullscreen>
                </iframe>
                '''
            )

if __name__ == "__main__":
    app.launch()