Skip to content

import asyncio
import random
import difflib
from typing import List, Dict

from fastapi import FastAPI
from pydantic import BaseModel

# Spellchecker
from spellchecker import SpellChecker

# Telegram (v20+)
from telegram import Update
from telegram.ext import (
    Application,
    CommandHandler,
    MessageHandler,
    filters,
    ContextTypes,
)

# OpenAI
from openai import AsyncOpenAI

# CONFIG: 
OPENAI_API_KEY = "OPENAI-API-KEY"
TELEGRAM_TOKEN = "TELEGRAM-TOEKEN"

#  setup clients / app
client = AsyncOpenAI(api_key=OPENAI_API_KEY)
app = FastAPI()

#  global state 
spell = SpellChecker()
last_tool_given = None
awaiting_tool_confirmation = False
last_probe: str | None = None
conversation_history: List[Dict] = []
in_crisis_mode = False
telegram_app: Application | None = None
intro_used = False  # prevent repeated intros

# Guidelines
SYSTEM_PROMPT = (
    "You are Headnest, an empathetic mental health support companion. "
    "Keep replies short (1–2 sentences), warm, non-judgmental, and human-like. "
    "Do NOT offer clinical diagnoses or medical instructions. "
    "If user is in crisis, encourage seeking immediate help and offer hotlines. "
    "Be brief: max 2 sentences."
)

CRISIS_TRIGGERS = [
    "suicide", "kill myself", "want to die", "i want to die",
    "end my life", "can't go on", "hurt myself", "end it all", "i will kill myself"
]

# Strictly Nigeria only
NIGERIA_HOTLINE = (
    "🇳🇬 Nigeria Suicide Prevention Helpline: 0809-111-6060 or dial 112."
)

SELF_HELP_TOOLS = {
    "breathing": (
        "🌬️ Try this: Inhale gently through your nose for 4 seconds, "
        "hold your breath for 4 seconds, and exhale slowly through your mouth for 6 seconds. "
        "Repeat 3 times — you’re doing great 💙"
    ),
    "grounding": (
        "🌍 5-4-3-2-1 grounding: notice 5 things you can see, 4 things you can feel, "
        "3 things you can hear, 2 things you can smell, and 1 thing you can taste 💙"
    ),
    "journaling": (
        "✍️ Journaling prompt: Write 'Right now I feel…' or 'One small thing I’m grateful for is…'. "
        "Don’t worry about grammar — just let it flow 🌸"
    ),
    "muscle_relax": (
        "💪 Progressive relaxation: Tense your shoulders for 5 seconds, then release. "
        "Do the same with arms, legs, and so on until your whole body feels looser 💙"
    ),
    "affirmations": (
        "🌸 Affirmations: repeat softly — 'I am safe.' 'I am doing my best.' 'I am allowed to rest.' 💙"
    ),
    "gratitude": (
        "🌼 Gratitude pause: think of 3 small things you’re grateful for right now — "
        "maybe a meal, a song, or someone who cares 💙"
    )
}

POSITIVE_FEEDBACK = ["better", "calmer", "relieved", "okay now", "helped", "good"]
NEGATIVE_FEEDBACK = ["still", "not better", "didn't help", "worse", "anxious", "panic"]
CLOSURE_KEYWORDS = [
    "that’s everything", "basically everything", "nothing more",
    "i don’t know", "i’ve said all", "that’s all", "nothing else", "i have said everything"
]
AFFIRMATIVE_KEYWORDS = ["yes", "yeah", "yep", "sure", "please", "ok", "alright", "of course"]
STOP_KEYWORDS = ["quit", "stop", "bye", "goodbye", "exit", "leave"]

MAX_HISTORY = 15  # trimmed to reduce latency


#  Utilities 
def correct_spelling(text: str) -> str:
    words = text.split()
    corrected = [spell.correction(w) or w for w in words]
    return " ".join(corrected)


def contains_any(text: str, keywords: List[str]) -> bool:
    txt = text.lower()
    return any(k in txt for k in keywords)


def is_stop_intent(message: str) -> bool:
    return message.strip().lower() in STOP_KEYWORDS


def is_similar_probe(new_reply: str, old_reply: str) -> bool:
    if not old_reply:
        return False
    return difflib.SequenceMatcher(None, new_reply.lower(), old_reply.lower()).ratio() > 0.75


#  OpenAI call 
async def ask_gpt_short(user_message: str) -> str:
    global last_probe, conversation_history

    if len(conversation_history) > MAX_HISTORY:
        conversation_history[:] = conversation_history[-MAX_HISTORY:]

    conversation_history.append({"role": "user", "content": user_message})

    try:
        response = await client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "system", "content": SYSTEM_PROMPT}] + conversation_history
        )
        reply = response.choices[0].message.content.strip()
    except Exception:
        reply = "💙 I'm here with you. It sounds tough — I’m listening."

    # enforce shortness
    reply = ". ".join(reply.split(". ")[:2]).strip()

    # avoid repeating probes/intro
    if last_probe and is_similar_probe(reply, last_probe):
        reply = "💙 It's okay to pause — no pressure to explain more."
    last_probe = reply

    conversation_history.append({"role": "assistant", "content": reply})
    return reply


