import gradio as gr
import requests

prefix = "https://smartfeed-custom-tools.hf.space/gradio_api/file="


def change_image_style(image_url, style_image_url):
    data = {
        "image_url": image_url,
        "style_image_url": style_image_url
    }

    response = requests.post(
        "https://api.hkhappymobile.com/tools/change-image-style",
        json=data,
        headers={"content-type": "application/json"}
    )

    if response.status_code == 200:
        resultImageUrl = response.json().get("data").get("image_url")
        return resultImageUrl
    else:
        raise Exception(f"Error: {response.status_code} - {response.text}")

def generate_image(input_image, style_images):

    if not input_image:
        raise gr.Error(f"Please upload an input image! Refer to step 1️⃣")

    if style_images is None:
        raise gr.Error(f"Cannot find any style image! Please refer to step 1️⃣")

    inputImageUrl = prefix + input_image

    print(f"Input Image URL: {inputImageUrl}")

    result_images = list[str]()
    for style_image in style_images:
        if not style_image:
            raise gr.Error(f"Cannot find any style image! Please refer to step 1️⃣")
        styleImageUrl = prefix + style_image[0]
        print(f"Style Image URL: {styleImageUrl}")
        result_images.append(change_image_style(inputImageUrl, styleImageUrl))

    return result_images
    

def swap_to_gallery(images):
    return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False)

def remove_back_to_files():
    return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Input Image", type="filepath", interactive=True)

            files = gr.File(
                        label="Drag (Select) 1 or more style images",
                        file_types=["image"],
                        file_count="multiple"
                    )
            uploaded_files = gr.Gallery(label="Your images", visible=False, columns=5, rows=1, height=200)
            with gr.Column(visible=False) as clear_button:
                remove_and_reupload = gr.ClearButton(value="Remove and upload new ones", components=files, size="sm")
            submit = gr.Button("Submit")

        with gr.Column():
            gallery = gr.Gallery(label="Generated Images")

        files.upload(fn=swap_to_gallery, inputs=files, outputs=[uploaded_files, clear_button, files])
        remove_and_reupload.click(fn=remove_back_to_files, outputs=[uploaded_files, clear_button, files])

        submit.click(
            fn=generate_image,
            inputs=[input_image, uploaded_files],
            outputs=[gallery]
        )

    
demo.launch()