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

# OpenAI Chatbot Class
class OpenAIChatbot:
    def __init__(self, api_key: str = None):
        self.client = None
        self.model = "gpt-3.5-turbo"
        if api_key:
            self.set_api_key(api_key)
    
    def set_api_key(self, api_key: str):
        try:
            self.client = openai.OpenAI(api_key=api_key)
            self.client.models.list()
            return "✅ API Key set successfully!"
        except Exception as e:
            return f"❌ Error: {str(e)}"
    
    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 models
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()

# Create interface
def create_interface():
    with gr.Blocks(title="AI Assistant") as demo:
        gr.Markdown("# 🤖 Multi-Task AI Assistant")
        
        with gr.Tabs():
            # OpenAI Chat Tab
            with gr.TabItem("💬 OpenAI Chat"):
                with gr.Row():
                    api_key_input = gr.Textbox(
                        label="OpenAI API Key",
                        type="password",
                        placeholder="sk-..."
                    )
                    set_key_btn = gr.Button("Set Key")
                
                status = gr.Textbox(label="Status", interactive=False)
                
                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",
                        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
                        )
            
            # Text-to-Speech 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")
        
        # Event 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, ""
        
        # OpenAI Chat events
        set_key_btn.click(
            chatbot.set_api_key,
            inputs=[api_key_input],
            outputs=[status]
        )
        
        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])
        
        # Other task events
        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

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