from langchain_google_genai import ChatGoogleGenerativeAI
from google import genai
from google.api_core import retry
from io import BytesIO
import gradio as gr
import base64
import time
import uuid
import os

# Define a retry policy
is_retriable = lambda e: (isinstance(e, genai.errors.APIError) and e.code in {429, 503})

if not hasattr(genai.models.Models.generate_content, '__wrapped__'):
  genai.models.Models.generate_content = retry.Retry(
      predicate=is_retriable)(genai.models.Models.generate_content)

os.environ["GOOGLE_API_KEY"] = os.getenv("GOOGLE_API_KEY")

vis_llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash-preview-image-generation")

# Image to Base64 Converter
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')

# Base64 to Image Converter
def base64_to_image(base64_img, filename):
  with open(filename, "wb") as f:
    f.write(base64.b64decode(base64_img))
  return filename

# Function that takes user inputs and displays it on chatUI
def query_message(message, history):
  if message["files"] is not None:
    for img in message["files"]:
      history.append({"role": "user", "content": {"path": img}})
  if message["text"] is not None:
    history.append({"role": "user", "content": message["text"]})
  return gr.MultimodalTextbox(value=None, interactive=False), history

# Function that takes user inputs, generate response and displays on chatUI
def chat_bot(history):
  messages = history.copy()

  # Format with system prompt
  system_prompt = """
                   system: You are CreativeMuse, a warm, imaginative, and encouraging creative assistant.
                   Your purpose is to help users overcome creative blocks and spark new ideas across various fields
                   — writing, design, music, visual arts, filmmaking, and more. When a user feels stuck, you offer unconventional prompts,
                   inspiring questions, or curated brainstorming techniques to get their creativity flowing again.

                   - Tailor your suggestions to their medium and mood, mixing practical strategies with whimsical, unexpected twists to encourage exploration.
                   Always be supportive, non-judgmental, and affirming—help them feel safe to experiment.
                   You can also suggest small creative exercises, constraints, "what if" scenarios, or draw connections between unrelated concepts to stimulate new thinking.
                   If a user shares part of a project, offer suggestions that expand on their work without taking it over.

                   - Keep your response short and straight to the point. You only generate an image on the user's request. 
                   """

  formatted_prompt = system_prompt + "\n\n"

  # Add conversation history
  for msg in messages:
    if "path" in msg["content"] and isinstance(msg["content"], dict):
      encoded_img = image_to_base64(msg["content"]["path"])
      formatted_prompt += f"{msg['role']}: data:image/jpeg.base64,{encoded_img}"
    else:
      formatted_prompt += f"{msg['role']}: {msg['content']}\n\n"

  # Get response from model
  response = vis_llm.invoke(formatted_prompt,
                            generation_config=dict(response_modalities=["TEXT", "IMAGE"]))

  # Add response to history
  history.append({"role": "assistant", "content": ""})

  for block in response.content:
    if isinstance(block, str):
      history[-1]["content"] += block
      time.sleep(0.02)
      yield history
    elif isinstance(block, dict) and block.get("image_url"):
      url = block["image_url"].get("url").split(",")[-1]
      image = base64_to_image(url, f"{uuid.uuid4()}.png")
      history.append({"role": "assistant", "content": {"path": image}})
      yield history

welcome_message = "Hey, I’m CreativeMuse—your creative companion when inspiration runs dry. Stuck on an idea? Let’s shake things up and spark something new together."

# Interface Code
with gr.Blocks() as app:
  chatbot = gr.Chatbot(
      value=[{"role": "assistant", "content": welcome_message}],
      label="CreativeMuse",
      show_label=True,
      type='messages',
      bubble_full_width=False,
      avatar_images=(None, "1eb05f325ec50a15c8b045f3428d6d5e.jpg")
    )
  text_box = gr.MultimodalTextbox(
          placeholder="Enter text or upload an image...",
          container=False,
          file_types=["image"],
          file_count="multiple",
          interactive=True,
      )

  chat_msg = text_box.submit(query_message, [text_box,chatbot], [text_box, chatbot])
  bot_msg = chat_msg.then(chat_bot, chatbot, chatbot)
  bot_msg.then(lambda: gr.MultimodalTextbox(interactive=True), None, text_box)

if __name__ == "__main__":
    app.queue() 
    app.launch(share=True)