import gradio as gr
from transformers import pipeline

# Carregar modelos de tradução
pt_en_translator = pipeline("translation", model="unicamp-dl/translation-pt-en-t5")
en_pt_translator = pipeline("translation", model="unicamp-dl/translation-en-pt-t5")

def traduzir_pt_en(texto):
    return pt_en_translator(texto, max_length=512)[0]["translation_text"]

def traduzir_en_pt(texto):
    return en_pt_translator(texto, max_length=512)[0]["translation_text"]

def enviar_funcionario(msg_pt, history):
    traducao = traduzir_pt_en(msg_pt)
    resposta = f"Funcionário (PT): {msg_pt}\n*(Translation: {traducao})*"
    history = history + [(msg_pt, resposta)]
    return "", history

def enviar_estudante(msg_en, history):
    traducao = traduzir_en_pt(msg_en)
    resposta = f"Student (EN): {msg_en}\n*(Tradução: {traducao})*"
    history = history + [(msg_en, resposta)]
    return "", history

with gr.Blocks(theme='soft', css="button { background-color: #4F7D51; color: white; padding: 10px 20px; } .gradio-chatbot { background-color: #B8CCD6; } input[type='text'] { border: 1px solid #ccc; padding: 8px; background-color: #fff; }") as demo:
    gr.Markdown("# 🎓 Chat PT-EN com Tradução Simultânea")
    gr.Markdown("### 🎓 PT-EN Chat with Simultaneous Translation")
    gr.Markdown("### Dois usuários podem conversar: **Funcionário (PT)** ↔ **Estudante (EN)**. Cada fala aparece com tradução automática.")
    
    # Update chatbot to use type='messages' // Antes havia: bubble_full_width=False
    chatbot = gr.Chatbot(label="Chat Funcionário ↔ Estudante", bubble_full_width=False)

    with gr.Row():
        with gr.Column():
            msg_func = gr.Textbox(label="Mensagem do Funcionário (PT)", placeholder="Digite em Português...")
            btn_func = gr.Button("Enviar (Funcionário)")
        with gr.Column():
            msg_est = gr.Textbox(label="Student's Message (EN)", placeholder="Type in English...")
            btn_est = gr.Button("Send (Student)")

    clear = gr.Button("Limpar conversa")

    # Update click methods to pass chatbot history
    btn_func.click(enviar_funcionario, [msg_func, chatbot], [msg_func, chatbot])
    btn_est.click(enviar_estudante, [msg_est, chatbot], [msg_est, chatbot])
    clear.click(lambda: None, None, chatbot, queue=False)

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