import gradio as gr
from utils import calculate_cost, augment_image

def update_btn_cost(w, h, req_type, selection, l=''):
    if selection == 'scale':
        after = l.split()[-1]
        w, h = map(int, after.split('x'))
    cost = calculate_cost(w, h, rmbg=(req_type == "移除背景"))
    return gr.Button(value=f"生成(预计消耗{cost}点数)")

def getWH(img, w, h, f):
    if img is None:
        return w, h, ''
    w, h = img.width, img.height
    l = f"{w}x{h} -> {int(w * f)}x{int(h * f)}"
    return w, h, l

def refresh_hint(l, f):
    if len(l) == 0:
        return ''
    before = l.split()[0]
    w, h = map(int, before.split('x'))
    l = f"{w}x{h} -> {int(w * f)}x{int(h * f)}"
    return l

def set_index(evt: gr.SelectData):
    return evt.index

def send_outputs(imgs, index):
    if imgs is None or len(imgs) == 0:
        return None
    elif len(imgs) == 1:
        return imgs[0][0]
    elif index >= 0:
        return imgs[index][0]
    else:
        return None
    
def director_ui():
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="上传图片", value=None, sources=["upload", "clipboard", "webcam"], interactive=True, type="pil", show_share_button=False)
            with gr.Row():
                from_t2i = gr.Button("使用文生图上一次生成的图片")
                from_out = gr.Button("使用定向修图上一次生成的图片")
        with gr.Column():
            output_image = gr.Gallery(label="输出图片", format='png', interactive=False, preview=True, show_share_button=False)
            selected_index = gr.Number(value=-1, visible=False)
            with gr.Row():
                send_i2i = gr.Button("发送到图生图")
                send_inp = gr.Button("发送到局部重绘")
                send_vib = gr.Button("发送到风格迁移")
    req_type = gr.Radio(["移除背景", "素描", "线稿", "上色", "更改表情", "去聊天框"], value="线稿", label="功能选择")
    with gr.Tab("自由调整") as resize:
        width = gr.Slider(label="调整宽度", value=1024, minimum=1, maximum=2048, step=1)
        height = gr.Slider(label="调整高度", value=1024, minimum=1, maximum=2048, step=1)
        reset = gr.Button("重置")
    with gr.Tab("等比例缩放") as scale:
        hint = gr.Textbox('', show_label=False, interactive=False, max_lines=1, container=False)
        factor = gr.Slider(label="缩放比例", minimum=0.1, maximum=2, step=0.1)
    defry = gr.State(0)
    prompt = gr.State('')
    selection = gr.Radio(["resize", "scale"], visible=False)
    resize.select(lambda: 'resize', inputs=None, outputs=selection)
    scale.select(lambda: 'scale', inputs=None, outputs=selection)
    @gr.render(inputs=req_type)
    def extra_paras(r):
        if r == "上色":
            with gr.Row():
                strength = gr.Radio(choices=range(6), value=0, label="强度")
                text = gr.Textbox(label="提示词(可选)", value='')
                strength.change(lambda x: x, inputs=strength, outputs=defry)
                text.change(lambda x: x, inputs=text, outputs=prompt)
        elif r == "更改表情":
            with gr.Row():
                emotion = gr.Dropdown(choices=["neutral", "happy", "sad", "angry", "scared", "surprised", "tired", "excited", "nervous", "thinking", "confused", "shy", "disgusted", "smug", "bored", "laughing", "irritated", "aroused", "embarrassed", "worried", "love", "determined", "hurt", "playful"], value="neutral", label="表情")
                strength = gr.Dropdown(choices=["正常", "偏弱", "弱", "更弱", "很弱", "最弱"], value="正常", label="强度")
                text = gr.Textbox(label="提示词(可选)", value='')
                emotion.change(lambda e, t: f"{e};;{t}", inputs=[emotion, text], outputs=prompt)
                text.change(lambda e, t: f"{e};;{t}", inputs=[emotion, text], outputs=prompt)
                strength.change(lambda s: {"正常": 0, "偏弱": 1, "弱": 2, "更弱": 3, "很弱": 4, "最弱": 5}[s], inputs=strength, outputs=defry)
        else:
            return
    output_image.select(set_index, inputs=None, outputs=selected_index)
    from_out.click(send_outputs, inputs=[output_image, selected_index], outputs=input_image)
    req_type.change(lambda r: "neutral;;" if r == "更改表情" else '', inputs=req_type, outputs=prompt)
    input_image.change(getWH, inputs=[input_image, width, height, factor], outputs=[width, height, hint])
    factor.change(refresh_hint, inputs=[hint, factor], outputs=[hint])
    tool_btn = gr.Button(value="生成", variant="primary", visible=False)
    dtool_stop = gr.Button(value="取消", variant="stop", visible=False)
    input_image.change(lambda i: gr.Button(visible=False) if i is None else gr.Button(visible=True), inputs=input_image, outputs=tool_btn)
    tool_btn.click(lambda: (gr.Button(visible=False), gr.Button(visible=True), gr.Gallery(selected_index=None), -1), inputs=None, outputs=[tool_btn, dtool_stop, output_image, selected_index])
    dtool_gen = tool_btn.click(augment_image, inputs=[input_image, width, height, req_type, selection, factor, defry, prompt], outputs=output_image, concurrency_id="generate")
    dtool_stop.click(lambda: (gr.Button(visible=True), gr.Button(visible=False)), inputs=None, outputs=[tool_btn, dtool_stop], cancels=[dtool_gen])
    output_image.change(lambda: (gr.Button(visible=True), gr.Button(visible=False)), inputs=None, outputs=[tool_btn, dtool_stop])
    reset.click(getWH, inputs=[input_image, width, height, factor], outputs=[width, height, hint])
    width.change(update_btn_cost, inputs=[width, height, req_type, selection, hint], outputs=tool_btn)
    height.change(update_btn_cost, inputs=[width, height, req_type, selection, hint], outputs=tool_btn)
    req_type.change(update_btn_cost, inputs=[width, height, req_type, selection, hint], outputs=tool_btn)
    hint.change(update_btn_cost, inputs=[width, height, req_type, selection, hint], outputs=tool_btn)
    selection.change(update_btn_cost, inputs=[width, height, req_type, selection, hint], outputs=tool_btn)
    return from_t2i, send_i2i, send_inp, send_vib, input_image, output_image, selected_index