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

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 = os.getenv("MODEL_PATH", "baseline_pipe.pkl")


def train_small_baseline(save_path: str = BASELINE_PATH,
                         max_samples: int = 10000):
    """
    Treina um baseline pequeno usando uma amostra do dataset amazon_polarity.
    Usado apenas se baseline_pipe.pkl nao existir no Space.
    """
    from datasets import load_dataset
    import pandas as pd

    # Carrega o split de treino inteiro
    ds = load_dataset("amazon_polarity", split="train")

    # Embaralha e pega apenas max_samples exemplos (para ficar leve)
    ds_small = ds.shuffle(seed=42).select(range(min(max_samples, len(ds))))

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

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

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



def load_or_bootstrap_baseline():
    """
    Se existir baseline_pipe.pkl, carrega.
    Se nao existir e DISABLE_AUTOTRAIN != 1, treina um baseline pequeno.
    """
    if os.path.exists(BASELINE_PATH):
        return joblib.load(BASELINE_PATH)

    disable_auto = os.getenv("DISABLE_AUTOTRAIN", "0")
    if disable_auto == "1":
        return None

    return train_small_baseline()


baseline_model = load_or_bootstrap_baseline()


def classify_only(text: str):
    """
    Apenas classifica o sentimento (positivo/negativo) e retorna JSON.
    """
    if not text or text.strip() == "":
        return {"erro": "Digite um texto."}

    if baseline_model is None:
        return {
            "erro": (
                "Modelo baseline nao encontrado. "
                "Envie baseline_pipe.pkl na aba Files ou remova DISABLE_AUTOTRAIN."
            )
        }

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

    return {"sentimento": label, "confianca": round(conf, 3)}


# ======================================================================
# 2. IA Generativa (FLAN-T5) para resposta ao cliente
# ======================================================================

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

generator = hf_pipeline("text2text-generation", model=GEN_MODEL_ID)

SYSTEM_PROMPT = (
    "Voce e um atendente virtual educado que responde em portugues do Brasil. "
    "Use poucas frases, tom empatico e pratico."
)


def generate_reply(user_text: str, sentimento_json):
    """
    Gera uma resposta em PT-BR condicionada ao sentimento detectado.
    """
    if not user_text or user_text.strip() == "":
        return "Digite uma mensagem."

    sentimento = None
    if isinstance(sentimento_json, dict) and "sentimento" in sentimento_json:
        sentimento = sentimento_json["sentimento"]

    if sentimento == "negativo":
        intent = (
            "A avaliacao do cliente e NEGATIVA. "
            "Peca desculpas, reconheca o problema, ofereca ajuda objetiva "
            "e solicite informacoes adicionais (numero do pedido, produto, contato)."
        )
    elif sentimento == "positivo":
        intent = (
            "A avaliacao do cliente e POSITIVA. "
            "Agradeca de forma calorosa, reforce os pontos positivos citados "
            "e convide o cliente a continuar comprando."
        )
    else:
        intent = (
            "O sentimento nao foi identificado. "
            "Responda de forma neutra, cordial e util."
        )

    # opcional: pegar a confianca pra colocar no prompt
    conf = None
    if isinstance(sentimento_json, dict) and "confianca" in sentimento_json:
        conf = sentimento_json["confianca"]

    prompt = f"""
Voce e um assistente de atendimento ao cliente.

Com base na mensagem do usuario e no sentimento detectado,
gere uma resposta educada, objetiva e natural em PORTUGUES DO BRASIL,
usando entre 2 e 4 frases.

Mensagem do usuario:
\"\"\"{user_text}\"\"\"

Sentimento identificado: {sentimento}
Confianca do classificador: {conf}

Sua resposta deve:
- demonstrar empatia,
- responder diretamente ao que o cliente escreveu,
- NUNCA mencionar que esta gerando uma "resposta curta",
- NUNCA explicar o que esta fazendo,
- NUNCA repetir estas instrucoes internas.

Agora responda ao cliente da forma mais natural possivel.
"""


    out = generator(
        prompt,
        max_length=128,
        do_sample=True,
        temperature=0.6,
        top_p=0.9,
    )[0]["generated_text"]

    return out


# ======================================================================
# 3. Interface Gradio - duas abas (Analise e Chatbot)
# ======================================================================

with gr.Blocks(title="Chatbot de Sentimentos - Professor Rodrigo") as demo:
    gr.Markdown(
        """
# Chatbot de Sentimentos (ML + IA Generativa)
**Professor Rodrigo** - Projeto Final ML & DL

- Classificacao: TF-IDF + Regressao Logistica (baseline).
- Geracao: `google/flan-t5-base` para redigir respostas em PT-BR.

> Dica: envie `baseline_pipe.pkl` na aba *Files* para usar um modelo treinado por voce.
        """
    )

    # Aba 1 - somente analise de sentimento
    with gr.Tab("Analise de Sentimento"):
        input_text = gr.Textbox(
            label="Digite uma avaliacao de produto",
            lines=4,
            placeholder="Ex.: O produto chegou rapido e superou minhas expectativas.",
        )
        output_json = gr.JSON(label="Resultado da classificacao")
        btn_analisar = gr.Button("Analisar")
        btn_analisar.click(classify_only, inputs=input_text, outputs=output_json)

    # Aba 2 - Chatbot (analise + resposta generativa)
    with gr.Tab("Chatbot (Analise + Resposta)"):
        chat_input = gr.Textbox(
            label="Mensagem do cliente",
            lines=4,
            placeholder="Ex.: Estou chateado, o produto chegou quebrado.",
        )
        with gr.Row():
            btn_analise_chat = gr.Button("1) Analisar sentimento")
            btn_resposta = gr.Button("2) Gerar resposta")

        chat_analysis = gr.JSON(label="Sentimento detectado")
        chat_reply = gr.Textbox(
            label="Resposta gerada pela IA",
            lines=6,
        )

        btn_analise_chat.click(
            classify_only, inputs=chat_input, outputs=chat_analysis
        )
        btn_resposta.click(
            generate_reply, inputs=[chat_input, chat_analysis], outputs=chat_reply
        )


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