import gradio as gr
import numpy as np
import random
import spaces
import os
import torch
import re
from PIL import Image

from diffusers import DiffusionPipeline, AutoencoderTiny
from huggingface_hub import login
from live_preview_helpers import flux_pipe_call_that_returns_an_iterable_of_images

# Ensure image_preview dir exists
os.makedirs("image_preview", exist_ok=True)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

# ✅ DO NOT CHANGE: Working pipeline using taef1
taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)

pipe = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    torch_dtype=dtype,
    vae=taef1
).to(device)

pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
pipe.load_lora_weights("ZennyKenny/flux_lora_natalie-diffusion")


def sanitize_filename(name):
    return re.sub(r"[^a-zA-Z0-9_-]", "_", name)[:80]


@spaces.GPU(duration=75)
def infer(user_token, prompt, seed=42, randomize_seed=False, width=1024, height=1024,
          guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
    login(token=user_token)

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)

    full_prompt = f"XTON {prompt}"

    for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
        prompt=full_prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        generator=generator,
        output_type="pil",
    ):
        safe_name = sanitize_filename(prompt)
        img_path = f"image_preview/{safe_name}_{seed}.jpg"
        img.convert("RGB").save(img_path, "JPEG", quality=60)

        previews = [f"image_preview/{f}" for f in sorted(os.listdir("image_preview")) if f.endswith(".jpg")]
        return img, seed, previews

# Wrapper to inject a fallback token if needed
def infer_with_fallback_token(user_token, prompt, *args):
    if not user_token.strip():
        user_token = "your_token_here"  # Replace with a real test token for dev, not in production
    return infer(user_token, prompt, *args)

# Prompt-only examples; token will be filled in by wrapper
prompt_examples = [
    "a man walking in the forest",
    "a viking ship sailing down a river",
    "a woman resting by an open fire",
    "a sword fight in a medieval village"
]

with gr.Blocks(css="style.css") as natalie_diffusion:
    with gr.Row():
        with gr.Column(scale=1, elem_id="left-column"):
            gr.Markdown("""
# ХТОНЬ: Natalie LoRA Image Generator

Generate images in the surreal style of artist [Natalie Kav](https://www.behance.net/nataliKav), adapted using a custom LoRA on the FLUX.1 [dev] model.

> This space is designed for prototyping concept art for a forthcoming game called **ХТОНЬ**. All outputs are generated locally in the browser using GPU acceleration.
""")

            hf_token_input = gr.Textbox(
                label="Your Hugging Face API Token",
                placeholder="Paste your token here",
                type="password"
            )

            with gr.Row():
                prompt = gr.Text(
                    label="Prompt",
                    show_label=False,
                    max_lines=1,
                    placeholder="Enter your prompt...",
                    container=False,
                )
                run_button = gr.Button("🎨 Generate", scale=0)

            with gr.Accordion("Advanced Settings", open=False):
                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

                with gr.Row():
                    width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
                    height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)

                with gr.Row():
                    guidance_scale = gr.Slider(label="Guidance Scale", minimum=1, maximum=15, step=0.1, value=3.5)
                    num_inference_steps = gr.Slider(label="Number of inference steps", minimum=1, maximum=50, step=1, value=28)

            result_example = gr.Image(visible=False)

            gr.Examples(
                examples=[[prompt] for prompt in prompt_examples],
                fn=lambda prompt: infer_with_fallback_token("", prompt),
                inputs=[prompt],
                outputs=[result_example, seed, gr.Gallery(visible=False)],
                cache_examples=False,
            )

        with gr.Column(scale=1, elem_id="right-column"):
            result = gr.Image(label="", show_label=False, elem_id="generated-image")

    with gr.Column():
        gr.Markdown("<h3 style='text-align:center;'>Generated Images Preview</h3>")
        gallery = gr.Gallery(label="", columns=4, height="auto", object_fit="cover")

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer_with_fallback_token,
        inputs=[hf_token_input, prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
        outputs=[result, seed, gallery],
    )

if __name__ == "__main__":
    natalie_diffusion.launch()