app.py

import os
import gradio as gr
import joblib
import numpy as np

from sklearn.pipeline import Pipeline
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression

from transformers import pipeline as hf_pipeline


# ============================================================
# 1. Baseline de Sentimentos (TF-IDF + Logistic Regression)
# ============================================================

BASELINE_PATH = "baseline_pipe.pkl"


def train_small_baseline(save_path=BASELINE_PATH, max_samples=8000):
    """
    Treina um baseline pequeno usando o dataset amazon_polarity.
    Só é executado se baseline_pipe.pkl não existir no Space.
    """
    from datasets import load_dataset
    import pandas as pd

    ds = load_dataset("amazon_polarity", split="train")
    ds_small = ds.shuffle(seed=42).select(range(max_samples))

    df = pd.DataFrame({"text": ds_small["content"], "label": ds_small["label"]})

    pipe = Pipeline(
        [
            ("tfidf", TfidfVectorizer(max_features=25000, ngram_range=(1, 2))),
            ("clf", LogisticRegression(max_iter=1200)),
        ]
    )

    pipe.fit(df["text"], df["label"])
    joblib.dump(pipe, save_path)
    return pipe


def load_baseline():
    if os.path.exists(BASELINE_PATH):
        return joblib.load(BASELINE_PATH)
    return train_small_baseline()


baseline_model = load_baseline()


def classify_sentiment(text: str):
    """
    Classifica o texto como positivo/negativo e retorna um dicionário.
    """
    text = text.strip()
    if not text:
        return {"erro": "Digite algo para analisar."}

    proba = baseline_model.predict_proba([text])[0]
    label = "positivo" if int(np.argmax(proba)) == 1 else "negativo"
    conf = float(np.max(proba))
    return {"sentimento": label, "confianca": round(conf, 3)}


# ============================================================
# 2. IA Generativa — FLAN-T5 (modelo aberto, sem token)
# ============================================================

GEN_MODEL = os.getenv("GEN_MODEL_ID", "google/flan-t5-base")

# text2text-generation é o modo ideal para FLAN-T5
generator = hf_pipeline(
    "text2text-generation",
    model=GEN_MODEL,
)

SYSTEM_PROMPT = (
    "Você é um atendente virtual educado, empático e profissional "
    "de uma loja online. Responda SEMPRE em português do Brasil, "
    "com 2 a 4 frases claras, naturais e diretas. "
    "Não mencione que é um modelo de IA e não explique o processo interno."
)


def build_prompt(history, user_msg: str, sentimento_json: dict) -> str:
    """
    Monta o prompt para o FLAN-T5 usando histórico + sentimento.
    Nada disso deve aparecer explicitamente na resposta.
    """
    sentimento = sentimento_json.get("sentimento")
    if not sentimento or "erro" in sentimento_json:
        sentimento = "nao identificado"

    if sentimento == "negativo":
        orientacao = (
            "O cliente está insatisfeito. Mostre empatia, peça desculpas, "
            "demonstre interesse em resolver e, se fizer sentido, peça "
            "dados adicionais (pedido, produto, canal de contato)."
        )
    elif sentimento == "positivo":
        orientacao = (
            "O cliente está satisfeito. Agradeça com entusiasmo, "
            "reforce os pontos positivos e convide o cliente a continuar "
            "comprando ou usando o serviço."
        )
    else:
        orientacao = (
            "O sentimento não está claro. Responda de forma neutra, cordial "
            "e prestativa, buscando ajudar o cliente."
        )

    # Monta o histórico em texto simples
    history_text = ""
    if history:
        for u, b in history:
            history_text += f"Cliente: {u}\nAtendente: {b}\n"

    prompt = (
        f"{SYSTEM_PROMPT}\n\n"
        "Use as instruções abaixo apenas como contexto, sem citar explicitamente:\n"
        f"- Sentimento detectado: {sentimento}\n"
        f"- Orientação de tom: {orientacao}\n\n"
        "Histórico da conversa (se houver):\n"
        f"{history_text}\n"
        "Nova mensagem do cliente:\n"
        f"Cliente: {user_msg}\n\n"
        "Escreva apenas a resposta do Atendente em português do Brasil."
    )

    return prompt


def generate_reply(history, user_msg: str):
    """
    - Classifica sentimento da nova mensagem
    - Gera resposta usando FLAN-T5
    - Atualiza o histórico
    """
    user_msg = user_msg.strip()
    if not user_msg:
        return history, "", "Digite uma mensagem antes de enviar."

    sentimento = classify_sentiment(user_msg)
    if "erro" in sentimento:
        return history, "", sentimento["erro"]

    prompt = build_prompt(history, user_msg, sentimento)

    out = generator(
        prompt,
        max_length=200,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
    )[0]["generated_text"].strip()

    # Atualiza histórico para o Chatbot
    new_history = history + [(user_msg, out)]
    # Limpa o campo de texto do usuário e também retornamos o JSON do sentimento
    return new_history, "", sentimento


# ============================================================
# 3. Interface Gradio (compatível com HF Spaces + queue)
# ============================================================

with gr.Blocks(
    title="Chatbot de Sentimentos - Prof. Rodrigo",
    theme=gr.themes.Default(),
) as demo:
    gr.Markdown(
        """
# Chatbot de Sentimentos (ML + IA Generativa)

**Professor Rodrigo — Projeto Final ML & DL**

- Classificação: TF-IDF + Regressão Logística (baseline, `amazon_polarity`)  
- Geração de resposta: `google/flan-t5-base` (modelo aberto, sem token)  
- Histórico de conversa mantido na aba de chatbot  
- Você pode enviar um `baseline_pipe.pkl` na aba **Files** do Space para usar
  um modelo de sentimentos treinado pelo seu grupo.
        """
    )

    # ----------------- Aba 1: Análise isolada -----------------
    with gr.Tab("Análise de Sentimento"):
        text_in = gr.Textbox(
            label="Digite um comentário",
            lines=4,
            placeholder="Ex.: O produto chegou quebrado, fiquei muito chateado.",
        )
        text_out = gr.JSON(label="Resultado da classificação")
        btn_analisar = gr.Button("Analisar sentimento", variant="primary")

        btn_analisar.click(classify_sentiment, inputs=text_in, outputs=text_out)

    # ----------------- Aba 2: Chatbot com histórico -----------------
    with gr.Tab("Chatbot (Análise + Resposta)"):
        chatbot = gr.Chatbot(
            label="Histórico de conversa com o atendente virtual",
            height=400,
        )
        user_box = gr.Textbox(
            label="Mensagem do cliente",
            lines=3,
            placeholder="Ex.: Estou chateado, o produto é ruim.",
        )
        sentimento_box = gr.JSON(
            label="Sentimento da última mensagem analisada",
        )
        send_btn = gr.Button("Enviar e gerar resposta", variant="primary")

        # Quando clicar, gera resposta + atualiza histórico + mostra sentimento
        send_btn.click(
            generate_reply,
            inputs=[chatbot, user_box],
            outputs=[chatbot, user_box, sentimento_box],
        )

# IMPORTANTE para Spaces: ativar fila (queue) antes de lançar
demo.queue().launch()