import gradio as gr
from huggingface_hub import InferenceClient
import base64
import os

MODEL = "Qwen/Qwen-Image-Edit-2511"
PROVIDER = "fal-ai"
NUM_IMAGES = 5


def image_to_data_uri(image_path):
    with open(image_path, "rb") as f:
        data = base64.b64encode(f.read()).decode("utf-8")
    ext = os.path.splitext(image_path)[1].lower()
    mime = "image/png" if ext == ".png" else "image/jpeg"
    return f"data:{mime};base64,{data}"


def edit_images(img1, img2, img3, img4, img5, prompt, token: gr.OAuthToken | None = None):
    if not token:
        raise gr.Error("Please sign in with your Hugging Face account first.")
    if not prompt or prompt.strip() == "":
        raise gr.Error("Please enter a prompt.")

    images = [img for img in [img1, img2, img3, img4, img5] if img is not None]
    if len(images) == 0:
        raise gr.Error("Please upload at least one image.")

    client = InferenceClient(provider=PROVIDER, token=token.token)

    image_urls = [image_to_data_uri(img) for img in images]

    result = client.image_to_image(
        model=MODEL,
        image=image_urls[0],
        prompt=prompt,
        extra_body={"image_urls": image_urls} if len(image_urls) > 1 else {},
    )

    return result


def swap(a, b):
    return b, a


def send_to(output_img):
    return output_img


with gr.Blocks(fill_height=True) as demo:
    with gr.Sidebar():
        gr.Markdown("# Qwen Image Edit")
        gr.Markdown(
            f"Upload one or more images and describe how you want them edited.\n\n"
            f"Model: **{MODEL}** via **{PROVIDER}**"
        )
        login_btn = gr.LoginButton("Sign in with Hugging Face")

    with gr.Row():
        with gr.Column(scale=3):
            with gr.Row():
                img = [
                    gr.Image(
                        label=f"Image {i+1}" if i == 0 else f"Image {i+1} (optional)",
                        type="filepath",
                    )
                    for i in range(NUM_IMAGES)
                ]
            with gr.Row():
                swap_btns = []
                for i in range(NUM_IMAGES - 1):
                    swap_btns.append(gr.Button(f"{i+1} ↔ {i+2}", size="sm"))
            prompt = gr.Textbox(
                label="Edit Prompt",
                placeholder="Describe how you want the image(s) edited...",
                lines=3,
            )
            submit_btn = gr.Button("Edit", variant="primary")

        with gr.Column(scale=2):
            output_image = gr.Image(label="Result", type="filepath")
            gr.Markdown("**Send result to input:**")
            with gr.Row():
                to_btns = [
                    gr.Button(f"→ {i+1}", size="sm") for i in range(NUM_IMAGES)
                ]

    # Swap buttons
    for idx, btn in enumerate(swap_btns):
        btn.click(fn=swap, inputs=[img[idx], img[idx + 1]], outputs=[img[idx], img[idx + 1]])

    # Send output to input
    for idx, btn in enumerate(to_btns):
        btn.click(fn=send_to, inputs=[output_image], outputs=[img[idx]])

    # Edit
    submit_btn.click(
        fn=edit_images,
        inputs=[*img, prompt],
        outputs=[output_image],
    )

demo.launch()