import gradio as gr
import joblib
from transformers import pipeline

# Load model dan pipeline
model = joblib.load("ensemble_model.pkl")
vectorizer = joblib.load("vectorizer.pkl")
qa_pipe = pipeline("question-answering", model="Rifky/IndoBERT-QA")
ner_pipe = pipeline("ner", model="cahya/bert-base-indonesian-NER", aggregation_strategy="simple")

# --- Fungsi ---
def detect_hoax(text):
    vec = vectorizer.transform([text])
    result = model.predict(vec)[0]
    if result == 1:
        return "<div style='background-color:#e74c3c; color:white; padding:10px; border-radius:5px'>HOAX</div>"
    else:
        return "<div style='background-color:#27ae60; color:white; padding:10px; border-radius:5px'>BUKAN HOAX</div>"

def qa_manual(message, history, context):
    if not context:
        return history + [[message, "Mohon isi teks berita terlebih dahulu."]]
    result = qa_pipe(question=message, context=context)
    return history + [[message, result["answer"]]]

def ner(text):
    entities = ner_pipe(text)
    styled = ""
    color_map = {
        "PER": "#ffd1dc", "ORG": "#d1e0ff", "LOC": "#d1ffd1", "MISC": "#fdfd96"
    }
    for ent in entities:
        color = color_map.get(ent["entity_group"], "#eee")
        styled += f"<mark style='background-color:{color}; padding:2px; margin:2px'>{ent['word']} <small>({ent['entity_group']})</small></mark> "
    return styled

# --- UI Gradio ---
with gr.Blocks() as demo:
    gr.Markdown("## Hoax Detector App")

    context_input = gr.Textbox(label="Teks Berita / Konteks", lines=5, placeholder="Masukkan teks berita di sini...")

    with gr.Tab("Deteksi Hoaks"):
        detect_btn = gr.Button("DETEKSI")
        hoax_output = gr.HTML()
        detect_btn.click(fn=detect_hoax, inputs=context_input, outputs=hoax_output)

    with gr.Tab("QA"):
        #gr.Markdown("### Tanya Jawab Berdasarkan Teks Berita")
        qa_question = gr.Textbox(placeholder="Tulis pertanyaan...", label="Pertanyaan")
        qa_btn = gr.Button("KIRIM")
        qa_history = gr.Chatbot(label="Riwayat Tanya Jawab")
        qa_state = gr.State([])

        qa_btn.click(
            fn=qa_manual,
            inputs=[qa_question, qa_state, context_input],
            outputs=[qa_history],
            show_progress=False
        ).then(fn=lambda h: h, inputs=qa_history, outputs=qa_state)

    with gr.Tab("NER"):
        ner_btn = gr.Button("Ekstrak Entitas")
        ner_result = gr.HTML()
        ner_btn.click(fn=ner, inputs=context_input, outputs=ner_result)

demo.launch()