import gradio as gr
from transformers import pipeline
import random

generator = pipeline("text2text-generation", model="King-8/affirmation-generator")

# Optional: keyword detector for typed input
category_keywords = {
    "confidence": ["confidence", "confident", "brave", "courage"],
    "calm": ["calm", "peace", "relax", "quiet", "anxious"],
    "focus": ["focus", "attention", "clarity", "concentrate"],
    "self-esteem": ["esteem", "worth", "value", "self-love"],
    "safety": ["safe", "secure", "protected"],
    "gratitude": ["grateful", "gratitude", "thankful"],
    "success": ["success", "goal", "achievement", "win"]
}

def detect_category(user_input):
    text = user_input.lower()
    for category, keywords in category_keywords.items():
        if any(keyword in text for keyword in keywords):
            return category
    return None

def generate_affirmation(category):
    prompt = f"Category: {category}"
    result = generator(prompt, max_new_tokens=20, no_repeat_ngram_size=2, repetition_penalty=1.2)
    return result[0]["generated_text"]

def chatbot_response(message, history):
    try:
        if not message.strip():
            return "Please enter a message."


        category = detect_category(message)
        if category:
            prompt = f"Category: {category}"
            result = generator(
                prompt,
                max_new_tokens=30,
                do_sample=True,
                top_k=50,
                top_p=0.95,
                temperature=0.9,
                repetition_penalty=1.2
            )
            
            return result[0]["generated_text"]
        else:
            return "Please mention a category like confidence, calm, or gratitude."

    except Exception as e:
        return f"⚠️ Something went wrong: {str(e)}"
        
gr.ChatInterface(
    fn=chatbot_response,
    title="💬 Affirmation Chatbot",
    description="Ask for an affirmation like 'Can I get something for self-esteem?' or 'Help me stay focused.'"
).launch()