import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
from lingua import Language, LanguageDetectorBuilder

# -------------------------------
# Chat Model (CPU)
# -------------------------------
tokenizer = AutoTokenizer.from_pretrained("facebook/blenderbot-400M-distill")
model = AutoModelForSeq2SeqLM.from_pretrained("facebook/blenderbot-400M-distill")

def chat_in_english(user_input, history):
    history.append(user_input)
    prompt = "</s><s>".join(history)
    tokens = tokenizer(prompt, return_tensors="pt")
    reply_ids = model.generate(**tokens, max_new_tokens=200)
    return tokenizer.batch_decode(reply_ids, skip_special_tokens=True)[0]

# -------------------------------
# Language Detection
# -------------------------------
detector = LanguageDetectorBuilder.from_languages(
    Language.ENGLISH, Language.FRENCH
).build()

def detect_language(text):
    lang = detector.detect_language_of(text)
    return "fr" if lang == Language.FRENCH else "en"

# -------------------------------
# Translators
# -------------------------------
en_fr = pipeline("translation_en_to_fr", model="t5-small")
fr_en = pipeline("translation_fr_to_en", model="Helsinki-NLP/opus-mt-fr-en")

def translate_to_en(text):
    return fr_en(text)[0]["translation_text"]

def translate_to_fr(text):
    return en_fr(text)[0]["translation_text"]

# -------------------------------
# Main Logic
# -------------------------------
def chatbot(user_input, history):
    lang = detect_language(user_input)

    if lang == "fr":
        input_en = translate_to_en(user_input)
    else:
        input_en = user_input

    output_en = chat_in_english(input_en, history)

    if lang == "fr":
        final = translate_to_fr(output_en)
    else:
        final = output_en

    return final, history

# -------------------------------
# Gradio UI
# -------------------------------
with gr.Blocks() as demo:
    gr.Markdown("## 🌍 Multilingual Chatbot (CPU Version)")

    history_state = gr.State([])
    user_input = gr.Textbox(label="Your message")
    response = gr.Textbox(label="Bot reply")
    send_btn = gr.Button("Send")

    send_btn.click(
        chatbot,
        inputs=[user_input, history_state],
        outputs=[response, history_state]
    )

demo.launch()