# ✅ OBLIGATOIRE EN PREMIER
import spaces

import torch
from diffusers import Flux2KleinPipeline
import gradio as gr
import os
from datetime import datetime
from ddgs import DDGS
import requests
from PIL import Image
from io import BytesIO
import re

# ⚠️ pas de empty_cache avant spaces

# =========================
# CSS
# =========================
custom_css = """
* { font-family: Arial, sans-serif !important; }
html, body, .gradio-container { background-color: #0f172a !important; }
h1, h2, h3, label, span, p { color: #22c55e !important; }
input, textarea, select {
    background-color: #1f2937 !important;
    color: white !important;
    border-radius: 8px !important;
    border: 1px solid #22c55e !important;
}
button { background: black !important; color: #22c55e !important; }
"""

# =========================
# PIPELINE
# =========================
pipe = Flux2KleinPipeline.from_pretrained(
    "lea97338/KTXFlux-2.0",
    token=os.environ["HF_Token"],  # ✅ en Space, token inutile si public
    cache_dir="./model",
    torch_dtype=torch.float16,
)

# ✅ meilleur mode pour Spaces
pipe.enable_model_cpu_offload()
pipe.enable_attention_slicing()

OUTPUT_DIR = "./outputs"

# =========================
# CLEAN TEXT
# =========================
def clean_text(text):
    text = re.sub(r"#\w+|http\S+|\d+", "", text.lower())
    return " ".join([w for w in text.split() if len(w) > 3][:10])

# =========================
# DDGS (light pour Spaces)
# =========================
def ddgs_search(prompt):

    texts = []

    try:
        with DDGS() as ddgs:
            for r in ddgs.text(prompt, max_results=2):
                t = clean_text(r.get("body", ""))
                if t:
                    texts.append(t)
    except:
        pass

    return " ".join(texts)

# =========================
# RESOLUTION
# =========================
def get_dimensions(res):
    return {
        "256x256": (256,256),
        "512x512": (512,512),
        "Panorama": (1024,512),
    }[res]

# =========================
# GENERATE
# =========================
@spaces.GPU
def generate_image(prompt, steps, guidance, resolution):

    width, height = get_dimensions(resolution)
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    extra_text = ddgs_search(prompt)
    final_prompt = f"{prompt}, {extra_text}"

    result = pipe(
        prompt=final_prompt,
        height=height,
        width=width,
        num_inference_steps=steps,   # ✅ obligatoire pour vitesse
        guidance_scale=float(guidance),
    )

    image = result.images[0]

    path = os.path.join(
        OUTPUT_DIR,
        datetime.now().strftime("img_%Y%m%d_%H%M%S.png")
    )
    image.save(path)

    return image, path

# =========================
# UI
# =========================
with gr.Blocks(css=custom_css) as demo:

    gr.Markdown("# ⚡ Flux Generator (HF Spaces optimisé)")

    with gr.Row():
        with gr.Column():

            prompt = gr.Textbox(value="minecraft village", label="Prompt")

            steps = gr.Slider(1, 30, 1, step=1, label="Steps")
            guidance = gr.Slider(1.0, 5.0, 2.5, step=0.5, label="Guidance")

            resolution = gr.Dropdown(
                ["256x256","512x512","Panorama"],
                value="512x512",
                label="Résolution"
            )

            btn = gr.Button("⚡ Générer")
            path_out = gr.Textbox(label="Fichier")

        with gr.Column():
            image_out = gr.Image(label="Résultat")

    btn.click(
        generate_image,
        inputs=[prompt, steps, guidance, resolution],
        outputs=[image_out, path_out]
    )

# =========================
# RUN (🔥 important)
# =========================
if __name__ == "__main__":
    demo.launch()  # ❌ PAS de share=True sur Spaces