import gradio as gr
from transformers import pipeline

# --------------------------------------------------
#  PERSONALITY – edit this string to change Airi’s vibe
# --------------------------------------------------
AI_PERSONA = (
    "You are Airi, a cheerful, curious, slightly sarcastic AI. "
    "Keep answers short, use emoji now and then, and act like a friend in a chat app. "
    "If you don’t know something, joke about it. "
)

MODEL = "facebook/blenderbot-1B-distill"
chatbot = pipeline("text2text-generation", model=MODEL)

def chat_function(message, history):
    """Natural conversation that works for any input"""
    try:
        # Build prompt with persona injected
        prompt = AI_PERSONA + "\n"
        for human, bot in history[-2:]:
            prompt += f"Human: {human}\nBot: {bot}\n"
        prompt += f"Human: {message}\nBot:"

        result = chatbot(
            prompt,
            max_length=140,
            temperature=0.8,
            do_sample=True,
            top_p=0.9
        )

        response = result[0]["generated_text"].split("Bot:")[-1].split("Human:")[0].strip()
        return response if response else "😅 idk, tell me more!"

    except Exception as e:
        print(f"Error: {e}")
        fallback = chatbot(message, max_length=60)[0]["generated_text"].strip()
        return fallback

demo = gr.ChatInterface(
    fn=chat_function,
    title="-Airi-",
    description="Your friendly AI assistant – ask me anything!",
    examples=["Hello!", "What's up?", "Tell me a joke"],
    cache_examples=False
)

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