import gradio as gr
import pandas as pd
import numpy as np
from transformers import BlenderbotTokenizer, BlenderbotForConditionalGeneration, M2M100ForConditionalGeneration, M2M100Tokenizer
from datasets import load_dataset
from deep_translator import GoogleTranslator

# Load Chatbot (BlenderBot 3B)
model_name = "facebook/blenderbot-3B"
tokenizer = BlenderbotTokenizer.from_pretrained(model_name)
chatbot_model = BlenderbotForConditionalGeneration.from_pretrained(model_name)

# Load Translation Model (Multilingual)
translate_model_name = "facebook/m2m100_418M"
translate_tokenizer = M2M100Tokenizer.from_pretrained(translate_model_name)
translate_model = M2M100ForConditionalGeneration.from_pretrained(translate_model_name)

# Load Hugging Face Dataset (Amazon Reviews)
dataset = load_dataset("amazon_us_reviews", split="train")
df = pd.DataFrame(dataset)

# Keep necessary columns
df = df[["product_category", "product_title", "star_rating", "review_body"]].dropna()
df["star_rating"] = df["star_rating"].astype(float)

# Function to translate text
def translate_text(text, target_lang="en"):
    if target_lang == "en":
        return text  # No translation needed for English
    inputs = translate_tokenizer(text, return_tensors="pt", src_lang="en")
    translated_tokens = translate_model.generate(**inputs, forced_bos_token_id=translate_tokenizer.get_lang_id(target_lang))
    return translate_tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]

# Function to recommend products based on user input with filters
def recommend_products(user_query, min_rating=3.5):
    keywords = user_query.lower().split()
    
    # Filter based on keywords & minimum rating
    recommended = df[(df["product_category"].str.lower().isin(keywords)) & (df["star_rating"] >= min_rating)]
    
    if recommended.empty:
        return "No recommendations found. Try searching for 'Electronics', 'Books', or 'Beauty products'."
    
    # Sort by highest rating
    recommended = recommended.sort_values(by="star_rating", ascending=False).head(5)
    
    return recommended[["product_title", "star_rating"]].to_string(index=False)

# Chatbot Response Function with improved answers
def chatbot_response(user_input, language="en", min_rating=3.5):
    # Translate input if not in English
    if language != "en":
        user_input = translate_text(user_input, target_lang="en")
    
    # Generate chatbot response
    inputs = tokenizer([user_input], return_tensors="pt")
    reply_ids = chatbot_model.generate(**inputs, max_length=100)
    response = tokenizer.decode(reply_ids[0], skip_special_tokens=True)

    # Get product recommendations
    recommendations = recommend_products(user_input, min_rating)

    # Translate output if needed
    if language != "en":
        response = translate_text(response, target_lang=language)
        recommendations = translate_text(recommendations, target_lang=language)
    
    return f"🤖 AI: {response}\n\n🔍 Recommended Products:\n{recommendations}"

# Gradio UI with Filters & Multi-Language
iface = gr.Interface(
    fn=chatbot_response,
    inputs=[
        gr.Textbox(label="Ask me about products!"),
        gr.Dropdown(["en", "es", "fr", "de", "hi"], label="Language", value="en"),  # Supports English, Spanish, French, German, Hindi
        gr.Slider(1, 5, value=3.5, step=0.5, label="Minimum Star Rating")
    ],
    outputs="text",
    title="🛍️ AI Shopping Assistant",
    description="Chat with an AI to get product recommendations with filters & multilingual support!",
    theme="default"
)

# Launch the App
iface.launch()