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

# Cargar las variables de entorno
load_dotenv()
API_KEY = os.getenv("GOOGLE_API_KEY")

if not API_KEY:
    raise ValueError("La clave de API 'GOOGLE_API_KEY' no está configurada en el archivo .env")

# Configuración del modelo Gemini
generation_config = {
    "temperature": 1,
    "top_p": 0.95,
    "top_k": 40,
    "max_output_tokens": 8192,
    "response_mime_type": "text/plain",
}

genai.configure(api_key=API_KEY)

model = genai.GenerativeModel(
    model_name="gemini-1.5-flash",
    generation_config=generation_config,
)

# Inicializar la sesión de chat
chat = model.start_chat(history=[])

# Función para transformar el historial de Gradio al formato de 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

# Función de respuesta que maneja el texto y los archivos multimodales
def response(message, history):
    global chat

    # Transformar el historial al formato esperado por Gemini
    chat.history = transform_history(history)
    
    # Enviar el mensaje al modelo y obtener la respuesta
    response = chat.send_message(message["text"])
    response.resolve()

    # Mostrar la respuesta carácter por carácter
    for i in range(len(response.text)):
        time.sleep(0.01)
        yield response.text[: i + 1]

# Función para manejar las imágenes cargadas
def handle_uploaded_images(files):
    """Procesa las imágenes cargadas, las redimensiona y las guarda"""
    image_paths = []
    for file in files:
        image = Image.open(file).convert('RGB')  # Abrir y convertir la imagen a RGB
        image_height = int(image.height * 512 / image.width)  # Mantener la proporción
        image_resized = image.resize((512, image_height))  # Redimensionar la imagen
        image_filename = f"{uuid.uuid4()}.jpeg"  # Crear un nombre único para la imagen
        image_path = f"/tmp/{image_filename}"  # Ruta de almacenamiento temporal
        image_resized.save(image_path, "JPEG")  # Guardar la imagen
        image_paths.append(image_path)  # Agregar la ruta al resultado
    return image_paths

# Función para contar las imágenes cargadas
def count_images(message, history):
    num_images = len(message["files"])  # Contar las imágenes cargadas en el mensaje actual
    total_images = 0
    for msg in history:  # Contar todas las imágenes en el historial
        if isinstance(msg["content"], tuple):  # Si el contenido es una tupla, es un archivo
            total_images += 1
    return f"You just uploaded {num_images} images, total uploaded: {total_images + num_images}"

# Función que se invoca al cargar imágenes
def upload_images_handler(files, history):
    # Procesar las imágenes cargadas
    image_paths = handle_uploaded_images(files)

    # Actualizar el historial con las imágenes procesadas
    history.append((None, image_paths))

    return history, count_images({"files": files}, history)

# Crear la interfaz de Gradio
demo = gr.ChatInterface(
    fn=response,  # Función de chat para manejar texto y archivos
    examples=[  # Ejemplos iniciales de mensajes
        {"text": "No files", "files": []}
    ],
    cache_examples=True,
    multimodal=True,  # Activar la modalidad multimodal
    textbox=gr.MultimodalTextbox(  # Configuración del cuadro de texto multimodal
        file_count="multiple",  # Permitir múltiples archivos
        file_types=["image"],  # Aceptar solo imágenes
        sources=["upload", "microphone"]  # Fuentes de entrada: carga de archivos y micrófono
    )
)

# Conectar la función de carga de imágenes con el evento de Gradio
demo.add_event("file_upload", upload_images_handler)

# Iniciar la interfaz
demo.launch()