# import required modules
import os
import time
import base64
import asyncio
import gradio as gr
from PIL import Image
from google import genai
from threading import Lock
from google.genai import types

# define which Gemini model version is going to be used
model_id = "gemini-2.0-flash-exp"

# Gemini model system instruction settings
system_instruction="""
    "You always reply in the same language the user sent the question. It is mandatory.",
    "You only change the response language if explicitly asked - otherwise, answer in the original language."
    "You are an assistant who helps people with their questions.",
    "You only provide answers in one paragraph or less.",
    "Your answers are long enough to not miss any information.",
    "You are always kind and use simple, pleasant language.",
"""


## helper functions
# convert image files to base64 data
def image_to_base64(image_path):
  with open(image_path, 'rb') as img:
    encoded_string = base64.b64encode(img.read())
  return encoded_string.decode('utf-8')


# show user message at the chatbot history
def query_message(history,txt,img):
 if not img:
  history += [(txt,None)]
  return history
 if img:
  base64 = image_to_base64(img)
  data_url = f"data:image/jpeg;base64,{base64}"
  history += [(f"{txt} ![]({data_url})", None)]
  return history

## gradio interface
# gradio page variables
TITLE = """<h1 align="center">Gemini 2.0 Chatbot 🤖</h1>"""
SUBTITLE = """<h2 align="center">A multimodal chatbot powered by Gradio and Gemini API</h2>"""

# gradio styles in css
css = """
.container {
    max-width: 100%;
    padding: 0 1rem;
}

.chatbot {
    height: calc(100vh - 250px) !important;
    overflow-y: auto;
}

.textbox {
    margin-top: 0.5rem;
}
"""

# gradio chatbot main function
def registry(name, token, examples=None, **kwargs):
  client = genai.Client(api_key=token)
  chat_locks = {}  # Dictionary to hold locks for each user's chat
  chat_sessions = {} # Dictionary to hold each user chat

  def create_chat():
    return client.chats.create(
      model=name,
      config=types.GenerateContentConfig(
        system_instruction=system_instruction,
        temperature=0.5,
      ),
    )

  # send a user message to Gemini, streams the response back to the chatbot
  # and updates the history
  def stream_response(history, text, img, request: gr.Request):
    user_id = request.client.host
    if user_id not in chat_locks:
      chat_locks[user_id] = Lock()
      chat_sessions[user_id] = create_chat()

    lock = chat_locks[user_id]
    chat = chat_sessions[user_id]

    try:
      with lock:
        if not img:
          response_stream = chat.send_message_stream(
            text
          )
        else:
          try:
            img = Image.open(img)
            response_stream = chat.send_message_stream(
              [text, img]
            )
          except Exception as e:
            print(f"Error processing image: {str(e)}")
            return

        # Initialize response text
        response_text = ""

        # Stream the response
        for chunk in response_stream:
          if chunk.text:
            response_text += chunk.text
            # Update the last message in history with the new content
            history[-1] = (history[-1][0], response_text)
            yield history

    except Exception as e:
      print(f"Error in stream_response: {str(e)}")
      return

  print("Building the gradio app...")
  with gr.Blocks(css=css) as app:
    gr.HTML(TITLE)
    gr.HTML(SUBTITLE)

    with gr.Row():
      image_box = gr.Image(type="filepath")
      chatbot = gr.Chatbot(
        scale=2,
        height=500,
        container=True
      )

    text_box = gr.Textbox(
      placeholder="Type your message and press enter and optionally upload an image",
      container=False,
    )

    btn = gr.Button("Send")

    # Update the event handlers to use streaming
    btn.click(
      fn=query_message,
      inputs=[chatbot, text_box, image_box],
      outputs=[chatbot],
    ).then(
      fn=stream_response,
      inputs=[chatbot, text_box, image_box],
      outputs=[chatbot],
      api_name="stream_response"
    ).then(
      fn=lambda: (None, ""),  # Clear the image and text inputs after sending
      inputs=None,
      outputs=[image_box, text_box],
    )

    # Add enter key handler
    text_box.submit(
      fn=query_message,
      inputs=[chatbot, text_box, image_box],
      outputs=[chatbot],
    ).then(
      fn=stream_response,
      inputs=[chatbot, text_box, image_box],
      outputs=[chatbot],
      api_name="stream_response"
    ).then(
      fn=lambda: (None, ""),  # Clear the image and text inputs after sending
      inputs=None,
      outputs=[image_box, text_box],
    )

  return app

if __name__ == "__main__":
  # launch the gradio chatbot
  gr.load(
  name=model_id,
  src=registry,
  accept_token=True
).launch()