# -*- coding: utf-8 -*-
"""
GTP-FLUX-BAT: Chat + Imágenes con Hugging Face (100% GRATIS)
- Solo modelos HF permitidos
- Bloqueo total a otras APIs
- Genera imágenes con FLUX.1
- Chat con gpt2-large
"""

import gradio as gr
import requests
import io
from PIL import Image
import os
from urllib.parse import urlparse

# --- VALIDACIÓN DE SEGURIDAD ---
def validate_hf_only(endpoint: str):
    allowed = ["api-inference.huggingface.co", "huggingface.co"]
    hostname = urlparse(endpoint).hostname
    if hostname not in allowed:
        raise ValueError("BLOQUEO: Solo se permite Hugging Face")

# --- MODELOS PERMITIDOS ---
HF_CHAT_MODEL = "gpt2-large"
HF_IMAGE_MODELS = [
    "black-forest-labs/FLUX.1-dev",
    "black-forest-labs/FLUX.1-schnell",
    "black-forest-labs/FLUX.1-schnell-4bit"
]

# --- GENERAR TEXTO CON HF ---
def generate_text_hf(prompt: str, hf_token: str):
    validate_hf_only("/static-proxy?url=https%3A%2F%2Fapi-inference.huggingface.co")
    url = f"/static-proxy?url=https%3A%2F%2Fapi-inference.huggingface.co%2Fmodels%2F%7BHF_CHAT_MODEL%7D"
    headers = {"Authorization": f"Bearer {hf_token}"}
    payload = {"inputs": prompt, "parameters": {"max_length": 100}}
    try:
        response = requests.post(url, headers=headers, json=payload, timeout=30)
        if response.status_code == 200:
            return response.json()[0]["generated_text"]
    except:
        pass
    return "No pude generar texto. Intenta más tarde."# --- GENERAR IMAGEN CON HF (solo lista blanca) ---
def generate_image_hf_only(prompt: str, hf_token: str):
    validate_hf_only("/static-proxy?url=https%3A%2F%2Fapi-inference.huggingface.co")
    for model in HF_IMAGE_MODELS:
        try:
            url = f"/static-proxy?url=https%3A%2F%2Fapi-inference.huggingface.co%2Fmodels%2F%7Bmodel%7D"
            headers = {"Authorization": f"Bearer {hf_token}"}
            payload = {
                "inputs": prompt,
                "parameters": {
                    "num_inference_steps": 20,
                    "guidance_scale": 7.5,
                    "width": 1024,
                    "height": 1024
                }
            }
            response = requests.post(url, headers=headers, json=payload, timeout=60)
            if response.status_code == 200:
                return response.content  # bytes
        except Exception as e:
            print(f"Modelo {model} falló: {e}")
            continue
    raise RuntimeError("Todos los modelos de imagen HF están temporalmente no disponibles")# --- DETECCIÓN DE IMAGEN ---
def needs_image(user_input: str) -> bool:
    keywords = ["imagen", "genera", "dibuja", "foto", "paisaje", "retrato", "crea"]
    return any(k in user_input.lower() for k in keywords)

# --- RESPUESTA COMPLETA ---
def respond_to_user(message: str, history: list, hf_token: str):
    if not hf_token or not hf_token.startswith("hf_"):
        return history + [(message, "Ingresa un token válido de Hugging Face.")], None

    bot_response = generate_text_hf(message, hf_token)
    image_bytes = None

    if needs_image(message):
        try:
            bot_response += "\n\nGenerando imagen..."
            image_bytes = generate_image_hf_only(message, hf_token)
            bot_response = bot_response.replace("Generando imagen...", "Imagen generada:")
        except Exception as e:
            bot_response += f"\n\n{str(e)}"

    history.append((message, bot_response))
    return history, image_bytes# --- INTERFAZ GRADIO ---
with gr.Blocks(title="GTP-FLUX-BAT", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # GTP-FLUX-BAT  
    **Chat + Imágenes 100% GRATIS con Hugging Face**  
    Modelos: `gpt2-large` + `FLUX.1-dev/schnell`  
    **SOLO Hugging Face · Nunca Replicate/OpenAI**
    """)

    with gr.Row():
        token_input = gr.Textbox(
            label="Token de Hugging Face",
            placeholder="hf_xxxxxxxxxxxxxxxxxxxxx",
            type="password",
            scale=3
        )
        gr.Markdown("**Obtén tu token en [huggingface.co/settings/tokens](https://huggingface.co/settings/tokens)**")

    chatbot = gr.Chatbot(height=500)
    msg = gr.Textbox(label="Tu mensaje", placeholder="Escribe aquí...", scale=4)
    clear = gr.Button("Limpiar chat")
    img_out = gr.Image(label="Imagen generada", height=400)

    # --- EVENTOS ---
    def submit(message, history, token):
        if not message.strip():
            return history, None
        new_history, img = respond_to_user(message, history, token)
        return new_history, "", img

    msg.submit(submit, [msg, chatbot, token_input], [chatbot, msg, img_out])
    clear.click(lambda: ([], None), outputs=[chatbot, img_out])# --- MENSAJE DE INICIO ---
    gr.Markdown("""
    ### Instrucciones:
    1. Pega tu **token de Hugging Face** arriba  
    2. Escribe: _"Genera un paisaje al atardecer"_  
    3. ¡El bot responde y genera la imagen automáticamente!  

    **Modelos usados:**  
    - Texto: `gpt2-large`  
    - Imagen: `FLUX.1-dev` → `schnell` → `4bit` (solo HF)  
    """)

# --- LANZAR ---
if __name__ == "__main__":
    demo.launch(share=True)