import gradio as gr

import os
import openai
import time
import json
import requests
import shutil

import matplotlib.pyplot as plt
from gradio_client import Client
from newsapi import NewsApiClient
from PIL import Image

from gpt_function_definitions import generate_image, generate_music, generate_caption, generate_caption_func, generate_music_func, generate_image_func, dict_plugin_functions

#Streaming endpoint
API_URL = "https://api.openai.com/v1/chat/completions" 
# Get the value of the openai_api_key from environment variable
openai_api_key = os.getenv("OPENAI_API_KEY") 
openai.api_key = os.getenv("OPENAI_API_KEY")
    
dicts_list = [value['dict'] for value in dict_plugin_functions.values()]

available_function_defns = {
    key.split('_func')[0]: value['func']
    for key, value in dict_plugin_functions.items()
    }
add_plugin_steps = """## Steps to add new Plugins to your Gradio ChatGPT Chatbot
Do you want to open this information in a separate tab instead? - <a href="https://huggingface.co/spaces/ysharma/ChatGPT-Plugins-in-Gradio/blob/main/README.md" target="_blank">Click here</a>.

1. **Acquire the API Endpoint**
    - You need an API which you can query, and for this example let's consider using a text-to-speech demo hosted on Huggingface Spaces.
    - **API Endpoint**: [https://gradio-neon-tts-plugin-coqui.hf.space/](https://gradio-neon-tts-plugin-coqui.hf.space/)

2. **Create a Function to Query the API**
    - You can access any Gradio demo as an API via the Gradio Python Client.
    ```python
    from gradio.client import Client

    def texttospeech(input_text):
        client = Client("https://gradio-neon-tts-plugin-coqui.hf.space/")
        result = client.predict(
            input_text,	# str  in 'Input' Textbox component
            "en",	    # str  in 'Language' Radio component
            api_name="/predict"
        )
        return result
    ```

3. **Describe the Function to GPT-3.5**
    - You need to describe your function to GPT3.5/4. This function definition will get passed to gpt and will suck up your token. GPT may or may not use this function based on user inputs later on.  
    - You can either use the Gradio demo for converting any given function to the required JSON format for GPT-3.5.
      - Demo: [Function to JSON](https://huggingface.co/spaces/ysharma/function-to-JSON)
      - Or, you can create the dictionary object on your own. Note that, the correct format is super important here.
    - MAke sure to name your JSON object description as `<function_name>_func`.
    ```python
    texttospeech_func = {
        "name": "texttospeech",
        "description": "generate speech from the given input text",
        "parameters": {
            "type": "object",
            "properties": {
                "input_text": {
                    "type": "string",
                    "description": "text that will be used to generate speech"
                }
            },
            "required": [
                "input_text"
            ]
        }
    }
    ```

4. **Add Function and JSON Object Details**
    - Add the function definition and description to the `gpt_function_definitions.py` file (simply copy and paste).
    - `dict_plugin_functions` is a dictionary of all available plugins. Add your plugin information to this dictionary in the required format.
    ```python
    'texttospeech_func': {
        'dict': texttospeech_func,
        'func': texttospeech
    }
    ```

5. **Update the Chatbot Layout**
    - Go to the Blocks Chatbot layout and add a new checkbox for your plugin as:
    ```python
    texttospeech = gr.Checkbox(label="📝🗣️Text-To-Speech", value=False)
    ```
    - Add the new checkbox component to your submit and click events for your chatbot and to the predict function accordingly.
    - And also to the `plugins` list in `predict`
      ```python
      plugins = [music_gen, stable_diff, image_cap, top_news, texttospeech]
      ```

Thats it! you are have added your own brand new CHATGPT Plugin for yourself. Go PLAY!!
"""

