import gradio as gr
import os
import joblib
import pandas as pd
from datasets import load_dataset
from sklearn.pipeline import Pipeline
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from transformers import pipeline


BASELINE_PATH = "baseline_pipe.pkl"


# ============================================================
# 1) Função para treinar automaticamente o modelo baseline
# ============================================================
def treinar_baseline():
    print("➡️ Nenhum baseline encontrado. Treinando modelo TF-IDF + Regressão Logística...")

    ds = load_dataset("amazon_polarity", split="train")
    ds_small = ds.shuffle(seed=42).select(range(8000))  # rápido para HF Spaces
    df = pd.DataFrame({"text": ds_small["content"], "label": ds_small["label"]})

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

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

    print("✔ Modelo baseline treinado e salvo!")
    return pipe


# ============================================================
# 2) Carregar modelo baseline
# ============================================================
def load_or_train():
    if os.path.exists(BASELINE_PATH):
        print("✔ baseline_pipe.pkl encontrado! Carregando...")
        return joblib.load(BASELINE_PATH)
    return treinar_baseline()


baseline_model = load_or_train()


# ============================================================
# 3) IA Generativa (sem token)
# ============================================================
generator = pipeline("text2text-generation", model="google/flan-t5-base")


def resposta_chatbot(texto):
    prompt = f"""
Você é um assistente de suporte educado, direto e profissional.
Responda de forma curta e natural.  
Pergunta: {texto}
Resposta:
"""
    out = generator(prompt, max_length=128)
    return out[0]["generated_text"]


# ============================================================
# 4) Classificação de sentimento
# ============================================================
def analisar_sentimento(texto):
    label = baseline_model.predict([texto])[0]
    prob = baseline_model.predict_proba([texto])[0].max()

    sentimento = "positivo" if label == 1 else "negativo"
    return {
        "sentimento": sentimento,
        "confianca": round(float(prob), 3)
    }


# ============================================================
# 5) Interface Gradio
# ============================================================
def interface_sentimento(texto):
    try:
        return analisar_sentimento(texto)
    except Exception as e:
        return {"erro": str(e)}


def interface_chatbot(texto, histórico):
    resposta = resposta_chatbot(texto)
    histórico.append((texto, resposta))
    return histórico, histórico


# ============================================================
# Aba de análise
# ============================================================
aba1 = gr.Interface(
    fn=interface_sentimento,
    inputs=gr.Textbox(label="Digite uma avaliação"),
    outputs=gr.JSON(label="Resultado"),
)


# ============================================================
# Aba de chatbot
# ============================================================
with gr.Blocks() as aba2:
    gr.Markdown("# Chatbot (IA Generativa)")

    chatbot = gr.Chatbot()
    entrada = gr.Textbox(label="Digite sua mensagem")

    entrada.submit(interface_chatbot, [entrada, chatbot], [chatbot, chatbot])


# ============================================================
# Tabs principais
# ============================================================
demo = gr.TabbedInterface(
    [aba1, aba2],
    ["Análise de Sentimentos", "Chatbot IA"]
)

demo.launch()