import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import random

model_name = "keshan/sinhala-t5-small"

# Load model with Flax weights
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name, from_flax=True)

# Fallbacks for unclear outputs
fallbacks = [
    "මට ඒකට උත්තරයක් දැනෙන්නෙ නැහැ 😅",
    "හරි, තවත් කියන්න ❤️",
    "හොඳයි, ඒ ගැන තව කියන්න 🔥",
    "ඔයාට උදව් කරන්න පුළුවන් 😇",
]

def sinhala_t5_chatbot(message, chat_history):
    # Short conversation memory (last 3 turns)
    context = ""
    for user, bot in chat_history[-3:]:
        context += f"User: {user}\nBot: {bot}\n"
    context += f"User: {message}\nBot:"

    # Encode prompt
    inputs = tokenizer.encode(context, return_tensors="pt", truncation=True, max_length=512)

    # Generate Sinhala reply
    outputs = model.generate(
        inputs,
        max_length=128,
        num_beams=4,
        temperature=0.8,
        top_p=0.9,
        repetition_penalty=1.2,
        early_stopping=True
    )

    bot_reply = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
    if not bot_reply or len(bot_reply) < 3:
        bot_reply = random.choice(fallbacks)

    chat_history.append((message, bot_reply))
    return "", chat_history


# --- Gradio UI ---
with gr.Blocks(title="සිංහල AI Chatbot (Sinhala-T5)") as demo:
    gr.Markdown("## 🧠 සිංහල AI චැට්බොට් (keshan/sinhala-t5-small - Flax Model)")
    chatbot = gr.Chatbot()
    msg = gr.Textbox(label="ඔබේ පණිවිඩය", placeholder="ඔබේ පණිවිඩය මෙතනට ලියන්න...")

    msg.submit(sinhala_t5_chatbot, [msg, chatbot], [msg, chatbot])

demo.launch()