import gradio as gr
import joblib
import numpy as np
import os
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM

# Caminho padrão do modelo de sentimentos
MODEL_PATH = os.getenv("MODEL_PATH", "baseline_pipe.pkl")

# Carrega o modelo de sentimentos (TF-IDF + Regressão Logística)
baseline = None
if os.path.exists(MODEL_PATH):
    try:
        baseline = joblib.load(MODEL_PATH)
        print(f"✅ Modelo de sentimento carregado: {MODEL_PATH}")
    except Exception as e:
        print(f"⚠️ Erro ao carregar modelo: {e}")
else:
    print("⚠️ baseline_pipe.pkl não encontrado. Usando modo simulado.")

# ================================================================
# 🔤 Carrega modelo generativo em português (T5-base da Unicamp)
# ================================================================
model_name = "unicamp-dl/ptt5-base-portuguese-vocab"

try:
    print(f"🔄 Carregando modelo generativo: {model_name} ...")
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
    generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
    print("✅ Modelo generativo carregado com sucesso!")
except Exception as e:
    generator = None
    print(f"⚠️ Erro ao carregar modelo T5: {e}")
    print("⚠️ IA generativa será desativada.")

# ================================================================
# 🧩 Função principal: análise + resposta empática
# ================================================================
def predict_and_respond(text: str):
    if not text or text.strip() == "":
        return {"sentimento": "n/a", "resposta": "Por favor, digite uma avaliação válida."}

    # --- Etapa 1: análise de sentimento ---
    if baseline:
        proba = baseline.predict_proba([text])[0]
        pred = int(np.argmax(proba))
        label = "positivo" if pred == 1 else "negativo"
        conf = float(np.max(proba))
    else:
        label = "positivo" if "bom" in text.lower() else "negativo"
        conf = 0.0

    # --- Etapa 2: geração de resposta empática ---
    if generator:
        prompt = (
            f"O cliente escreveu: '{text}'. "
            f"O sentimento identificado é {label}. "
            "Gere uma resposta curta, educada e empática em português, "
            "como se fosse um atendente cordial da empresa."
        )
        try:
            response = generator(prompt, max_new_tokens=80)[0]["generated_text"]
            resposta = response.strip()
        except Exception as e:
            resposta = f"(Erro ao gerar resposta: {e})"
    else:
        # Fallback se o modelo T5 não carregar
        if label == "positivo":
            resposta = "Ficamos felizes em saber que você gostou! 😊"
        else:
            resposta = "Sentimos muito pela sua experiência 😞. Nossa equipe está à disposição para ajudar!"
    
    return {
        "sentimento": label,
        "confiança": round(conf, 3),
        "resposta": resposta
    }

# ================================================================
# 🎨 Interface do Gradio
# ================================================================
demo = gr.Interface(
    fn=predict_and_respond,
    inputs=gr.Textbox(label="Digite uma avaliação de produto"),
    outputs=gr.JSON(label="Resultado"),
    title="🤖 Análise de Sentimentos + IA Generativa (Português - T5)",
    description="Analisa o sentimento e gera uma resposta empática e natural com o modelo T5 em português (Unicamp).",
    examples=[
        ["Gostei do produto, mas chegou com um pequeno defeito."],
        ["Horrível, o suporte não resolveu meu problema."],
        ["Entrega rápida e produto excelente!"]
    ],
    theme="soft"
)

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