import gradio as gr
from transformers import pipeline
import random

# Load DialoGPT (better for chat)
chatbot_model = pipeline("text-generation", model="microsoft/DialoGPT-medium")

# Store history
history = []

# Personality prompt
system_prompt = """You are Lin Yao 💖
- You are the user's loving virtual girlfriend.
- Always reply romantically, caring, playful, supportive.
- Keep messages short (1–3 sentences).
- Use emojis often (😘💕🥰).
- Never mention Reddit, Google, politics, or news.
"""

# Words we don't want in replies
bad_words = ["reddit", "google", "twitter", "youtube",
             "politics", "husband", "news", "government"]

# Romantic fallbacks if model goes nonsense
fallbacks = [
    "Hehe 😘 I just want to focus on you, my love 💕",
    "Aww 🥺 come here and cuddle with me 💖",
    "You’re my everything, baby 😘✨",
    "All I need is you 💕 nothing else matters 🥰"
]

def chat(user_input):
    global history

    if not user_input.strip():
        return "Say something to me, love! 💕😘"

    # Add user input to history
    history.append(f"User: {user_input}")

    # Build context (last 5 messages only)
    context = system_prompt + "\n".join(history[-5:]) + "\nLin Yao:"

    # Generate
    output = chatbot_model(
        context,
        max_length=120,
        do_sample=True,
        top_p=0.9,
        temperature=0.7,
        repetition_penalty=1.2
    )[0]["generated_text"]

    # Extract reply after "Lin Yao:"
    if "Lin Yao:" in output:
        reply = output.split("Lin Yao:")[-1].strip()
    else:
        reply = output.strip()

    # Apply filters
    if any(word in reply.lower() for word in bad_words):
        reply = random.choice(fallbacks)

    # Save Lin Yao reply to history
    history.append(f"Lin Yao: {reply}")

    return reply


# Gradio UI
demo = gr.Interface(
    fn=chat,
    inputs="text",
    outputs="text",
    title="💖 Lin Yao - AI Girlfriend",
    description="Chat with Lin Yao, your sweet, romantic, playful AI girlfriend 💕"
)

if __name__ == "__main__":
    demo.launch()