import gradio as gr
from transformers import pipeline


sentiment_pipeline = pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english")
flagged_words = ["hate", "stupid", "ugly", "bad"]
positive_flags = ["love", "amazing", "great"]
negative_flags = ["awful", "terrible", "worst"]

def analyze_message(message, custom_flagged_words, custom_positive_flags, custom_negative_flags):
    flagged_list = custom_flagged_words.split(",") if custom_flagged_words else []
    positive_list = custom_positive_flags.split(",") if custom_positive_flags else []
    negative_list = custom_negative_flags.split(",") if custom_negative_flags else []

    sentiment_result = sentiment_pipeline(message)
    sentiment_label = sentiment_result[0]['label']
    flagged_terms = [word for word in flagged_list if word in message.lower()]
    positive_terms = [word for word in positive_list if word in message.lower()]
    negative_terms = [word for word in negative_list if word in message.lower()]

    flagged_status = "Yes" if flagged_terms else "No"

    # Rule-based flagging
    if sentiment_label == "POSITIVE" and positive_terms:
        sentiment_label = "FLAGGED POSITIVE"
    elif sentiment_label == "NEGATIVE" and negative_terms:
        sentiment_label = "FLAGGED NEGATIVE"


    return sentiment_label, ", ".join(flagged_terms) if flagged_terms else "None"

with gr.Blocks() as demo:
    gr.Markdown("## Messaging Sentiment analysis")

    with gr.Row():
        input_text = gr.Textbox(label="Enter Message")
        submit_btn = gr.Button("submit")

    sentiment_output = gr.Textbox(label="Sentiment Category", interactive=False)
    flagged_words_output = gr.Textbox(label="Flagged Words", interactive=False)

    gr.Markdown("### Customize Flagging Rules")
    custom_flagged_words = gr.Textbox(label="Custom Flagged Words", value=", ".join(flagged_words))
    custom_positive_flags = gr.Textbox(label="Words to Flag Positive Messages", value=", ".join(positive_flags))
    custom_negative_flags = gr.Textbox(label="Words to Flag Negative Messages", value=", ".join(negative_flags))

    submit_btn.click(analyze_message,
                     inputs=[input_text, custom_flagged_words, custom_positive_flags, custom_negative_flags],
                     outputs=[sentiment_output, flagged_words_output])

demo.launch()