import gradio as gr
from transformers import pipeline
from gtts import gTTS
import time
import os

# Load models
sentiment_pipeline = pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english")
summarization_pipeline = pipeline("summarization", model="RussianNLP/FRED-T5-Summarizer")

# Task logic
def perform_task(task, text):
    if not text.strip():
        return "⚠️ Please enter some text to analyze.", None, gr.update(visible=False)

    time.sleep(0.5)

    if task == "Sentiment Analysis":
        result = sentiment_pipeline(text)[0]
        label = result['label']
        score = round(result['score'], 3)
        emoji = "😊" if label == "POSITIVE" else "😟"
        confidence_bar = "█" * int(score * 10) + "░" * (10 - int(score * 10))
        output = f"""
{emoji} **Sentiment Analysis Results**

**Label:** {label}
**Confidence:** {score} ({score*100:.1f}%)
**Visual:** {confidence_bar}

**Interpretation:** This text expresses a {label.lower()} sentiment with {score*100:.1f}% confidence.
        """.strip()
        return output, None, gr.update(visible=False)

    elif task == "Summarization":
        result = summarization_pipeline(text, max_length=100, min_length=30, do_sample=False)
        summary = result[0]['summary_text']
        original_words = len(text.split())
        summary_words = len(summary.split())
        compression_ratio = round((1 - summary_words/original_words) * 100, 1)
        output = f"""
📝 **Text Summarization Results**

**Summary:**
{summary}

**Statistics:**
• Original: {original_words} words
• Summary: {summary_words} words
• Compression: {compression_ratio}% reduction
        """.strip()
        return output, None, gr.update(visible=False)

    elif task == "Text-to-Speech":
        tts = gTTS(text)
        filename = "tts_output.mp3"
        tts.save(filename)
        word_count = len(text.split())
        char_count = len(text)
        output = f"""
🔊 **Text-to-Speech Generated Successfully!**

**Input Statistics:**
• Words: {word_count}
• Characters: {char_count}
• Estimated duration: ~{word_count * 0.5:.1f} seconds

**Audio file ready for playback below** ⬇️
        """.strip()
        return output, filename, gr.update(visible=True, value=filename)

# Handle button click
def handle_task_processing(task, text):
    if not text.strip():
        return "⚠️ Please enter some text to get started!", None, gr.update(visible=False)
    result_text, audio_file, audio_update = perform_task(task, text)
    return result_text, audio_file, audio_update

# Custom CSS
custom_css = """<style>
/* Same CSS as before, trimmed for brevity */
body, .gradio-container {
    background-color: #0a0a0a !important;
    color: #ffffff !important;
}
</style>"""

# UI Layout
with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
    gr.HTML(custom_css)
    gr.Markdown("# 🤖 Multi-Task AI Assistant", elem_id="title")

    with gr.Tabs():
        # === Main Assistant Tab ===
        with gr.Tab("🧠 Multi-Task Interface"):
            gr.Markdown("""
            <div style="text-align: center; margin-bottom: 2rem; color: #b0b0b0; font-size: 1.2rem;">
                Harness the power of AI for <strong>sentiment analysis</strong>, <strong>text summarization</strong>, and <strong>text-to-speech</strong> conversion
            </div>
            """)

            with gr.Row():
                with gr.Column(scale=1, min_width=250):
                    task_selector = gr.Dropdown(
                        choices=["Sentiment Analysis", "Summarization", "Text-to-Speech"],
                        label="🔧 Select AI Task",
                        value="Sentiment Analysis",
                        info="Choose the AI capability you want to use"
                    )

                    gr.Markdown("""
                    **📊 Sentiment Analysis**: Analyze emotional tone and polarity of text  
                    **✂️ Summarization**: Generate concise summaries of long text  
                    **🔊 Text-to-Speech**: Convert text into natural-sounding audio
                    """, elem_classes=["task-info"])

                with gr.Column(scale=2):
                    textbox = gr.Textbox(
                        lines=8,
                        label="📝 Input Text",
                        placeholder="Enter your text here...",
                        info="Type or paste the text you want to process"
                    )

            with gr.Row():
                run_button = gr.Button("🚀 Process with AI", size="lg", variant="primary")

            gr.Markdown("## 📋 Results", elem_classes=["results-header"])

            with gr.Row():
                with gr.Column(scale=2):
                    output_text = gr.Textbox(
                        label="📊 Analysis Results",
                        lines=8,
                        interactive=False
                    )

                with gr.Column(scale=1):
                    output_audio = gr.Audio(
                        label="🔊 Generated Audio",
                        type="filepath",
                        visible=False,
                    )

            run_button.click(
                fn=handle_task_processing,
                inputs=[task_selector, textbox],
                outputs=[output_text, output_audio, output_audio]
            )

        # === Sentiment Chatbot Tab ===
        with gr.Tab("💬 Sentiment Chatbot"):
            gr.ChatInterface(
                fn=lambda message, history: (
                    history + [
                        (
                            message,
                            f"🧠 Sentiment: {sentiment_pipeline(message)[0]['label']} (Confidence: {round(sentiment_pipeline(message)[0]['score'], 2)})"
                        )
                    ],
                    history + [
                        (
                            message,
                            f"🧠 Sentiment: {sentiment_pipeline(message)[0]['label']} (Confidence: {round(sentiment_pipeline(message)[0]['score'], 2)})"
                        )
                    ]
                ),
                title="Sentiment Chatbot",
                description="Chat with the AI to detect the sentiment of your messages.",
            )

    gr.Markdown("""
    <div style="text-align: center; margin-top: 3rem; padding: 2rem; color: #666; border-top: 1px solid #333;">
        <p>✨ <strong>Multi-Task AI Assistant</strong> - Powered by Transformers & Gradio</p>
        <p>Built with ❤️ for seamless AI interaction</p>
    </div>
    """)

# Launch
if __name__ == "__main__":
    demo.launch(
        share=True,
        debug=True,
        show_error=True
    )