import torch
from transformers import pipeline
from huggingface_hub import login
import re
from runware import Runware, IPromptEnhance, IImageInference 
import asyncio
import gradio as gr
import os
from dotenv import load_dotenv

# Cargar las variables del archivo .env
load_dotenv()

# Obtener las claves desde las variables de entorno
hf_token = os.getenv("HF_TOKEN")
RUNWARE_API_KEY = os.getenv("RUNWARE_API_KEY")

# Variables de configuración del modelo
model_id = "meta-llama/Llama-3.2-1B-Instruct"

# Autenticación en Hugging Face
login(hf_token)

# Configuración del pipeline
pipe = pipeline(
    "text-generation",
    model=model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    temperature=0.2,  # Reducir creatividad
    top_p=0.8,
)

# Definir la función para generar texto mejorado
def generar_texto_mejorado(contenido_usuario, max_new_tokens=256):
    messages = [
        {"role": "system", "content": "Eres un editor profesional y experimentado que ayuda a escritores a mejorar su estilo de escritura. Tu objetivo es revisar el texto proporcionado por el usuario, mejorando su claridad, fluidez, gramática y precisión. Haz sugerencias que mantengan el tono y el estilo general del autor, pero que hagan el texto más atractivo y fácil de leer. Recuerda no agregar mas texto del que ya existe. Y al final del texto, introduce una lista con las mejoras que has realizado en el texto empezando con <mejoras>"},
        {"role": "user", "content": contenido_usuario},
    ]
    
    # Generar la respuesta con el modelo
    outputs = pipe(
        messages,
        max_new_tokens=max_new_tokens
    )
    
    # Extraer y retornar solo el contenido generado por el asistente
    texto_generado = outputs[0]["generated_text"]
    
    # Buscar y extraer únicamente el contenido después del primer ':', omitiendo información redundante
    # Ajustamos el contenido generado al esquema esperado
    if isinstance(texto_generado, list) and "content" in texto_generado[-1]:
        texto_generado = texto_generado[-1]["content"]
    
    # Buscar el contenido dentro de las etiquetas <mejoras> y extraerlo
        match = re.search(r"<mejoras>", texto_generado, re.IGNORECASE)
    
    if match:
        # Separar el contenido principal y las mejoras
        contenido_principal = texto_generado[:match.start()].strip()
        mejoras = texto_generado[match.end():].strip()  # Tomar desde el final de <mejoras> hasta el final del texto
        return contenido_principal, mejoras
    else:
        mejoras="No se encontraron mejoras en el texto."
        return texto_generado, mejoras

# Definir la función para generar texto mejorado
def generar_prompt_imagen(texto_generado, max_new_tokens=256):
    messages = [
        {"role": "system", "content": "Utilizando el siguiente texto de un artículo de noticia, crea un prompt visual detallado para generar una imagen que represente de manera precisa el tema y tono del artículo. Asegúrate de capturar los aspectos clave mencionados en el artículo, como el entorno, las personas, el contexto o los eventos, y expresar visualmente los sentimientos o conceptos que se mencionan. Mantén en cuenta la estética y los detalles descriptivos de lo que se discute"},
        {"role": "user", "content": texto_generado},
    ]
    
    # Generar la respuesta con el modelo
    outputs = pipe(
        messages,
        max_new_tokens=max_new_tokens
    )
    
    # Extraer y retornar solo el contenido generado por el asistente
    prompt_generado = outputs[0]["generated_text"]
    
    # Buscar y extraer únicamente el contenido después del primer ':', omitiendo información redundante
    # Ajustamos el contenido generado al esquema esperado
    if isinstance(prompt_generado, list) and "content" in prompt_generado[-1]:
        prompt_generado = prompt_generado[-1]["content"]
    return prompt_generado

# Función para generar imagen desde texto usando la API de Runware
async def generar_imagen(prompt_generado):
    if not (3 <= len(prompt_generado) <= 2000):
        return "Error: El texto debe tener entre 3 y 2000 caracteres."

    runware = Runware(api_key=RUNWARE_API_KEY)
    await runware.connect()

    request_image = IImageInference(
        positivePrompt=prompt_generado,
        model="civitai:36520@76907",
        numberResults=1,
        negativePrompt="cloudy, rainy",
        height=512,
        width=512,
    )

    images = await runware.imageInference(requestImage=request_image)
    if images:
        return images[0].imageURL
    else:
        return "No se generó ninguna imagen."


# Función principal que se pasará a Gradio
def interfaz_gradio(text):
    # Ejecutar la tarea asincrónica de procesar el texto y generar la imagen
    texto_mejorado, mejoras= generar_texto_mejorado(text)
    prompt_generado = generar_prompt_imagen(texto_mejorado)
    imagen_url = asyncio.run(generar_imagen(prompt_generado))

    # Retornar los resultados para mostrarlos en la interfaz
    return texto_mejorado, mejoras, imagen_url


iface = gr.Interface(
    fn=interfaz_gradio,  # Función que será llamada al interactuar con la interfaz
    inputs=gr.Textbox(label="Texto para mejorar"),  # Campo de texto para el usuario
    outputs=[
        gr.Textbox(label="Texto Mejorado"),  # Mostrar el texto mejorado
        gr.Textbox(label="Mejoras realizadas"),  # Mostrar las mejoras
        gr.Image(label="Imagen Generada")  # Mostrar la imagen generada
    ],
    live=False,  # Desactivar la actualización en tiempo real
    allow_flagging="never",  # Opcional: desactivar las banderas de los usuarios
)


# Iniciar la interfaz
iface.launch()