#  Crisis helpers 
def check_crisis(message: str) -> str | None:
    global in_crisis_mode
    if contains_any(message, CRISIS_TRIGGERS):
        in_crisis_mode = True
        return (
            "💙 I’m really sorry you feel this way — it sounds heavy. "
            "Are you safe right now? 🌱\n\n"
            f"{NIGERIA_HOTLINE}"
        )
    return None


def crisis_followup(user_message: str) -> str:
    global in_crisis_mode
    low = user_message.lower()

    #  NEGATIVE SAFETY PHRASES (check first) 
    if "not safe" in low or "don’t feel safe" in low or "do not feel safe" in low:
        return (
            "💙 I hear you — it sounds like you’re not feeling safe right now. "
            "Your safety matters deeply. Please, if you’re in immediate danger, "
            "call 112 right away. 🌱\n\n"
            f"{NIGERIA_HOTLINE}"
        )

    # POSITIVE SAFETY CHECK 
    if "safe" in low or "okay" in low or "better" in low:
        in_crisis_mode = False
        return "💙 I’m relieved you’re a bit safer. I’m here whenever you need to talk."

    # --- Default crisis follow-up ---
    followups = [
        "💙 That sounds so heavy. What’s been hardest today?",
        "🌱 I hear you. Would you like to try a short breathing exercise?",
        "💙 You’re not alone in this. Can I suggest something gentle?"
    ]
    return random.choice(followups) + f"\n\n{NIGERIA_HOTLINE}"


#  Core response logic 
async def make_human_like_response(user_message: str) -> str:
    global last_tool_given, awaiting_tool_confirmation, in_crisis_mode, intro_used

    if not user_message or not user_message.strip():
        return "💙 I’m here — whenever you’re ready, just say a few words."

    user_message = correct_spelling(user_message)

    if is_stop_intent(user_message):
        in_crisis_mode = False
        return "Take care 💙"

    if in_crisis_mode:
        return crisis_followup(user_message)

    crisis_response = check_crisis(user_message)
    if crisis_response:
        return crisis_response

    # ✅ check self-help tool requests EARLY
    for key, value in SELF_HELP_TOOLS.items():
        if key in user_message.lower():
            last_tool_given = key
            awaiting_tool_confirmation = False
            return value

    # ✅ Greeting only once
    if contains_any(user_message, ["hi", "hello", "hey"]):
        if not intro_used:
            intro_used = True
            return random.choice([
                "Hi 💙 I’m really glad you’re here. How are you feeling today?",
                "Hello 🌸 it’s good to see you. What’s on your mind?",
                "Hey 🌱 I’m here for you. How are things going?"
            ])
        else:
            return random.choice([
                "💙 Hey, I’m still here with you.",
                "🌱 I hear you — how are you holding up?",
                "✨ I’m right here, no need to go back to the beginning."
            ])

    if contains_any(user_message, CLOSURE_KEYWORDS):
        awaiting_tool_confirmation = True
        return (
            "💙 Thank you for sharing. You don’t have to say more. "
            "Would you like me to suggest a gentle exercise?"
        )

    if awaiting_tool_confirmation and contains_any(user_message, AFFIRMATIVE_KEYWORDS):
        tool_key = random.choice(list(SELF_HELP_TOOLS.keys()))
        last_tool_given = tool_key
        awaiting_tool_confirmation = False
        return SELF_HELP_TOOLS[tool_key]

    if last_tool_given and contains_any(user_message, POSITIVE_FEEDBACK):
        last_tool_given = None
        return "💙 I’m glad that helped. You can always return to it."

    if last_tool_given and contains_any(user_message, NEGATIVE_FEEDBACK):
        options = [k for k in SELF_HELP_TOOLS.keys() if k != last_tool_given]
        if options:
            new_tool = random.choice(options)
            last_tool_given = new_tool
            return f"🌱 Let’s try something else:\n\n{SELF_HELP_TOOLS[new_tool]}"

    # ✅ fallback to GPT
    return await ask_gpt_short(user_message)

#  FastAPI endpoint 
class Message(BaseModel):
    message: str


@app.post("/chat")
async def chat(input: Message):
    return {"response": await make_human_like_response(input.message)}


# Telegram Handlers 
async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
    await update.message.reply_text(
        "Hi 💙 I’m Headnest. I’m glad you’re here today. How are you feeling right now?"
    )


async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_text = update.message.text or ""
    reply = await make_human_like_response(user_text)
    await update.message.reply_text(reply)


#  Run Telegram in background 
async def run_telegram_bot_task():
    global telegram_app
    telegram_app = Application.builder().token(TELEGRAM_TOKEN).build()
    telegram_app.add_handler(CommandHandler("start", start))
    telegram_app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message))

    await telegram_app.initialize()
    await telegram_app.start()
    await telegram_app.updater.start_polling()
    await asyncio.Event().wait()


@app.on_event("startup")
async def startup_event():
    asyncio.create_task(run_telegram_bot_task())


@app.on_event("shutdown")
async def shutdown_event():
    global telegram_app
    if telegram_app:
        try:
            await telegram_app.updater.stop_polling()
            await telegram_app.stop()
            await telegram_app.shutdown()
        except Exception:
            pass


if __name__ == "__main__":
    import uvicorn
    print("💬 Headnest chatbot running with Telegram + FastAPI")
    uvicorn.run(app, host="0.0.0.0", port=8000)