import gradio as gr
from datasets import load_dataset

# Load the Tamil Kavithai Tanglish dataset
hf_dataset = load_dataset("abishekmahi/tamil-kavithai-tanglish", split="train")
df = hf_dataset.to_pandas()

# Ensure columns are strings
df["Content"] = df["Content"].astype(str)
df["Title"] = df["Title"].astype(str)
df["Category"] = df["Category"].astype(str)

df["TanglishTitle"] = df["TanglishTitle"].astype(str)
df["TanglishContent"] = df["TanglishContent"].astype(str)
df["TanglishCategory"] = df["TanglishCategory"].astype(str)

# Define search logic
def search_kavithai(query):
    query = query.lower()
    result = df[df.apply(lambda row:
        query in row["TanglishTitle"].lower() or
        query in row["TanglishContent"].lower() or
        query in row["TanglishCategory"].lower(), axis=1)]

    if result.empty:
        return "🙏 Sorry, no matching Kavithai found."

    kavithai = result.iloc[0]
    return f"""
🏷️ **Title**: {kavithai['Title']}

📜 **Category**: {kavithai['Category']}

📝 **Kavithai**:
{kavithai['Content']}
"""

# Gradio interface
chat_interface = gr.Interface(
    fn=search_kavithai,
    inputs=gr.Textbox(lines=2, placeholder="Enter a theme, title or word...", label="🔍 Your Query"),
    outputs=gr.Textbox(label="📖 Tamil Kavithai"),
    title="தமிழ் கவிதை Chatbot ✍️",
    description="Search by theme, title or keywords to discover Tamil Kavithai in Tanglish.",
    allow_flagging="never"
)

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