import gradio as gr
import spaces
from transformers import pipeline

MODEL_NAME = "distilbert-base-uncased-finetuned-sst-2-english"

classifier = pipeline("sentiment-analysis", model=MODEL_NAME, device=0)


@spaces.GPU
def analyze_single(text):
    if not text or not text.strip():
        return "Please enter some text.", None
    result = classifier(text)[0]
    label = result["label"]
    score = result["score"]
    emoji = "😊" if label == "POSITIVE" else "😞"
    label_display = f"{emoji} {label} ({score:.1%} confidence)"
    return label_display, {label: score, ("NEGATIVE" if label == "POSITIVE" else "POSITIVE"): 1 - score}


@spaces.GPU
def analyze_batch(batch_text):
    lines = [line.strip() for line in batch_text.split("\n") if line.strip()]
    if not lines:
        return [["Please enter at least one line of text.", "", ""]]
    results = classifier(lines)
    rows = []
    for line, r in zip(lines, results):
        rows.append([line, r["label"], f"{r['score']:.1%}"])
    return rows


with gr.Blocks(title="Sentiment Analysis Demo") as demo:
    gr.Markdown(
        f"# 🤗 Sentiment Analysis Demo\n"
        f"Classify text as **positive** or **negative** using "
        f"[`{MODEL_NAME}`](https://huggingface.co/{MODEL_NAME})."
    )

    with gr.Tab("Single text"):
        text_input = gr.Textbox(
            label="Enter text to analyze",
            placeholder="I really loved this movie, the acting was fantastic!",
            lines=3,
        )
        analyze_btn = gr.Button("Analyze", variant="primary")
        label_output = gr.Textbox(label="Result", interactive=False)
        confidence_output = gr.Label(label="Confidence breakdown")

        analyze_btn.click(
            fn=analyze_single,
            inputs=text_input,
            outputs=[label_output, confidence_output],
        )
        text_input.submit(
            fn=analyze_single,
            inputs=text_input,
            outputs=[label_output, confidence_output],
        )

    with gr.Tab("Batch analysis"):
        gr.Markdown("Enter multiple lines of text (one per line) to classify them all at once.")
        batch_input = gr.Textbox(
            label="Batch text",
            placeholder="This product exceeded my expectations!\nWorst customer service I've ever had.\nIt was okay, nothing special.",
            lines=6,
        )
        batch_btn = gr.Button("Analyze batch", variant="primary")
        batch_output = gr.Dataframe(
            headers=["Text", "Sentiment", "Confidence"],
            label="Results",
        )

        batch_btn.click(
            fn=analyze_batch,
            inputs=batch_input,
            outputs=batch_output,
        )

    gr.Markdown(
        "---\nBuilt with [Gradio](https://gradio.app) and "
        "[🤗 Transformers](https://huggingface.co/docs/transformers)."
    )

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