import os
import json
import numpy as np
import gradio as gr

# ---------- 1) Sentiment Baseline (TF-IDF + LogisticRegression) ----------
import joblib
from sklearn.pipeline import Pipeline
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression

BASELINE_PATH = os.getenv("MODEL_PATH", "baseline_pipe.pkl")

def train_small_baseline(save_path=BASELINE_PATH, sample_pct="train[:0.5%]"):
    from datasets import load_dataset
    import pandas as pd
    ds = load_dataset("amazon_polarity", split=sample_pct)
    df = pd.DataFrame({"text": ds["content"], "label": ds["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():
    if os.path.exists(BASELINE_PATH):
        return joblib.load(BASELINE_PATH)
    if os.getenv("DISABLE_AUTOTRAIN", "0") == "1":
        return None
    return train_small_baseline()

baseline = load_or_bootstrap_baseline()

def classify_only(text):
    if not text or not text.strip():
        return {"erro": "Digite um texto."}
    if baseline is None:
        return {"erro": "Modelo baseline nao encontrado. Envie baseline_pipe.pkl ou remova DISABLE_AUTOTRAIN."}
    proba = baseline.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) Generative Assistant (FLAN-T5) ----------
from transformers import pipeline
GEN_MODEL_ID = os.getenv("GEN_MODEL_ID", "google/flan-t5-base")
generator = 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, sentimento_json):
    if not user_text or not 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."
        )
    elif sentimento == "positivo":
        intent = (
            "A avaliacao do cliente e POSITIVA. "
            "Agradeca de forma calorosa, reforce pontos fortes mencionados e convide para novas compras."
        )
    else:
        intent = "O sentimento nao foi identificado. Responda de forma neutra, util e cordial."

    prompt = (
        f"{SYSTEM_PROMPT}\n\n{intent}\n\n"
        f"Mensagem do cliente:\n\"{user_text}\"\n\n"
        f"Escreva uma resposta curta (2-4 frases)."
    )
    out = generator(prompt, max_length=128, do_sample=True, temperature=0.6, top_p=0.9)[0]["generated_text"]
    return out

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

        - Classificacao: TF-IDF + Regressao Logistica (autotreino minimo se nao houver `baseline_pipe.pkl`).  
        - Geracao: `google/flan-t5-base` para redigir respostas educadas em PT-BR.

        > Dica: Envie `baseline_pipe.pkl` em *Files* para pular o autotrain.
        \"\"\"
    )

    with gr.Tab("Analise de Sentimento"):
        in_text = gr.Textbox(label="Digite uma avaliacao de produto", lines=4, placeholder="Ex.: O produto chegou rapido e funciona muito bem.")
        out_json = gr.JSON(label="Resultado")
        btn = gr.Button("Analisar")
        btn.click(classify_only, inputs=in_text, outputs=out_json)

    with gr.Tab("Chatbot (Analise + Resposta)"):
        chat_in = gr.Textbox(label="Mensagem do cliente", lines=4)
        with gr.Row():
            analyze_btn = gr.Button("1) Analisar sentimento")
            gen_btn = gr.Button("2) Gerar resposta")
        analysis_box = gr.JSON(label="Sentimento detectado")
        reply_box = gr.Textbox(label="Resposta gerada", lines=6)
        analyze_btn.click(classify_only, inputs=chat_in, outputs=analysis_box)
        gen_btn.click(generate_reply, inputs=[chat_in, analysis_box], outputs=reply_box)

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