import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch
from langdetect import detect

# --- CONFIG ---
MODEL_NAME = "distilgpt2"  # Modelo ligero para CPU
TRANSLATE_TO_ES_MODEL = "Helsinki-NLP/opus-mt-mul-es"
TRANSLATE_FROM_ES_MODEL = "Helsinki-NLP/opus-mt-es-mul"

# --- Cargar modelos ---
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)

translator_to_es = pipeline("translation", model=TRANSLATE_TO_ES_MODEL, device=-1)
translator_from_es = pipeline("translation", model=TRANSLATE_FROM_ES_MODEL, device=-1)

# --- Funciones de traducción ---
def translate_to_es(text):
    try:
        lang = detect(text)
    except:
        lang = "es"
    if lang != "es":
        translated = translator_to_es(text)[0]["translation_text"]
        return translated, lang
    return text, lang

def translate_from_es(text, lang):
    if lang != "es":
        translated = translator_from_es(text)[0]["translation_text"]
        return translated
    return text

# --- Función principal del chatbot ---
def answer(history, message):
    if not message.strip():
        return history, ""

    # Detectar idioma y traducir a español si es necesario
    msg_es, lang = translate_to_es(message)

    # Construir contexto
    context = ""
    for user, bot in history[-6:]:
        context += f"Usuario: {user}\nIA: {bot}\n"
    context += f"Usuario: {msg_es}\nIA:"

    # Generar respuesta en español
    output = generator(
        context,
        max_new_tokens=150,
        do_sample=True,
        top_k=50,
        top_p=0.9,
        temperature=0.8
    )[0]["generated_text"]

    if "IA:" in output:
        response_es = output.split("IA:")[-1].strip()
    else:
        response_es = output

    # Traducir de vuelta al idioma original
    response_final = translate_from_es(response_es, lang)

    # Actualizar historial
    history.append((message, response_final))
    return history, ""

# --- Interfaz Gradio ---
with gr.Blocks() as demo:
    gr.Markdown("# 🤖 Chatbot Multilenguaje (Traducción automática)")
    chat = gr.Chatbot()
    msg = gr.Textbox(placeholder="Escribe tu mensaje…")
    clear_btn = gr.Button("Limpiar")
    state = gr.State([])

    msg.submit(answer, [state, msg], [chat, msg])
    clear_btn.click(lambda: [], None, chat)

demo.launch()