import gradio as gr
from gradio_client import Client, handle_file
import os
import tempfile
import shutil

HF_TOKEN = os.environ.get("HF_TOKEN", "")

REMOTE_SPACE = "signsur4739379373/qwen-image-edit-rapid-aio-nsfw-v23"

_client = None

def get_client():
    global _client
    if _client is None:
        _client = Client(REMOTE_SPACE, hf_token=HF_TOKEN)
    return _client

def extract_paths(images):
    paths = []
    if not images:
        return paths
    for item in images:
        if isinstance(item, dict) and "path" in item:
            paths.append(item["path"])
        elif isinstance(item, dict) and "url" in item:
            paths.append(item["url"])
        elif isinstance(item, (list, tuple)) and len(item) > 0:
            if isinstance(item[0], dict) and "path" in item[0]:
                paths.append(item[0]["path"])
            elif isinstance(item[0], dict) and "url" in item[0]:
                paths.append(item[0]["url"])
            elif isinstance(item[0], str):
                paths.append(item[0])
        elif isinstance(item, str):
            paths.append(item)
    return paths

def generate(images, prompt, steps, guidance, seed, randomize, rewrite, width, height, chain_in):
    if not images:
        raise gr.Error("Upload at least one source image first.")
    if not prompt or not prompt.strip():
        raise gr.Error("Write an edit prompt first.")

    img_paths = extract_paths(images)
    if not img_paths:
        raise gr.Error("Could not extract image file paths from upload.")

    # Convert dims: 256 means auto/None
    w = width if width and width > 256 else None
    h = height if height and height > 256 else None

    try:
        client = get_client()

        # Try named API first, then positional
        result = None
        last_err = None

        try:
            result = client.predict(
                images=img_paths,
                prompt=prompt.strip(),
                seed=int(seed) if seed else 0,
                randomize_seed=bool(randomize),
                num_inference_steps=int(steps),
                true_guidance_scale=float(guidance),
                rewrite_prompt=bool(rewrite),
                width=w,
                height=h,
                api_name="/infer",
                hf_token=HF_TOKEN,
            )
        except Exception as e1:
            last_err = e1
            try:
                # Positional fallback
                result = client.predict(
                    img_paths,
                    prompt.strip(),
                    int(seed) if seed else 0,
                    bool(randomize),
                    int(steps),
                    float(guidance),
                    bool(rewrite),
                    w,
                    h,
                )
            except Exception as e2:
                last_err = e2

        if result is None:
            raise gr.Error(f"Remote Space call failed: {last_err}")

        # Parse result - expected format: (gallery_images, seed, ui_update)
        gen_images = None
        used_seed = seed

        if isinstance(result, (list, tuple)):
            if len(result) >= 1:
                gen_images = result[0]
            if len(result) >= 2:
                used_seed = result[1]
        elif isinstance(result, dict):
            gen_images = result.get("images") or result.get("gallery")
            used_seed = result.get("seed", seed)

        if not gen_images:
            raise gr.Error("No images returned from remote Space.")

        # Normalize to list of file paths
        out_paths = []
        if isinstance(gen_images, list):
            for img in gen_images:
                if isinstance(img, dict) and "path" in img:
                    out_paths.append(img["path"])
                elif isinstance(img, dict) and "url" in img:
                    out_paths.append(img["url"])
                elif isinstance(img, (list, tuple)) and len(img) > 0:
                    if isinstance(img[0], dict) and "path" in img[0]:
                        out_paths.append(img[0]["path"])
                    elif isinstance(img[0], str):
                        out_paths.append(img[0])
                elif isinstance(img, str):
                    out_paths.append(img)
        elif isinstance(gen_images, str):
            out_paths.append(gen_images)

        if not out_paths:
            raise gr.Error("Could not parse generated image paths.")

        # If chaining, return output as new input gallery format
        if chain_in:
            chain_gallery = []
            for p in out_paths:
                chain_gallery.append({"path": p, "url": p})
            return out_paths, str(used_seed), chain_gallery
        else:
            return out_paths, str(used_seed), None

    except gr.Error:
        raise
    except Exception as e:
        raise gr.Error(f"Error: {e}")

# Template helpers
TPL = {
    "preserve": "Keep the subject's facial features, hair, skin tone, and costume details identical to the source image. ",
    "lighting": "Match the original scene's warm stage lighting and color grading exactly. ",
    "realism": "Professional digital photography. ",
    "group": "Keep all subjects' identities identical. ",
    "pose": "Same character, same identity. She is now in a new position. ",
}

