import os
import gradio as gr
import joblib
import numpy as np
from transformers import pipeline
from sklearn.pipeline import Pipeline
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression


# ============================================================
# 1. Baseline de Sentimentos
# ============================================================

BASELINE_PATH = "baseline_pipe.pkl"

def train_small_baseline(save_path=BASELINE_PATH, max_samples=8000):
    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=20000)),
        ("clf", LogisticRegression(max_iter=1000)),
    ])

    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_only(text):
    text = (text or "").strip()
    if not text:
        return "Erro: digite um texto."

    proba = baseline_model.predict_proba([text])[0]
    pred = int(np.argmax(proba))
    lbl = "positivo" if pred == 1 else "negativo"
    conf = float(np.max(proba))

    return f"Sentimento: {lbl}\nConfiança: {conf:.3f}"


# ============================================================
# 2. IA Generativa - FLAN T5 (modelo aberto)
# ============================================================

generator = pipeline(
    "text2text-generation",
    model="google/flan-t5-base",
)

SYSTEM_PROMPT = (
    "Você é um atendente virtual educado e empático. "
    "Responda SEMPRE em português do Brasil."
)


def chatbot_logic(user_message, history):
    """
    Função usada pelo ChatInterface — history é automático.
    """

    # Detecta sentimento
    proba = baseline_model.predict_proba([user_message])[0]
    pred = int(np.argmax(proba))
    sent = "positivo" if pred == 1 else "negativo"
    conf = float(np.max(proba))

    if sent == "negativo":
        orient = "O cliente está insatisfeito. Mostre empatia, peça desculpas e ofereça ajuda."
    else:
        orient = "O cliente está satisfeito. Agradeça e incentive a continuar usando o serviço."

    prompt = (
        f"{SYSTEM_PROMPT}\n\n"
        f"Contexto de sentimento:\n- Sentimento: {sent}\n- Confiança: {conf:.3f}\n"
        f"Orientação: {orient}\n\n"
        f"Histórico prévio:\n{history}\n\n"
        f"Mensagem do cliente: {user_message}\n\n"
        "Gere uma resposta educada e profissional."
    )

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

    return out


# ============================================================
# 3. Interface Gradio — ChatInterface = sem bugs!
# ============================================================


with gr.Blocks(title="Chatbot de Sentimentos") as demo:

    gr.Markdown("""
# Chatbot de Sentimentos (ML + IA Generativa)
**Professor Rodrigo — Projeto Final ML & DL**

- Classificação: TF-IDF + Regressão Logística  
- Geração: `google/flan-t5-base`  
- Sem API aberta (evita erro "No API found")
    """)

    with gr.Tab("Análise de Sentimento"):
        txt = gr.Textbox(lines=4, label="Digite um comentário")
        out = gr.Textbox(lines=3, label="Resultado")
        btn = gr.Button("Analisar")
        btn.click(classify_only, txt, out)

    with gr.Tab("Chatbot (Análise + Resposta)"):
        gr.ChatInterface(
            fn=chatbot_logic,
            title="Atendente Virtual",
            examples=["O produto chegou quebrado", "Amei o serviço!"],
            retry_btn=None,
            undo_btn=None,
        )


demo.queue(api_open=False).launch()