app/agent.py

import base64
import os
from io import BytesIO
from PIL import Image
from pydub import AudioSegment
from pydub.playback import play
import openai
from app.utils import tools, gradio_to_openai_history
from app.config import MODEL_TEXT, MODEL_IMAGE, MODEL_AUDIO, OPENAI_API_KEY, SYSTEM_MESSAGE
from app.tools import handle_tool_call  # Asegúrate de que esté definido ahí

def initialize_agents():
    openai.api_key = os.environ["OPENAI_API_KEY"]
    return {
        "chat": chat,
        "artist": artist,
        "talker": talker
    }


def artist(city):
    image_response = openai.images.generate(
            model="dall-e-3",
            prompt=f"Una imagen que representa unas vacaciones en {city}, mostrando lugares turísticos y todo lo único de {city}, en un vibrante estilo pop-art",
            size="1024x1024",
            n=1,
            response_format="b64_json",
        )
    image_base64 = image_response.data[0].b64_json
    image_data = base64.b64decode(image_base64)
    return Image.open(BytesIO(image_data))
    
def talker(message):
    if not message:
        print("Mensaje vacío, no se genera audio")
        return None

    # Generar audio con OpenAI
    response = openai.audio.speech.create(
        model="tts-1",
        voice="onyx",
        input=message
    )

    # Guardar en BytesIO en formato mp3
    audio_stream = BytesIO(response.content)
    audio = AudioSegment.from_file(audio_stream, format="mp3")
    audio_bytes = BytesIO()
    audio.export(audio_bytes, format="mp3")
    audio_bytes.seek(0)

    return audio_bytes

def chat(gr_history, user_msg):
    """
    Función principal del chatbot.

    Args:
        gr_history: historial de Gradio, lista de tuplas (user_msg, assistant_msg)
        user_msg: mensaje actual del usuario (string)

    Returns:
        gr_history actualizado, imagen generada (PIL.Image o None), audio_bytes (BytesIO o None)
    """
    # Convertir el historial de Gradio al formato esperado por OpenAI
    messages = gradio_to_openai_history(gr_history)
    
    # Agregar el mensaje del usuario actual
    messages.append({"role": "user", "content": user_msg})

    # Llamada al modelo de chat de OpenAI
    response = openai.chat.completions.create(
        model=MODEL_TEXT,
        messages=messages,
        functions=tools,        # Lista de funciones tipo dict
        function_call="auto"
    )

    image = None

    # Manejo de tool calls
    if response.choices[0].finish_reason == "tool_calls":
        message = response.choices[0].message
        tool_response, city = handle_tool_call(message)
        
        # Agregar los mensajes correctamente al historial de OpenAI
        messages.append({"role": "assistant", "content": message.get("content", "")})
        messages.append({"role": "assistant", "content": tool_response})
        
        # Generar imagen
        image = artist(city)
        
        # Volver a llamar al modelo para la respuesta final
        response = openai.chat.completions.create(model=MODEL_TEXT, messages=messages)

    # Extraer la respuesta final del asistente
    reply = response.choices[0].message.content

    # Actualizar el historial de Gradio
    gr_history.append((user_msg, reply))

    # Generar audio como BytesIO
    audio_bytes = talker(reply)

    return gr_history, image, audio_bytes