def add_template(tpl_key, current_prompt):
    return current_prompt + TPL.get(tpl_key, "")

with gr.Blocks() as demo:
    gr.HTML("""
    <div style='text-align:center; padding:8px 0 4px 0'>
        <h1 style='color:#e91e63; font-size:1.4em; margin:0'>🔥 Qwen Edit Studio — NSFW v23</h1>
        <p style='color:#888; font-size:0.85em'>Client for Qwen-Image-Edit-Rapid-AIO NSFW v2.3</p>
    </div>
    """)
    with gr.Accordion("❓ How to Use", open=False):
        gr.HTML("""
        <div style='font-size:0.85em; color:#999; line-height:1.6'>
        <b>1.</b> Upload your source image(s) in the Image Gallery below.<br>
        <b>2.</b> Write your edit prompt in natural language (e.g. "Keep the subject's face identical. Change outfit to red dress").<br>
        <b>3.</b> Adjust parameters if needed — defaults are optimized for NSFW v23.<br>
        <b>4.</b> Click Generate. Results appear in the Output Gallery.<br>
        <b>5.</b> Click "Use Output as Input" to chain edits iteratively.<br>
        <b>Template buttons</b> insert preset prompt fragments.<br>
        <b>Tips:</b> Use natural language (not tags). Always specify what to keep. Match output size to input size.
        </div>
        """)

    with gr.Row():
        with gr.Column():
            gr.Markdown("### 📥 Input")
            input_gallery = gr.Gallery(
                label="Source Images",
                columns=2,
                height=300,
                type="filepath",
                interactive=True,
            )
            prompt_box = gr.Textbox(
                label="Edit Prompt",
                placeholder="Describe the edit in natural language. E.g. 'Keep the subject's face and identity identical. She is now on her knees performing oral sex. Realistic photography, warm lighting.'",
                lines=4,
            )

            with gr.Row():
                btn_preserve = gr.Button("Preserve ID", size="sm")
                btn_lighting = gr.Button("Match Light", size="sm")
                btn_realism = gr.Button("+Realism", size="sm")
                btn_group = gr.Button("Group Tpl", size="sm")
                btn_pose = gr.Button("Pose Tpl", size="sm")

            btn_generate = gr.Button("⚡ Generate", variant="primary", size="lg")

        with gr.Column():
            gr.Markdown("### 🌅 Output")
            output_gallery = gr.Gallery(
                label="Generated Images",
                columns=2,
                height=300,
                type="filepath",
                interactive=False,
            )
            seed_output = gr.Textbox(label="Seed Used", interactive=False)
            btn_chain = gr.Button("🔄 Use Output as Input", size="sm")

    with gr.Accordion("⚙️ Advanced Parameters", open=False):
        with gr.Row():
            steps = gr.Slider(4, 16, value=6, step=1, label="Inference Steps")
            guidance = gr.Slider(0.5, 3.0, value=1.0, step=0.1, label="CFG (Guidance Scale)")
        with gr.Row():
            seed = gr.Number(value=42, label="Seed (0=random)")
            randomize = gr.Checkbox(value=True, label="Randomize Seed")
        with gr.Row():
            rewrite = gr.Checkbox(value=True, label="Rewrite Prompt (LMS)")
            chain_in = gr.Checkbox(value=False, label="Chain Output→Input")
        with gr.Row():
            width = gr.Number(value=256, label="Width (256=auto)")
            height = gr.Number(value=256, label="Height (256=auto)")

    # Event handlers
    btn_preserve.click(add_template, [gr.Textbox(visible=False), prompt_box], prompt_box) if False else None

    # Fix: template buttons append to prompt
    for btn, key in [
        (btn_preserve, "preserve"),
        (btn_lighting, "lighting"),
        (btn_realism, "realism"),
        (btn_group, "group"),
        (btn_pose, "pose"),
    ]:
        btn.click(lambda p, k=key: p + TPL[k], [prompt_box], [prompt_box])

    btn_generate.click(
        generate,
        inputs=[input_gallery, prompt_box, steps, guidance, seed, randomize, rewrite, width, height, chain_in],
        outputs=[output_gallery, seed_output, input_gallery],
    )

    btn_chain.click(
        lambda out: out,
        inputs=[output_gallery],
        outputs=[input_gallery],
    )

demo.launch()