import re
from langdetect import detect, DetectorFactory
from transformers import pipeline
import gradio as gr

# Garante consistência nos resultados do langdetect
DetectorFactory.seed = 0

# Carrega modelo de texto em português
generator = pipeline(
    "text2text-generation",
    model="unicamp-dl/ptt5-base-portuguese-vocab",
    tokenizer="unicamp-dl/ptt5-base-portuguese-vocab"
)

# Mapeamento de comandos (em português)
COMMANDS = {
    "resumo": ["resuma", "resumo", "resumir"],
    "reescrever": ["reescreva", "reformule", "reformular"],
    "expandir": ["expanda", "expansão", "expandir", "detalhe"],
    "corrigir": ["corrija", "corrigir", "melhore", "revise"]
}

PROMPTS = {
    "resumo": "Resuma o texto a seguir:\n\n",
    "reescrever": "Reescreva com mais clareza e estilo:\n\n",
    "expandir": "Expanda este texto, adicionando detalhes e explicações:\n\n",
    "corrigir": "Corrija a gramática, ortografia e estilo do seguinte texto:\n\n",
    "gerar": ""
}

def detect_language(text):
    try:
        lang = detect(text)
        return lang if lang in ["pt", "en", "fr"] else "pt"
    except:
        return "pt"

def find_command(text, lang):
    text_lower = text.lower()
    for cmd, keywords in COMMANDS.items():
        for kw in keywords:
            if kw in text_lower:
                return cmd
    return "gerar"

def clean_text(text):
    text = re.sub(r"\s+", " ", text).strip()
    for keywords in COMMANDS.values():
        for kw in keywords:
            text = re.sub(rf"\b{kw}\b", "", text, flags=re.IGNORECASE)
    return text.strip(": ").strip()

def build_prompt(core, cmd):
    return PROMPTS.get(cmd, "") + core

def gerar_resposta(texto):
    lang = detect_language(texto)
    cmd = find_command(texto, lang)
    core = clean_text(texto)
    prompt = build_prompt(core, cmd)

    output = generator(prompt, max_new_tokens=256, temperature=0.7, top_p=0.9)[0]["generated_text"]
    resposta = output.replace(prompt, "").strip()

    if resposta.lower().startswith(core.lower()):
        resposta = resposta[len(core):].strip()

    return resposta or "⚠️ A IA não conseguiu processar esse texto. Tente reformular."

# Interface
demo = gr.Interface(
    fn=gerar_resposta,
    inputs=gr.Textbox(lines=6, placeholder="Digite algo como: Corrija: O menino está brincano...", label="Entrada"),
    outputs=gr.Textbox(label="Resposta da IA"),
    title="🧠 IA Multilingue Sr. Nicolas",
    description="Detecta comando embutido e devolve resposta humanizada em português.",
    allow_flagging="never"
)

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