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

# ✅ Set your OpenAI API key here (keep private!)
openai.api_key = "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"  # <<< REPLACE THIS WITH YOUR KEY

# OpenAI Chatbot Class
class OpenAIChatbot:
    def __init__(self):
        self.client = openai
        self.model = "gpt-3.5-turbo"

    def stream_chat(self, message: str, history: list, system_prompt: str = ""):
        if not self.client:
            history.append([message, "Please set your OpenAI API key first!"])
            yield history
            return

        try:
            messages = []
            if system_prompt.strip():
                messages.append({"role": "system", "content": system_prompt})

            for chat_pair in history:
                if len(chat_pair) >= 2:
                    messages.append({"role": "user", "content": chat_pair[0]})
                    messages.append({"role": "assistant", "content": chat_pair[1]})

            messages.append({"role": "user", "content": message})
            history.append([message, ""])

            stream = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                max_tokens=1000,
                temperature=0.7,
                stream=True
            )

            bot_response = ""
            for chunk in stream:
                if chunk.choices[0].delta.content is not None:
                    bot_response += chunk.choices[0].delta.content
                    history[-1] = [message, bot_response]
                    yield history
                    time.sleep(0.02)

        except Exception as e:
            history[-1] = [message, f"Error: {str(e)}"]
            yield history

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

# Task functions
def analyze_sentiment(text):
    if not text.strip():
        return "Please enter text to analyze."
    result = sentiment_pipeline(text)[0]
    return f"**Sentiment:** {result['label']}\n**Confidence:** {result['score']:.3f}"

def summarize_text(text):
    if not text.strip():
        return "Please enter text to summarize."
    result = summarization_pipeline(text, max_length=100, min_length=30, do_sample=False)
    return result[0]['summary_text']

def text_to_speech(text):
    if not text.strip():
        return "Please enter text for TTS.", None
    tts = gTTS(text)
    filename = "tts_output.mp3"
    tts.save(filename)
    return f"Audio generated for {len(text.split())} words.", filename

# Initialize chatbot
chatbot = OpenAIChatbot()

# Build the interface
def create_interface():
    with gr.Blocks(title="AI Assistant") as demo:
        gr.Markdown("# 🤖 Multi-Task AI Assistant")

        with gr.Tabs():
            # Chatbot Tab
            with gr.TabItem("💬 Chatbot"):
                with gr.Row():
                    model_dropdown = gr.Dropdown(
                        choices=["gpt-3.5-turbo", "gpt-4"],
                        value="gpt-3.5-turbo",
                        label="Model"
                    )
                    system_prompt = gr.Textbox(
                        label="System Prompt (Optional)",
                        placeholder="You are a helpful assistant..."
                    )

                chatbot_interface = gr.Chatbot(label="Chat", height=400)

                with gr.Row():
                    msg_input = gr.Textbox(
                        placeholder="Type your message...",
                        show_label=False,
                        scale=4
                    )
                    send_btn = gr.Button("Send", variant="primary")
                    clear_btn = gr.Button("Clear")

            # Sentiment Analysis Tab
            with gr.TabItem("🧠 Sentiment Analysis"):
                with gr.Row():
                    with gr.Column():
                        sentiment_input = gr.Textbox(
                            label="Text to analyze",
                            lines=5,
                            placeholder="Enter text to analyze sentiment..."
                        )
                        sentiment_btn = gr.Button("Analyze", variant="primary")

                    with gr.Column():
                        sentiment_output = gr.Textbox(
                            label="Results",
                            lines=5,
                            interactive=False
                        )

            # Summarization Tab
            with gr.TabItem("📰 Summarization"):
                with gr.Row():
                    with gr.Column():
                        summary_input = gr.Textbox(
                            label="Text to summarize",
                            lines=8,
                            placeholder="Enter long text to summarize..."
                        )
                        summary_btn = gr.Button("Summarize", variant="primary")

                    with gr.Column():
                        summary_output = gr.Textbox(
                            label="Summary",
                            lines=8,
                            interactive=False
                        )

            # TTS Tab
            with gr.TabItem("🔈 Text-to-Speech"):
                with gr.Row():
                    with gr.Column():
                        tts_input = gr.Textbox(
                            label="Text to convert",
                            lines=5,
                            placeholder="Enter text to convert to speech..."
                        )
                        tts_btn = gr.Button("Generate Speech", variant="primary")

                    with gr.Column():
                        tts_status = gr.Textbox(label="Status", interactive=False)
                        tts_audio = gr.Audio(label="Generated Audio")

        # Handlers
        def send_message(message, history, system_prompt):
            if not message.strip():
                return history, ""
            for updated_history in chatbot.stream_chat(message, history, system_prompt):
                yield updated_history, ""

        # Events
        send_btn.click(
            send_message,
            inputs=[msg_input, chatbot_interface, system_prompt],
            outputs=[chatbot_interface, msg_input]
        )

        msg_input.submit(
            send_message,
            inputs=[msg_input, chatbot_interface, system_prompt],
            outputs=[chatbot_interface, msg_input]
        )

        clear_btn.click(lambda: None, outputs=[chatbot_interface])

        sentiment_btn.click(
            analyze_sentiment,
            inputs=[sentiment_input],
            outputs=[sentiment_output]
        )

        summary_btn.click(
            summarize_text,
            inputs=[summary_input],
            outputs=[summary_output]
        )

        tts_btn.click(
            text_to_speech,
            inputs=[tts_input],
            outputs=[tts_status, tts_audio]
        )

    return demo

# Run app
if __name__ == "__main__":
    demo = create_interface()
    demo.launch(share=True)