import gradio as gr
from transformers import pipeline
from wordcloud import WordCloud
import matplotlib
matplotlib.use("Agg")  # fix for HF Spaces
import matplotlib.pyplot as plt
import io
import pandas as pd

# Load Hugging Face pipelines with smaller stable models
sentiment_model = pipeline(
    "sentiment-analysis",
    model="distilbert-base-uncased-finetuned-sst-2-english"
)
summarizer = pipeline(
    "summarization",
    model="sshleifer/distilbart-cnn-12-6"
)

# Word cloud generator
def generate_wordcloud(text):
    wordcloud = WordCloud(width=800, height=400, background_color="white").generate(text)
    img = io.BytesIO()
    plt.figure(figsize=(8, 4))
    plt.imshow(wordcloud, interpolation="bilinear")
    plt.axis("off")
    plt.savefig(img, format="png")
    plt.close()
    img.seek(0)
    return img

# Process single comment
def process_comment(comment):
    # Sentiment
    sentiment = sentiment_model(comment)[0]

    # Summary
    try:
        summary = summarizer(comment, max_length=60, min_length=10, do_sample=False)[0]['summary_text']
    except Exception:
        summary = "Summary not available for short text."

    # Word cloud
    wc_img = generate_wordcloud(comment)

    return f"{sentiment['label']} ({sentiment['score']:.2f})", summary, wc_img

# Batch processing from CSV
def analyze_batch(file_path):
    df = pd.read_csv(file_path)
    results, summaries = [], []
    combined_text = ""

    for comment in df["comment"].astype(str):
        sentiment, summary, _ = process_comment(comment)
        results.append(sentiment)
        summaries.append(summary)
        combined_text += " " + comment

    # One word cloud for all comments
    wc_img = generate_wordcloud(combined_text)

    df["Sentiment"] = results
    df["Summary"] = summaries
    return df, wc_img

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 📝 E-Consultation Sentiment Analysis")

    with gr.Tab("Single Comment"):
        inp = gr.Textbox(label="Enter Comment")
        out1 = gr.Textbox(label="Sentiment")
        out2 = gr.Textbox(label="Summary")
        out3 = gr.Image(label="Word Cloud")
        inp.submit(process_comment, inp, [out1, out2, out3])

    with gr.Tab("Batch Upload"):
        file_in = gr.File(label="Upload CSV with 'comment' column", type="filepath")
        out_df = gr.Dataframe(label="Results")
        out_wc = gr.Image(label="Word Cloud (All Comments)")
        file_in.change(analyze_batch, file_in, [out_df, out_wc])

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