# managing conversation with Plugins
def run_conversation(user_input, function_call_decision):
    FLAG_MUSIC, FLAG_IMAGE, FLAG_GEN, FLAG_FUN = False, False, False, False
    # Step 1: send the conversation and available functions to GPT
    messages = [{"role": "user", "content": user_input}]
    functions = dicts_list # example values - [ generate_music_func, generate_image_func]

    # Attempt to make a request to GPT3.5/4 with retries
    max_retries = 3
    retry_delay = 5  # seconds

    for attempt in range(max_retries):
        try:
            response = openai.ChatCompletion.create(
                model="gpt-3.5-turbo-0613",
                messages=messages,
                functions=functions,
                function_call=function_call_decision,
            )
            response_message = response["choices"][0]["message"]
            print(f"response message ^^ -{response_message}")
            break  # If successful, exit the loop

        except openai.error.ServiceUnavailableError as e:
            print(f"OpenAI Server is not available. Error: {e}")
            if attempt < max_retries - 1:
                print(f"Retrying in {retry_delay} seconds...")
                time.sleep(retry_delay)
            else:
                print("Max retries reached. Exiting.")
                return None, None, None, False, False, False, False

        except openai.error.APIError as e:
            # This will catch API errors from OpenAI
            print(f"An API error occurred: {e}")
            if attempt < max_retries - 1:
                print(f"Retrying in {retry_delay} seconds...")
                time.sleep(retry_delay)
            else:
                print("Max retries reached. Exiting.")
                return None, None, None, False, False, False, False

        except Exception as e:
            # This will catch any other exceptions that are raised.
            print(f"An unexpected error occurred: {e}")
            return None, None, None, False, False, False, False

    # Step 2: check if GPT wanted to call a function
    if response_message.get("function_call"):
        FLAG_FUN = True
        # Step 3: call the function
        # Note: the JSON response may not always be valid; be sure to handle errors
        available_functions = available_function_defns
        # only one function in this example, but you can have multiple
        function_name = response_message["function_call"]["name"]
        print(f"function_name - {function_name}")

        try:
          function_to_call = available_functions[function_name]
          function_args = json.loads(response_message["function_call"]["arguments"])
          print(f"Logging: fuction_name is - {function_name}")
          print(f"Logging: fuction_to_call is - {function_to_call}")
          print(f"Logging: function_args is - {function_args}")
          function_response = function_to_call(**function_args)
          print(f"Logging: function_response ^^ is -{function_response}")
        
        except KeyError as e:
            print(f"Function not found: {e}")
            return response_message, None, None, False, False, False, False
        
        except Exception as e:
            print(f"An error occurred while calling the function: {e}")
            return response_message, None, None, False, False, False, False

        if isinstance(function_response, str):
          if function_response.split('.')[-1] == 'png':
              FLAG_IMAGE = True
          elif function_response.split('.')[-1] in ['mp4', "wav", "mp3"]:
              FLAG_MUSIC = True
          else:
              FLAG_GEN = True
        else:
            print("PLUGIN FUNCTION RETURNS A NON-STRING OUTPUT: FIX IT TO A STRING OUTPUT TO GET A RESPONSE FROM GPT")

        # Step 4: send the info on the function call and function response to GPT
        messages.append(response_message)  # extend conversation with assistant's reply
        messages.append(
            {
                "role": "function",
                "name": function_name,
                "content": function_response,
            }
        )
        print(f"Logging: messages is - {messages}")
        # extend conversation with function response
        second_response = openai.ChatCompletion.create(
            model="gpt-3.5-turbo-0613",
            messages=messages,
        )  # get a new response from GPT where it can see the function response

        print(f"Logging: second_response is - {second_response}")
        print(f"Logging: values of Music, Image, and General flags are respectively - {FLAG_MUSIC}, {FLAG_IMAGE}, {FLAG_GEN}")
        return response_message, second_response, function_response, FLAG_MUSIC, FLAG_IMAGE, FLAG_GEN, FLAG_FUN

    else:
        return response_message, None, None, False, False, False, False #second_response, function_response, FLAG_MUSIC, FLAG_IMAGE, FLAG_GEN, FALG_FUN


