import os
import openai
import gradio as gr

openai.api_key = os.getenv('OPEN_AI_KEY')
hf_api_key = os.getenv('HF_TOKEN_KEY')

def predict(message, history):
    history_openai_format = []
    for human, assistant in history:
        history_openai_format.append({"role": "user", "content": human })
        history_openai_format.append({"role": "assistant", "content":assistant})
    history_openai_format.append({"role": "user", "content": message})

    response = openai.ChatCompletion.create(
        model='gpt-3.5-turbo',
        messages= history_openai_format,
        temperature=1.0,
        stream=True
    )

    partial_message = ""
    for chunk in response:
        if len(chunk['choices'][0]['delta']) != 0:
            partial_message = partial_message + chunk['choices'][0]['delta']['content']
            yield partial_message

A1 = gr.ChatInterface(predict,
                title="PeachChat+Plus",
                description="An AI Powered Productivity Assistant with Vision and Image Generation Capabilities Built By Peach State Innovation and Technology. Select The Corresponding Tab For Tool Accessibility",
                textbox=gr.Textbox(placeholder="Enter your question/prompt here..."),
                theme= gr.themes.Glass(primary_hue="amber", neutral_hue="lime"),
                retry_btn=None,
                clear_btn="Clear")

A2 = gr.load(
             "huggingface/Salesforce/blip-image-captioning-large",
              title="20/20-AI",
              description="Take and Upload a Photo or Existing Image, I'll Give You Its Description",
              outputs=[gr.Textbox(label="I see...")],
              theme= gr.themes.Glass(primary_hue="amber", neutral_hue="lime"))

A3 = gr.load(
             "huggingface/stabilityai/stable-diffusion-xl-base-1.0",
              inputs=[gr.Textbox(label="Enter Your Image Description")],
              outputs=[gr.Image(label="Image")],
              title="Reverie",
              description="Bring Your Imagination Into Existence On The Digital Canvas With Reverie, Powered With Stable Diffusion",
              allow_flagging="never", 
              examples=["A monster wandering the streets of downtown Atlanta","A robot in a Brazilian favela"])

pcp = gr.TabbedInterface([A1, A2, A3], ["Chat", "Describe", "Create"], theme= gr.themes.Glass(primary_hue="amber", neutral_hue="lime"))
pcp.queue().launch()