import os
import time
import itertools
import gradio as gr
import google.generativeai as genai

# Credentials
genai.configure(api_key=os.getenv('PALM_API_KEY'))

# Gradio
chat_defaults = {
  'model': 'models/chat-bison-001',
  'temperature': 0.25,
  'candidate_count': 1,
  'top_k': 40,
  'top_p': 0,
}

chat_history = []

def generate_chat(prompt: str, chat_messages=chat_history):
    print(chat_messages)
    context = "You are an intelligent chatbot powered by biggest technology company."
    print("Generating Chat Message...")
    print(f"User Message:\n{prompt}\n")
    chat_messages.append(prompt)
    try:
        response = genai.chat(
          **chat_defaults,
          context=context,
          messages=chat_messages
        )
        result = response.last
        if result is None:
            result = "Apologies but something went wrong. Please try again later."
            chat_messages = chat_messages[:-1]
        else:
            chat_messages.append(result)
    except Exception as e:
        result = "Apologies but something went wrong. Please try again later."
        chat_messages = chat_messages[:-1]
        print(f"Exception {e} occured\n")
    chat_history = chat_messages
    print(f"Bot Message:\n{result}\n")
    return result

with gr.Blocks() as app:
    chatbot = gr.Chatbot(height=400, bubble_full_width=False, container=False)
    msg = gr.Textbox(label="Type your message...", value="Hi Gerard, can you introduce yourself?", container=False)
    clear = gr.Button("Clear")

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

    def bot(history):
        bot_message = generate_chat(history[-1][0])
        history[-1][1] = ""
        for character in bot_message:
            history[-1][1] += character
            time.sleep(0.01)
            yield history

    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, chatbot, chatbot
    )
    clear.click(lambda: None, None, chatbot, queue=False)
    
    gr.Markdown(
    f"""
    Hosted on 🤗 Spaces. Powered by Gradio & Google PaLM.
    """)
    
app.queue()
app.launch()