# driver 
def predict(inputs, top_p, temperature, chat_counter, music_gen, stable_diff, image_cap, top_news, file_output, plugin_message, chatbot=[], history=[]):  #repetition_penalty, top_k

    #openai.api_key = os.getenv("OPENAI_API_KEY")
    
    payload = {
    "model": "gpt-3.5-turbo-0613",
    "messages": [{"role": "user", "content": f"{inputs}"}],
    "temperature" : 1.0,
    "top_p":1.0,
    "n" : 1,
    "stream": True,
    "presence_penalty":0,
    "frequency_penalty":0,
    }

    headers = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {openai_api_key}"
    }

    print(f"chat_counter - {chat_counter}")
    print(f"music_gen is {music_gen}, stable_diff is {stable_diff}")

    # file handling
    print(f"Logging: file_output is - {file_output}")
    if file_output is not None:
      files_avail =  [f.name for f in file_output ]
      print(f"Logging: files_available are - {files_avail} ")
    else:
      print("Logging: No files available at the moment!")

    if chat_counter != 0 :
        messages=[]
        for data in chatbot:
          temp1 = {}
          temp1["role"] = "user"
          temp1["content"] = data[0]
          temp2 = {}
          temp2["role"] = "assistant"
          temp2["content"] = data[1]
          messages.append(temp1)
          messages.append(temp2)
        temp3 = {}
        temp3["role"] = "user"
        temp3["content"] = inputs
        messages.append(temp3)
        #messages
        payload = {
        "model": "gpt-3.5-turbo",
        "messages": messages, #[{"role": "user", "content": f"{inputs}"}],
        "temperature" : temperature, #1.0,
        "top_p": top_p, #1.0,
        "n" : 1,
        "stream": True,
        "presence_penalty":0,
        "frequency_penalty":0,
        }

    chat_counter+=1
    history.append(inputs)
    print(f"Logging: payload is - {payload}")

    plugins = [music_gen, stable_diff, image_cap, top_news, ] 
    function_call_decision = "auto" if any(plugins) else "none"
    #function_call_decision = "none" if not (music_gen or stable_diff) else "auto"
    #function_call_decision = "auto" if (music_gen or stable_diff or image_cap) else "none"
    print(f"Logging: function_call_decision flag (auto/none) is - {function_call_decision}")
    IS_FUN = False
    first_response = None

    if function_call_decision == "auto":
        first_response, second_response, function_response, IS_MUSIC, IS_IMAGE, IS_GEN, IS_FUN = run_conversation(inputs, function_call_decision)
        print(f"Logging: first_response return value - {first_response}")
        print(f"Logging: second_response return value - {second_response}")
        print(f"Logging: function_response return value - {function_response}")
        print(f"Logging: IS_MUSIC, IS_IMAGE, IS_GEN, IS_FUN, respectively return value - {IS_MUSIC}, {IS_IMAGE}, {IS_GEN}, {IS_FUN}")

        if (second_response is None) and (first_response is None):
          bot_response_using_plugins_error = 'Something went wrong! It was either your query or the OpenAI server. I would suggest you can either try again from the start or just reword your last message for more appropriate response.'

          history.append(bot_response_using_plugins_error)
          print(f"Logging: history with plugins is - {history}")
          chat = [(history[i], history[i+1]) for i in range(0, len(history)-1, 2)] + ([(history[-1],)] if len(history) % 2 != 0 else [])
          print(f"Logging: chat with plugins is - {chat}")

          yield chat, history, chat_counter, gr.update(visible=True), gr.update(visible=True), gr.update(visible=False)
          #yield {chatbot: chat, state:history, chat_counter:chat_counter, plugin_message: gr.update(visible=False) }
        
        if (second_response is not None): # and (first_response is not None):
          bot_response_using_plugins = second_response['choices'][0]['message']['content']
          print(f"Logging: bot_response_using_plugins using plugins is - {bot_response_using_plugins}")
          bot_response_using_plugins = bot_response_using_plugins.replace("sandbox:", "")
            
          history.append(bot_response_using_plugins)
          print(f"Logging: history with plugins is - {history}")
          chat = [(history[i], history[i+1]) for i in range(0, len(history)-1, 2)] + ([(history[-1],)] if len(history) % 2 != 0 else [])
          print(f"Logging: chat with plugins is - {chat}")

          if IS_MUSIC:
            yield chat, history, chat_counter, gr.update(value=function_response), gr.update(visible=True), gr.update(value="<big><b>⏳ Using MusicGen Plugin</big></b>")
            #yield {chatbot: chat, state:history, chat_counter:chat_counter, gen_music:gr.update(value=function_response), plugin_message: gr.update(value="**## ⏳ Using MusicGen Plugin**") }
          elif IS_IMAGE:
            yield chat, history, chat_counter, gr.update(visible=True), gr.update(value=function_response), gr.update(value="<big><b>⏳ Using Diffusers Plugin</big></b>")
            #yield {chatbot: chat, state:history, chat_counter:chat_counter, gen_image:gr.update(value=function_response), plugin_message: gr.update(value="**## ⏳ Using Diffusers Plugin**") }
          elif IS_GEN:
            yield chat, history, chat_counter, gr.update(visible=True), gr.update(visible=True), gr.update(value="<big><b>⏳ Using ImageCaption/News Plugin</big></b>")
            #yield {chatbot: chat, state:history, chat_counter:chat_counter, plugin_message: gr.update(value="**## ⏳ Using ImageCaption/News Plugin**") }
          
 
    # When no plugins are chosen; or when plugins are chosen but none was used
    if (function_call_decision == "none") or (first_response is not None and IS_FUN == False):
        # make a POST request to the API endpoint using the requests.post method, passing in stream=True
        response = requests.post(API_URL, headers=headers, json=payload, stream=True)
        #response = requests.post(API_URL, headers=headers, json=payload, stream=True)
        token_counter = 0
        partial_words = ""

        counter=0
        for chunk in response.iter_lines():
            #Skipping first chunk
            if counter == 0:
              counter+=1
              continue
            #counter+=1
            # check whether each line is non-empty
            if chunk.decode() :
              chunk = chunk.decode()
              # decode each line as response data is in bytes
              if len(chunk) > 12 and "content" in json.loads(chunk[6:])['choices'][0]['delta']:
                  #if len(json.loads(chunk.decode()[6:])['choices'][0]["delta"]) == 0:
                  #  break
                  partial_words = partial_words + json.loads(chunk[6:])['choices'][0]["delta"]["content"]
                  if token_counter == 0:
                    history.append(" " + partial_words)
                  else:
                    history[-1] = partial_words
                  chat = [(history[i], history[i + 1]) for i in range(0, len(history) - 1, 2) ]  # convert to tuples of list
                  token_counter+=1
                  yield chat, history, chat_counter, gr.update(visible=True), gr.update(visible=True), gr.update(visible=False)  
                  #yield {chatbot: chat, state:history, chat_counter:chat_counter, plugin_message: gr.update(visible=False) }


