import gradio as gr
from huggingface_hub import InferenceClient
import os

# --- PASO 1: AUTENTICACIÓN ---
# Leemos el Token que guardaste en los "Secrets" del Space
hf_token = os.environ.get("HF_TOKEN")

# Verificación de que el token existe
if not hf_token:
    raise ValueError(
    )

# --- PASO 2: DEFINIR NUESTROS "MOTORES" (CLIENTES DE API) ---

# MOTOR 1: El "Artista de Prompts" (LLM)
# Usamos el modelo 'Instruct' de Llama 3 8B. Es el estándar de la industria.
llm_client = InferenceClient(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    token=hf_token
)

# MOTOR 2: El "Artista de Imágenes" (Texto a Imagen)
# Usamos el modelo base de Stable Diffusion XL 1.0. Es potente y fiable.
image_client = InferenceClient(
    "stabilityai/stable-diffusion-xl-base-1.0",
    token=hf_token
)

# --- PASO 3: LA LÓGICA DE "ENCADENAMIENTO" ---

def generate_magic_image(simple_prompt):
    
    # --- ETAPA 1: MEJORAR EL PROMPT USANDO LLAMA 3 ---
    
    # Esta es la instrucción "secreta" que le damos a Llama 3
    system_prompt = """
    Tú eres un 'Prompt Artist' experto. Tu único trabajo es reescribir un prompt
    de usuario simple y convertirlo en una descripción visual vívida, detallada y
    fotorrealista para un modelo de generación de imágenes.
    
    REGLAS:
    - Responde SÓLO con el prompt mejorado.
    - No saludes ni digas nada más.
    - El prompt mejorado debe estar en INGLÉS.
    - Ejemplo: si el usuario dice "un gato", tú respondes:
      "A cinematic, photorealistic portrait of a majestic black cat with 
      piercing green eyes, studio lighting, 8k, ultra-detailed texture"
    """
    
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": simple_prompt}
    ]
    
    print(f"[Debug] Prompt simple recibido: {simple_prompt}")
    
    try:
        # --- Llamada al Modelo 1 (LLM) ---
        response = llm_client.chat_completion(
            messages=messages,
            max_tokens=200, 
            temperature=0.7
        )
        enhanced_prompt = response.choices[0].message.content.strip()
        
    except Exception as e:
        print(f"[Error] Error con el LLM: {e}")
        # Si Llama falla, al menos no rompemos la app
        raise gr.Error(f"Error al contactar con Llama 3: {e}")
    
    print(f"[Debug] Prompt mejorado: {enhanced_prompt}")
    
    try:
        # Usamos el prompt que nos dio Llama 3
        image_bytes = image_client.text_to_image(
            prompt=enhanced_prompt,
            model="stabilityai/stable-diffusion-xl-base-1.0", # Especificamos el modelo por si acaso
            num_inference_steps=25,
            negative_prompt="blurry, ugly, deformed, text, watermark" # Evitar resultados feos
        )
        
        print("[Debug] Imagen generada con éxito.")
        
    except Exception as e:
        print(f"[Error] Error con Stable Diffusion: {e}")
        raise gr.Error(f"Error al generar la imagen con SDXL: {e}")
    
    return enhanced_prompt, image_bytes


with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown("# 🎨 Creador de Arte Mágico (Llama 3 + SDXL) 🚀")
    gr.Markdown("Escribe una idea simple. La IA la convertirá en un 'prompt' de artista y generará una obra maestra.")
    
    with gr.Row():
        with gr.Column(scale=1):
            simple_prompt_input = gr.Textbox(
                label="Tu idea simple", 
                placeholder="Ej: un perro en el espacio"
            )
            submit_btn = gr.Button("¡Crear Magia!", variant="primary")
        
        with gr.Column(scale=2):
            enhanced_prompt_output = gr.Textbox(
                label="El Prompt Mejorado por Llama 3 (La magia)", 
                interactive=False
            )
            image_output = gr.Image(label="Tu Obra Maestra (SDXL 1.0)", type="pil")
    
    gr.Examples(
        ["Un gato con sombrero", "Un bosque por la noche", "Madrid en el futuro"],
        inputs=simple_prompt_input
    )
    
    submit_btn.click(
        fn=generate_magic_image,
        inputs=simple_prompt_input,
        outputs=[enhanced_prompt_output, image_output]
    )

app.launch()