import os
import openai
import gradio as gr
from youtube_transcript_api import YouTubeTranscriptApi
from langchain_openai import ChatOpenAI
from langchain.agents import AgentExecutor
from langchain.memory import ConversationBufferWindowMemory
from youtube_FC_14 import YouTubeTranscriptTool, MainPointsExtractor, SummaryExtractor, YouTubeAgent

import logging
logging.getLogger().setLevel(logging.ERROR)

import warnings
warnings.filterwarnings("ignore")

class ChatBot:
    def __init__(self):
        self.youtube_agent = None
        self.api_key = None

    def initialize_agent(self, api_key):
        if api_key:
            os.environ['OPENAI_API_KEY'] = api_key
            openai.api_key = api_key
            self.api_key = api_key
            self.youtube_agent = YouTubeAgent()
            return "API key set successfully. Agent initialized."
        else:
            return "Please provide a valid API key."

    def chat(self, message, history):
        if not self.youtube_agent:
            return "Please set your OpenAI API key first."
        
        try:
            response = self.youtube_agent.invoke(message)
            return response
        except Exception as e:
            return f"An error occurred: {str(e)}"

chatbot = ChatBot()  # Create an instance of ChatBot

def set_api_key(api_key):
    return chatbot.initialize_agent(api_key)

def user_message(message, history):
    return "", history + [[message, None]]

def bot_message(history):
    user_message = history[-1][0]
    bot_response = chatbot.chat(user_message, history)
    history[-1][1] = bot_response
    return history

def use_example(example, text_input):
    return example

# Example messages
example_messages = [
    "What tools are available for use?",
    "What is the following video about? https://www.youtube.com/watch?v=dZxbVGhpEkI",
    "Can you summarize this video? https://www.youtube.com/watch?v=hM8unyUM6KA",
    "Extract the main points from this video: https://www.youtube.com/watch?v=UF8uR6Z6KLc"
]

with gr.Blocks() as demo:
    gr.Markdown("# YouTube Video Analysis Chatbot")
    
    with gr.Row():
        api_key_input = gr.Textbox(type="password", label="Enter your OpenAI API key")
        api_key_button = gr.Button("Set API Key")
    
    api_key_status = gr.Textbox(label="API Key Status", interactive=False)
    
    chatbot_interface = gr.Chatbot()
    msg = gr.Textbox(label="Message")
    
    with gr.Row():
        submit_btn = gr.Button("Submit")
        clear_btn = gr.Button("Clear")
    
    gr.Markdown("## Example Messages")
    example_btns = [gr.Button(i) for i in example_messages]

    api_key_button.click(set_api_key, inputs=api_key_input, outputs=api_key_status)
    
    submit_btn.click(user_message, [msg, chatbot_interface], [msg, chatbot_interface], queue=False).then(
        bot_message, chatbot_interface, chatbot_interface
    )
    
    msg.submit(user_message, [msg, chatbot_interface], [msg, chatbot_interface], queue=False).then(
        bot_message, chatbot_interface, chatbot_interface
    )
    
    clear_btn.click(lambda: None, None, chatbot_interface, queue=False)
    
    for btn, example in zip(example_btns, example_messages):
        btn.click(use_example, inputs=[gr.Textbox(value=example, visible=False)], outputs=msg)

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