def reset_textbox():
    return gr.update(value='')

def add_image(file_to_save, file_output):
    print(f"Logging: image file_to_save is - {file_to_save}")
    print(f"Logging: files available in directory are -{file_output}")

    if file_output is not None:
      file_output = [f.name for f in file_output]
    if file_to_save is None:
      return file_output
    file_output = [file_to_save] if file_output is None else file_output + [file_to_save]
    print(f"Logging: Updated file directory - {file_output}")
    return file_output #gr.update(value="dog1.jpg")

def add_audio(file_to_save, file_output):
    print(f"Logging: audio file_to_save is - {file_to_save}")
    print(f"Logging: files available in directory are -{file_output}")

    if file_output is not None:
      file_output = [f.name for f in file_output]
    if file_to_save is None:
      return file_output
    file_output = [file_to_save] if file_output is None else file_output + [file_to_save]
    print(f"Logging: Updated file directory - {file_output}")
    return file_output #gr.update(value="dog1.jpg")

def upload_file(file, file_output):
    print(f"Logging: all files available - {file_output}")
    print(f"Logging: file uploaded is - {file}")

    img_orig_name = file.name.split('/')[-1]
    shutil.copy2(file.name, img_orig_name)

    file_output = [file] if file_output is None else file_output + [file]
    file_output = [f.name for f in file_output]
    print(f"Logging: Updated file list is - {file_output}")
    return file_output

