import gradio as gr
from huggingface_hub import InferenceClient
import os
import time

# Obtén el token de manera segura desde el entorno
hf_token = os.getenv("HF_API_TOKEN")

# Clase para manejar múltiples modelos
class ModelHandler:
    def __init__(self, model_names, token):
        self.clients = {model_name: InferenceClient(model_name, token=token) for model_name in model_names}
        self.current_model = model_names[0]

    def switch_model(self, model_name):
        if model_name in self.clients:
            self.current_model = model_name
        else:
            raise ValueError(f"Modelo {model_name} no está disponible.")

    def generate_response(self, input_text):
        prompt = f"Debes de responder a cualquier pregunta:\nPregunta: {input_text}"
        try:
            messages = [{"role": "user", "content": prompt}]
            client = self.clients[self.current_model]
            response = client.chat_completion(messages=messages, max_tokens=500)
            if hasattr(response, 'choices') and response.choices:
                return response.choices[0].message.content
            else:
                return str(response)
        except Exception as e:
            return f"Error al realizar la inferencia: {e}"

# Lista de modelos disponibles
model_names = [
    "microsoft/Phi-3-mini-4k-instruct"
]

# Inicializa el manejador de modelos
model_handler = ModelHandler(model_names, hf_token)

# Define la función para generación de imágenes con progreso
def generate_image_with_progress(prompt):
    """
    Genera una imagen utilizando el modelo de "stabilityai/stable-diffusion-2" y muestra un progreso.
    """
    try:
        client = InferenceClient("stabilityai/stable-diffusion-2", token=hf_token)
        
        # Simular progreso
        for progress in range(0, 101, 20):
            time.sleep(0.5)
            yield f"Generando imagen... {progress}% completado", None

        image = client.text_to_image(prompt, width=512, height=512)
        yield "Imagen generada con éxito", image
    except Exception as e:
        yield f"Error al generar la imagen: {e}", None

# Configura la interfaz en Gradio con selección de modelos y generación de imágenes
with gr.Blocks(title="Multi-Model LLM Chatbot with Image Generation") as demo:
    gr.Markdown(
        """
        ## Chatbot Multi-Modelo LLM con Generación de Imágenes
        Este chatbot permite elegir entre múltiples modelos de lenguaje para responder preguntas o generar imágenes 
        a partir de descripciones.
        """
    )
    with gr.Row():
        model_dropdown = gr.Dropdown(
            choices=model_names + ["Generación de Imágenes"],
            value=model_names[0],
            label="Seleccionar Acción/Modelo",
            interactive=True
        )
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                lines=5,
                placeholder="Escribe tu consulta o descripción para la imagen...",
                label="Entrada"
            )
        with gr.Column():
            output_display = gr.Textbox(
                lines=5,
                label="Estado",
                interactive=False
            )
            output_image = gr.Image(
                label="Imagen Generada",
                interactive=False
            )
    submit_button = gr.Button("Enviar")
    
    # Define la función de actualización
    def process_input(selected_action, user_input):
        try:
            if selected_action == "Generación de Imágenes":
                # Manejamos el generador de progreso
                progress_generator = generate_image_with_progress(user_input)
                last_status = None
                last_image = None
                for status, image in progress_generator:
                    last_status = status
                    last_image = image
                return last_status, last_image
            else:
                model_handler.switch_model(selected_action)
                response = model_handler.generate_response(user_input)
                return response, None
        except Exception as e:
            return f"Error: {e}", None
    
    # Conecta la función a los componentes
    submit_button.click(
        fn=process_input,
        inputs=[model_dropdown, input_text],
        outputs=[output_display, output_image]
    )

# Lanza la interfaz
demo.launch()