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 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."

    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 response.choices[0].message.content.strip()

chatbot = gr.ChatInterface(
    respond,
    title="🌱 Teen Mental Wellness Chatbot",
    description="Feeling stressed, anxious, or overwhelmed? Ask me questions about mental wellness, coping strategies, and emotional well-being.",
    examples=[ # give the user example prompts
        "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?"    
    ]
                          
)


chat_theme = gr.themes.Soft(
    primary_hue="ice blue",
    secondary_hue = "sage",
    neutral_hue = "gray",
    spacing_size = "lg",
    radius_size = "lg"   
).set (
    #Input area
    input_background_fill = "*neutral_50",
    input_border_color_focus = "*primary_300",
    #Button styling 
    button_primary_background_fill = "*primary_500",
    button_primary_background_fill_hover = "*primary_400"
)

chatbot = gr.ChatInterface(respond, type="messages", theme=chat_theme)

chatbot.launch(ssr_mode=False)