import gradio as gr
from transformers import pipeline
import pandas as pd
import spaces

# Load dataset
from datasets import load_dataset
ds = load_dataset('ZennyKenny/demo_customer_nps')
df = pd.DataFrame(ds['train'])

# Initialize model pipeline
from huggingface_hub import login
import os

# Login using the API key stored as an environment variable
hf_api_key = os.getenv("API_KEY")
login(token=hf_api_key)

classifier = pipeline("text-classification", model="distilbert/distilbert-base-uncased-finetuned-sst-2-english")
generator = pipeline("text2text-generation", model="google/flan-t5-base")

# Function to classify customer comments
@spaces.GPU
def classify_comments(categories):
    global df  # Ensure we're modifying the global DataFrame
    sentiments = []
    assigned_categories = []
    for comment in df['customer_comment']:
        # Classify sentiment
        sentiment = classifier(comment)[0]['label']
        # Generate category
        category_str = ', '.join(categories)
        prompt = f"What category best describes this comment? '{comment}' Please answer using only the name of the category: {category_str}."
        category = generator(prompt, max_length=30)[0]['generated_text']
        assigned_categories.append(category)
        sentiments.append(sentiment)
    df['comment_sentiment'] = sentiments
    df['comment_category'] = assigned_categories
    return df.to_html(index=False)  # Return all fields with appended sentiment and category

# Function to add a category
def add_category(categories, new_category):
    if new_category.strip() != "" and len(categories) < 5:  # Limit to 5 categories
        categories.append(new_category.strip())
    return categories, "", f"**Categories:**\n" + "\n".join([f"- {cat}" for cat in categories])

# Function to reset categories
def reset_categories():
    return [], "**Categories:**\n- None"

# Function to load data from uploaded CSV
def load_data(file):
    global df  # Ensure we're modifying the global DataFrame
    if file is not None:
        file.seek(0)  # Reset file pointer
        if file.name.endswith('.csv'):
            custom_df = pd.read_csv(file, encoding='utf-8')
        else:
            return "Error: Uploaded file is not a CSV."
        # Check for required columns
        required_columns = ['customer_comment']
        if not all(col in custom_df.columns for col in required_columns):
            return f"Error: Uploaded CSV must contain the following column: {', '.join(required_columns)}"
        df = custom_df
        return "Custom CSV loaded successfully!"
    else:
        return "No file uploaded."

# Function to use template categories
def use_template():
    template_categories = ["Product Experience", "Customer Support", "Price of Service", "Other"]
    return template_categories, f"**Categories:**\n" + "\n".join([f"- {cat}" for cat in template_categories])

# Custom CSS for button styling
custom_css = """
/* Style buttons only */
button:not(.file-upload-button) {
    background-color: #4c6bfd !important;
    border: none !important;
    color: white !important;
    padding: 10px 20px !important;
    text-align: center !important;
    text-decoration: none !important;
    display: inline-block !important;
    font-size: 16px !important;
    border-radius: 5px !important;
    transition: background-color 0.3s ease !important;
}

button:not(.file-upload-button):hover {
    background-color: #3a52cc !important;
}

/* Hide footer */
footer {
    visibility: hidden;
}
"""

# Gradio Interface
with gr.Blocks(css=custom_css) as nps:
    # State to store categories
    categories = gr.State([])

    # App title
    gr.Markdown("# 🎉 Customer Comment Classifier 🎉")

    # Short explanation
    gr.Markdown("""
    This app classifies customer comments into categories and assigns sentiment labels (Positive/Negative). 
    You can upload your own dataset or use the provided template. The app will append the generated 
    `comment_sentiment` and `comment_category` fields to your dataset.
    """)

    # File upload and instructions
    with gr.Row():
        with gr.Column(scale=1):
            uploaded_file = gr.File(label="📂 Upload CSV", type="filepath", scale=1)
        with gr.Column(scale=1):
            gr.Markdown("""
            **📝 Instructions:**
            - Upload a CSV file with at least one column: `customer_comment`.
            - If you don't have your own data, click **Use Template** to load a sample dataset.
            """)
            template_btn = gr.Button("✨ Use Template", size="sm")

    gr.Markdown("---")

    # Category section
    with gr.Row():
        with gr.Column(scale=1):
            # Category input and buttons
            category_input = gr.Textbox(label="📝 New Category", placeholder="Enter category name", scale=1)
            with gr.Row():
                add_category_btn = gr.Button("➕ Add Category", size="sm")
                reset_btn = gr.Button("🔄 Reset Categories", size="sm")
            # Category display
            category_status = gr.Markdown("**📂 Categories:**\n- None")
        with gr.Column(scale=1):
            gr.Markdown("""
            **📝 Instructions:**
            - Enter a category name and click **Add Category** to add it to the list.
            - Click **Reset Categories** to clear the list.
            - The `customer_comment` field will be categorized based on the categories you provide.
            """)

    gr.Markdown("---")

    # Classify button and output
    with gr.Row():
        with gr.Column(scale=1):
            classify_btn = gr.Button("🔍 Classify", size="sm")
        with gr.Column(scale=3):  # Center the container and make it 75% of the window width
            output = gr.HTML()

    # Event handlers
    add_category_btn.click(
        fn=add_category,
        inputs=[categories, category_input],
        outputs=[categories, category_input, category_status]
    )
    reset_btn.click(
        fn=reset_categories,
        outputs=[categories, category_status]
    )
    uploaded_file.change(
        fn=load_data,
        inputs=uploaded_file,
        outputs=output
    )
    template_btn.click(
        fn=use_template,
        outputs=[categories, category_status]
    )
    classify_btn.click(
        fn=classify_comments,
        inputs=categories,
        outputs=output
    )

nps.launch(share=True)