import os
import gradio as gr
from PIL import Image

from core import qwen_editor, wanx_editor

PRESETS = ["Smooth skin", "Slim face", "Remove wrinkles", "Brighten eyes"]


def make_preset_fn(p):
    def fn():
        return p
    return fn


def run_edit(image, instruction, model):
    if image is None:
        raise gr.Error("Please upload an image first.")
    if not instruction.strip():
        raise gr.Error("Please enter an instruction.")
    if not os.environ.get("DASHSCOPE_API_KEY"):
        raise gr.Error("DASHSCOPE_API_KEY is not configured.")

    pil_image = Image.open(image)
    editor = qwen_editor if model == "Qwen (qwen-image-2.0-pro)" else wanx_editor
    result = editor.edit_image(pil_image, instruction)
    return pil_image, result


with gr.Blocks(title="QwenImageEdit", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# QwenImageEdit\nInstruction-based image editing via DashScope API.")

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="filepath", label="Upload Photo")
            instruction_box = gr.Textbox(
                label="Instruction",
                placeholder='e.g. "smooth the skin and reduce dark circles"',
            )
            with gr.Row():
                for preset in PRESETS:
                    gr.Button(preset, size="sm").click(
                        fn=make_preset_fn(preset),
                        outputs=instruction_box,
                        api_name=False,
                    )
            model_radio = gr.Radio(
                choices=["Qwen (qwen-image-2.0-pro)", "Wanx (wanx2.1-imageedit)"],
                value="Qwen (qwen-image-2.0-pro)",
                label="Model",
            )
            submit_btn = gr.Button("Edit Image", variant="primary")

        with gr.Column():
            original_out = gr.Image(label="Original", interactive=False)
            result_out = gr.Image(label="Edited Result", interactive=False)

    submit_btn.click(
        fn=run_edit,
        inputs=[input_image, instruction_box, model_radio],
        outputs=[original_out, result_out],
        api_name=False,
    )

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