import gradio as gr
from huggingface_hub import InferenceClient
from sentence_transformers import SentenceTransformer, util

client = InferenceClient("Qwen/Qwen2.5-7B-Instruct")

with open("knowledge.txt", "r", encoding="utf-8") as file:
    knowledge_text = file.read()
# print(knowledge_text)

def preprocess_text(text):
    cleaned_text = text.strip()

    chunks = cleaned_text.split("\n")

    cleaned_chunks = []

    for chunk in chunks:
        stripped_chunk = chunk.strip()
        if len(stripped_chunk) > 0:
            cleaned_chunks.append(stripped_chunk)
    # print(cleaned_chunks)
    # print(len(cleaned_chunks))
    return cleaned_chunks
cleaned_chunks = preprocess_text(knowledge_text)


model = SentenceTransformer('all-MiniLM-L6-v2')

def create_embeddings(text_chunks):
    chunk_embeddings = model.encode(text_chunks, convert_to_tensor=True)
    # print(chunk_embeddings)
    # print(chunk_embeddings.shape)
    return chunk_embeddings
chunk_embeddings = create_embeddings(cleaned_chunks)

def get_top_chunk(query):
    query_embedding = model.encode(query, convert_to_tensor=True)

    similarities = util.cos_sim(query_embedding, chunk_embeddings)

    best_chunk_index = similarities.argmax()

    best_chunk = cleaned_chunks[best_chunk_index]

    return best_chunk

mental_health_keywords = [ # we implement a list of keywords to ensure chatbot only assists with mental health
    "stress", "stressed", "anxiety", "anxious",
    "sad", "sadness", "depressed", "depression",
    "lonely", "alone", "angry", "anger",
    "overwhelmed", "burnout", "burned out",
    "school", "grades", "homework", "test", "exam",
    "family", "parents", "siblings",
    "friend", "friends", "friendship",
    "relationship", "breakup",
    "sleep", "tired", "exhausted",
    "panic", "worried", "worry",
    "mental", "emotion", "emotional",
    "feel", "feeling", "feelings",
    "coping", "calm",
    "therapy", "therapist", "counselor",
    "motivation", "confidence", "self-esteem",
    "crying", "numb", "overthinking",
    "focus", "concentrate", "concentration"
]

def detect_mood(message):
    text = message.lower()

    if any(word in text for word in ["stress", "stressed", "overwhelmed", "burnout", "pressure", "finals", "exam"]):
        return "😟 Stressed / Overwhelmed"
    elif any(word in text for word in ["anxiety", "anxious", "panic", "worried", "worry", "nervous", "fear", "scared"]):
        return "😰 Anxious"
    elif any(word in text for word in ["sad", "lonely", "depressed", "crying", "numb", "hopeless", "down", "upset"]):
        return "😢 Sad / Low"
    elif any(word in text for word in ["angry", "anger", "mad", "frustrated", "annoyed", "irritated"]):
        return "😠 Angry / Frustrated"
    elif any(word in text for word in ["happy", "excited", "proud", "grateful", "confident"]):
        return "😊 Positive"
    else:
        return "🙂 General Support"

def respond(message, history):

    if not any(word in message.lower() for word in mental_health_keywords):
        return "I'm mainly here to support teen mental wellness. Try asking me about stress, anxiety, school pressure, friendships, emotions, sleep, or coping strategies."

    mood = detect_mood(message)
    
    relevant_chunk = get_top_chunk(message)
    print("Retrieved chunk:", relevant_chunk) # we are using this line to test RAG system (the terminal will show what knowledge chunk chatbot used)
    
    messages = [
    {
        "role": "system",
        "content": """
You are a friendly, supportive mental wellness chatbot for teens.

Respond in a warm and encouraging way
Use the provided information to answer the user's question.
Use simple language
Give clear coping strategies.
Keep the response helpful but not too long.
You are not a replacement for a therapist or doctor.
If the information does not contain the answer, say so instead of making up information.
If the user mentions self-harm, suicide, abuse, or danger, encourage them to call or text 988, a trusted adult or emergency support right away.
"""
    }
]
    messages.append({
        "role": "user",
        "content": f"""
Use this information to help answer the user's question:

{relevant_chunk}

User question:
{message}
    """
    })

    response = client.chat_completion(
        messages=messages,
        max_tokens=400 # allow longer responses
    )

    return f"🧠 MindMate Mood Check: {mood}\n\n" + response.choices[0].message.content.strip()

custom_css = """
body {
    background: #e8f5f1;
}

.gradio-container {
    background: #e8f5f1;
}

textarea {
    background-color: #ffffff !important;
    border: 2px solid #7fc8a9 !important;
}

button {
    background-color: #7fc8a9 !important;
    color: white !important;
}
"""

with gr.Blocks(css=custom_css) as app:
    gr.ChatInterface(
        respond,
        title="🧠 MindMate: Teen Mental Wellness Assistant",
        description="""
        Welcome to MindMate! 🌱

        Feeling stressed, anxious, or overwhelmed? Ask me questions about mental wellness, coping strategies, and emotional well-being.

        MindMate is for support and education and is not a replacement for professional mental health care.
        """,
        examples=[
            "I've been feeling stressed about school lately.",
            "What are some coping strategies for anxiety?",
            "How can I manage test anxiety?",
            "What should I do when I'm feeling overwhelmed?",
            "My parents put a lot of pressure on me."
        ]
    )

app.launch()