import time
import gradio as gr
import google.generativeai as genai
import os
from PIL import Image
from dotenv import load_dotenv


# Cargar clave API desde variables de entorno
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")

if not GOOGLE_API_KEY:
    raise ValueError("GOOGLE_API_KEY is not set in environment variables.")

genai.configure(api_key=GOOGLE_API_KEY)

# Transformar el historial a formato Gemini
def transform_history(history):
    new_history = []
    for chat in history:
        new_history.append({"parts": [{"text": chat[0]}], "role": "user"})
        new_history.append({"parts": [{"text": chat[1]}], "role": "model"})
    return new_history

def add_message(history, message):
    for x in message["files"]:
        history.append({"role": "user", "content": {"path": x}})
    if message["text"] is not None:
        history.append({"role": "user", "content": message["text"]})
    return history, gr.MultimodalTextbox(value=None, interactive=False)

def bot(history):
    chat_history = transform_history(history)
    response = genai.ChatCompletion.create(
        model="gemini-1.5-flash",  # Usar el modelo adecuado
        messages=chat_history + [{"role": "user", "content": history[-1][0]}],
        max_tokens=150,
    )

    reply = response['choices'][0]['message']['content']
    
    # Responder con cada carácter progresivamente
    for i in range(len(reply)):
        time.sleep(0.05)
        yield history + [{"role": "assistant", "content": reply[:i + 1]}]

with gr.Blocks() as demo:
    chatbot = gr.Chatbot(elem_id="chatbot", bubble_full_width=False, type="messages")
    
    chat_input = gr.MultimodalTextbox(
        interactive=True,
        file_count="multiple",
        placeholder="Escribe un mensaje o sube un archivo...",
        show_label=False,
        sources=["microphone", "upload"],
    )
    
    chat_msg = chat_input.submit(add_message, [chatbot, chat_input], [chatbot, chat_input])
    bot_msg = chat_msg.then(bot, chatbot, chatbot, api_name="bot_response")
    bot_msg.then(lambda: gr.MultimodalTextbox(interactive=True), None, [chat_input])

    chatbot.like(lambda x: print(x.index, x.value, x.liked), None, None, like_user_message=True)

demo.launch()