messaging = """
How does a Language Model like GPT makes discerning choices regarding which plugins to run? Well, this is done using the Language Model as a reasoning agent and allowing it to assess and process information intelligently:<br><br>
<b>Function Calling</b>: Interacting with external APIs via free-form text isn't optimal; instead, employing JSON format proves to be a more efficient method.<br>
<b>Gradio Chatbots</b>: Using Gradio and Function Calling you can create chatbots designed to respond to queries by communicating with external APIs. The API responses are fed back to the Language Model for processing and a new response is generated for the user.<br>
<b>Describe your functions to GPT</b>: When integrating with GPT-3.5, specific instructions on how to utilize a particular function or plugin are essential; this encompasses specifying the name, description, and required parameters or inputs. Look at gpt_function_definitions.py for more context.<br>
<b>Caution</b>: Such function definitions would be conveyed to GPT, so when duplicating to build your own Plugins, proceed with caution as functions consume tokens.<br>
<b>Gradio's Usefulness</b>: The versatility of this using Gradio to build LLM applications is immense; In this Gradio app, you can have an array of functions tailored for various purposes, enhancing the breadth and depth of interactions with your Language Model.
"""
howto = """
Welcome to the <b>ChatGPT-Plugins</b> demo, built using Gradio! This interactive chatbot employs the GPT3.5-turbo-0613 model from OpenAI and boasts custom plugins to enhance your chatting experience. Here’s a quick guide to get you started:<br><br>
<b>Getting Started</b>: Simply type your messages in the textbox to chat with ChatGPT just like you would in the original app.<br>
<b>Using Plugins</b>: Want to try out a plugin? Check the checkbox next to the plugin you want to use.<br><br>

<b>DIFFUSERS PLUGIN:</b><br>
<b>What it does:</b> Generates images based on your text descriptions.<br>
<b>How to use:</b> Type a text description of the image you want to generate, and the plugin will create it for you.<br>
<b>Example input:</b> "Generate an image of a sunset over the mountains."<br><br>

<b>MUSIC-GEN PLUGIN:</b><br>
<b>What it does:</b> Generates music based on your descriptions.<br>
<b>How to use:</b> Describe the type of music you want and select an input melody. Remember to upload a melody first!<br>
<b>Example input:</b> "Generate music for a parade using bach.mp3 as input melody."<br><br>

<b>IMAGE CAPTION PLUGIN:</b><br>
<b>What it does:</b> Describes images that you upload.<br>
<b>How to use:</b> Upload an image and ask ChatGPT to describe it by name.<br>
<b>Example input:</b> "Describe the image dog.jpg."<br><br>

<b>NEWS PLUGIN:</b><br>
<b>What it does:</b> Provides the top 3 news articles based on your search query.<br>
<b>How to use:</b> Simply type in a search query and the plugin will present the top 3 news articles matching your query based on relevance.<br>
<b>Example input:</b> "Show me the top news about space exploration."<br><br>

Access Generated Content: Find all generated images and audio in the Gradio Files component located below the input textbox.<br>
Have Fun!: Explore and enjoy the versatile features of this <b>Gradio-ChatGPT-PLUGIN demo</b>.<br>
Now you’re all set to make the most of this ChatGPT demo. Happy chatting!
"""

