import gradio as gr
from langdetect import detect
from transformers import pipeline

# Load a free bilingual-friendly model (small for free CPU)
chatbot = pipeline("text-generation", model="microsoft/phi-2")

# System behavior
SYSTEM_PROMPT = """
You are a bilingual (Vietnamese + English) therapeutic chatbot that helps with overthinking.
Rules:
1. Detect user’s language each message and reply in the same language.
2. Follow this flow: LISTEN → COMFORT → CLARIFY → ASK → (if asked) GIVE ONE TIP.
3. Do NOT jump straight into solutions. Start by listening.
4. Keep answers short and supportive. One step per turn.
5. If crisis/self-harm → give safe message with hotline 111 (Vietnam).
"""

def chat(user_input, history):
    try:
        lang = detect(user_input)
    except:
        lang = "en"

    # Format input
    conversation = SYSTEM_PROMPT + "\n"
    for h in history:
        conversation += f"User: {h[0]}\nAI: {h[1]}\n"
    conversation += f"User: {user_input}\nAI:"

    # Generate
    response = chatbot(conversation, max_new_tokens=200, temperature=0.7)[0]['generated_text']
    reply = response.split("AI:")[-1].strip()
    
    return reply

with gr.Blocks() as demo:
    gr.Markdown("# 🧠 Overthinking Coach AI (VN/EN) — Free on HF Spaces")
    chatbot_ui = gr.Chatbot()
    msg = gr.Textbox(placeholder="Type in Vietnamese or English...")
    clear = gr.Button("Clear")

    def respond(message, chat_history):
        bot_message = chat(message, chat_history)
        chat_history.append((message, bot_message))
        return "", chat_history

    msg.submit(respond, [msg, chatbot_ui], [msg, chatbot_ui])
    clear.click(lambda: None, None, chatbot_ui, queue=False)

demo.launch()