with gr.Blocks(css = """#col_container { margin-left: auto; margin-right: auto;}
                #chatbot {height: 520px; overflow: auto;}""") as demo:   
    gr.HTML('<h1 align="center">Build Your Own 🧩Plugins For ChatGPT using 🚀Gradio</h1>')

    with gr.Accordion("Create Plugins for ChatGPT using Gradio in less than 5 minutes!", open=False ):
        gr.Markdown(add_plugin_steps)

    with gr.Accordion("How to use the demo and other useful stuff:", open=False):
        with gr.Accordion("How to use the demo?", open=False):
            gr.HTML(howto)                    
        with gr.Accordion("What is happening?", open=False):
            gr.HTML(messaging)                    
                    
    gr.HTML('''<center><a href="https://huggingface.co/spaces/ysharma/ChatGPT-Plugins-in-Gradio?duplicate=true"><img src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>Duplicate the Space and run securely with your OpenAI API Key</center>''')
    #with gr.Column(elem_id = "col_container"):
    with gr.Row():
      with gr.Column():
        with gr.Accordion("OpenAI API KEY🔑"):
          openai_api_key_tb = gr.Textbox(label="Enter your OpenAI API key here", value="🎁GPT3.5 keys are provided by HuggingFace for Free🥳 Don't need to enter yours!😉🙌")
          plugin_message = gr.Markdown()
      with gr.Column():
        with gr.Accordion("Plug-ins🛠️: Check the box against the plugins you want to use (can select all or few or none)",):
          music_gen = gr.Checkbox(label="🎵MusicGen", value=False)
          stable_diff = gr.Checkbox(label="🖼️Diffusers", value=False)
          image_cap = gr.Checkbox(label="🎨Describe Image", value=False)
          top_news = gr.Checkbox(label="📰News", value=False)

    with gr.Row():
      with gr.Column(scale=0.7):
        chatbot = gr.Chatbot(elem_id='chatbot')
      with gr.Column(scale=0.3):
        #with gr.Group():
        gen_audio = gr.Audio(label="generated audio")
        gen_image = gr.Image(label="generated image", type="filepath")

    with gr.Row():
      with gr.Column(scale=0.85):
        inputs = gr.Textbox(placeholder= "Hi there!", label= "Type an input and press Enter")
      with gr.Column(scale=0.15, min_width=0):
          btn = gr.UploadButton("📁Upload", file_types=["image", "audio"],  file_count="single")

    state = gr.State([]) #s
    b1 = gr.Button("🏃Run")

    with gr.Row():
      with gr.Accordion("Parameters", open=False):
        top_p = gr.Slider( minimum=-0, maximum=1.0, value=1.0, step=0.05, interactive=True, label="Top-p (nucleus sampling)",)
        temperature = gr.Slider( minimum=-0, maximum=5.0, value=1.0, step=0.1, interactive=True, label="Temperature",)
        chat_counter = gr.Number(value=0, visible=False, precision=0)
      with gr.Accordion("Files", open=False):
        file_output = gr.File(file_count="multiple", file_types=["image", "audio"])


    inputs.submit( predict,
                  [inputs, top_p, temperature, chat_counter, music_gen, stable_diff, image_cap, top_news, file_output, plugin_message, chatbot, state],
                   [chatbot, state, chat_counter, gen_audio, gen_image, plugin_message],)
    b1.click( predict,
            [inputs, top_p, temperature, chat_counter, music_gen, stable_diff, image_cap, top_news, file_output, plugin_message, chatbot, state],
             [chatbot, state, chat_counter, gen_audio, gen_image, plugin_message],)

    b1.click(reset_textbox, [], [inputs])
    inputs.submit(reset_textbox, [], [inputs])

    btn.upload(upload_file, [btn, file_output], file_output)
    gen_image.change(add_image, [gen_image, file_output], file_output)
    gen_audio.change(add_audio, [gen_audio, file_output], file_output)
    
    gr.HTML("""Bonus! Follow these steps for adding your own Plugins to this chatbot: <a href="https://huggingface.co/spaces/ysharma/ChatGPT-Plugins-in-Gradio/blob/main/README.md" target="_blank">How to add new Plugins in ChatGPT in 5 mins!!</a> or open the accordion given on top.""")

    
demo.queue(concurrency_count=2, max_size=10).launch(debug=True, height = '1000')