##!/usr/bin/python3
# -*- coding: utf-8 -*-
# @Time    : 2025-03-12
# @Author  : Junjie He
import os
import time
import uuid

import gradio as gr
import numpy as np
from PIL import Image

from src.anystory import call_anystory
from src.matting import ImageUniversalMatting
from src.util import upload_pil_2_oss

if not os.path.exists("models/tf_matting.pb"):
    os.makedirs("models", exist_ok=True)
    os.system(f"wget -O  models/tf_matting.pb {os.getenv('MATTING_PATH')}")
universal_matting = ImageUniversalMatting("models/tf_matting.pb")


def image_matting(pil_image):
    if pil_image.mode == "RGBA":
        mask = np.array(pil_image)[..., -1] > 200
        if np.all(mask):
            mask = ((universal_matting(pil_image)[..., -1] > 200) * 255).astype(np.uint8)
        else:
            mask = ((np.array(pil_image)[..., -1] > 200) * 255).astype(np.uint8)
    else:
        mask = ((universal_matting(pil_image.convert("RGB"))[..., -1] > 200) * 255).astype(np.uint8)
    pil_mask = Image.fromarray(mask)

    np_image = np.array(pil_image.convert("RGB"))
    np_mask = np.array(pil_mask)[..., None] / 255.
    pil_masked_image = Image.fromarray((np_mask * np_image + (1 - np_mask) * 255.).astype(np.uint8))
    return pil_masked_image, pil_mask


def process(
        pil_subject_A_image=None,
        pil_subject_A_mask=None,
        pil_subject_B_image=None,
        pil_subject_B_mask=None,
        prompt="",
):
    request_id = time.strftime('%Y%m%d-', time.localtime(time.time())) + str(uuid.uuid4())

    if prompt == "":
        raise gr.Error("Please enter your prompt")

    if pil_subject_A_image is None and pil_subject_B_image is None:
        raise gr.Error("Please upload your reference image(s)")

    image_urls = []
    if pil_subject_A_image is not None:
        if pil_subject_A_mask is not None:
            if pil_subject_A_mask.size != pil_subject_A_image.size:
                raise gr.Error("Subject [A] image & mask size mismatch")
            pil_subject_A_image = pil_subject_A_image.convert("RGB")
            pil_subject_A_mask = pil_subject_A_mask.convert("L")
            pil_subject_A_image = Image.merge("RGBA", (*pil_subject_A_image.split(), pil_subject_A_mask))
        image_urls.append(upload_pil_2_oss(pil_subject_A_image, name=request_id + "_A.png"))
    if pil_subject_B_image is not None:
        if pil_subject_B_mask is not None:
            if pil_subject_B_mask.size != pil_subject_B_image.size:
                raise gr.Error("Subject [B] image & mask size mismatch")
            pil_subject_B_image = pil_subject_B_image.convert("RGB")
            pil_subject_B_mask = pil_subject_B_mask.convert("L")
            pil_subject_B_image = Image.merge("RGBA", (*pil_subject_B_image.split(), pil_subject_B_mask))
        image_urls.append(upload_pil_2_oss(pil_subject_B_image, name=request_id + "_B.png"))

    res = call_anystory(image_urls, prompt)[0]

    return res


def interface():
    with gr.Row(variant="panel"):
        gr.HTML(description + "<br>" + tips)
    with gr.Row(variant="panel"):
        with gr.Column(scale=2, min_width=100):
            with gr.Row(equal_height=False):
                with gr.Column(scale=1, min_width=100):
                    with gr.Tab(label="Subject [A]"):
                        with gr.Group():
                            with gr.Row(equal_height=True):
                                with gr.Column(min_width=100):
                                    pil_subject_A_image = gr.Image(type="pil", label="Subject [A] Reference Image",
                                                                   format="png", show_label=True, image_mode="RGBA")
                                with gr.Column(min_width=100):
                                    with gr.Group():
                                        pil_subject_A_mask = gr.Image(type="pil",
                                                                      label="Subject [A] Mask (upload supported)",
                                                                      format="png", show_label=True, image_mode="L")
                                        seg_subject_A = gr.Button(value="Segment Subject")

                with gr.Column(scale=1, min_width=100):
                    with gr.Tab(label="Subject [B]"):
                        with gr.Group():
                            with gr.Row(equal_height=True):
                                with gr.Column(min_width=100):
                                    pil_subject_B_image = gr.Image(type="pil", label="Subject [B] Reference Image",
                                                                   format="png", show_label=True, image_mode="RGBA")
                                with gr.Column(min_width=100):
                                    with gr.Group():
                                        pil_subject_B_mask = gr.Image(type="pil",
                                                                      label="Subject [B] Mask (upload supported)",
                                                                      format="png", show_label=True, image_mode="L")
                                        seg_subject_B = gr.Button(value="Segment Subject")

            with gr.Group():
                prompt = gr.Textbox(value="", label='Prompt', lines=6, show_label=True)

        with gr.Column(scale=1, min_width=100):
            result_gallery = gr.Image(type="pil", label="Generated Image", visible=True, height=450)
            # result_gallery = gr.Gallery(label='Generated Image', show_label=True, elem_id="gallery", preview=True,
            #                             format="png", height=450)
            run_button = gr.Button(value="🧑‍🎨 RUN")

            generated_information = gr.Markdown(label="Generation Details", value="", visible=False)

    seg_subject_A.click(
        fn=set_image_seg_unfinished, outputs=generated_information
    ).then(
        fn=image_matting, inputs=[pil_subject_A_image], outputs=[pil_subject_A_image, pil_subject_A_mask]
    ).then(
        fn=set_image_seg_finished, outputs=generated_information
    )

    seg_subject_B.click(
        fn=set_image_seg_unfinished, outputs=generated_information
    ).then(
        fn=image_matting, inputs=[pil_subject_B_image], outputs=[pil_subject_B_image, pil_subject_B_mask]
    ).then(
        fn=set_image_seg_finished, outputs=generated_information
    )

    run_button.click(
        fn=set_image_generate_unfinished, outputs=generated_information
    ).then(
        fn=process,
        inputs=[pil_subject_A_image, pil_subject_A_mask, pil_subject_B_image, pil_subject_B_mask, prompt],
        outputs=[result_gallery]
    ).then(
        fn=set_image_generate_finished, outputs=generated_information
    )

    with gr.Row():
        examples = [
            [
                "assets/examples/1.webp",
                "assets/examples/1_mask.webp",
                None,
                None,
                "Cartoon style. A sheep is riding a skateboard and gliding through the city, holding a wooden sign that says \"TongYi\".",
                "assets/examples/1_output.webp",
            ],
            [
                "assets/examples/2.webp",
                "assets/examples/2_mask.webp",
                None,
                None,
                "Cartoon style. Sun Wukong stands on a tank, holding up an ancient wooden sign high in the air. The sign reads 'AnyStory'. The background is a cyberpunk-style city sky filled with towering buildings.",
                "assets/examples/2_output.webp",
            ],
            [
                "assets/examples/3.webp",
                "assets/examples/3_mask.webp",
                None,
                None,
                "A modern and stylish Nezha playing an electric guitar, dynamic pose, vibrant colors, fantasy atmosphere, mythical Chinese character with a rock-and-roll twist, red scarf flowing in the wind, traditional elements mixed with contemporary design, cinematic lighting, 4k resolution",
                "assets/examples/3_output.webp",
            ],
            [
                "assets/examples/4.webp",
                "assets/examples/4_mask.webp",
                None,
                None,
                "a man riding a bike on the road",
                "assets/examples/4_output.webp",
            ],
            [
                "assets/examples/7.webp",
                "assets/examples/7_mask.webp",
                None,
                None,
                "Nezha is surrounded by a mysterious purple glow, with a pair of eyes glowing with an eerie red light. Broken talismans and debris float around him, highlighting his demonic nature and authority.",
                "assets/examples/7_output.webp",
            ],
            [
                "assets/examples/8.webp",
                "assets/examples/8_mask.webp",
                None,
                None,
                "The car is driving through a cyberpunk city at night in the middle of a heavy downpour.",
                "assets/examples/8_output.webp",
            ],
            [
                "assets/examples/9.webp",
                "assets/examples/9_mask.webp",
                None,
                None,
                "This cosmetic is placed on a table covered with roses.",
                "assets/examples/9_output.webp",
            ],
            [
                "assets/examples/10.webp",
                "assets/examples/10_mask.webp",
                None,
                None,
                "A little boy model is posing for a photo.",
                "assets/examples/10_output.webp",
            ],
            # [
            #     "assets/examples/5_1.webp",
            #     "assets/examples/5_1_mask.webp",
            #     "assets/examples/5_2.webp",
            #     "assets/examples/5_2_mask.webp",
            #     "两个小孩骑着一辆炫酷的双人电动车，在热闹的菜市场中穿梭。周围是琳琅满目的蔬菜摊、水果筐和忙碌的摊主，他们表情专注又带点嬉笑，车篮里还装着几根胡萝卜和一把青菜，传统与现代元素在烟火气息中完美交融。",
            #     "assets/examples/5_output.webp",
            # ],
            [
                "assets/examples/6_1.webp",
                "assets/examples/6_1_mask.webp",
                "assets/examples/6_2.webp",
                "assets/examples/6_2_mask.webp",
                "Two men are sitting by a wooden table, which is laden with delicious food and a pot of wine. One of the men holds a wine glass, drinking heartily with a bold expression; the other smiles as he pours wine for his companion, both of them engaged in cheerful conversation. In the background is an ancient pavilion surrounded by emerald bamboo groves, with sunlight filtering through the leaves to cast dappled shadows.",
                "assets/examples/6_output.webp",
            ],
        ]
        gr.Examples(
            label="Examples",
            examples=examples,
            inputs=[pil_subject_A_image, pil_subject_A_mask, pil_subject_B_image, pil_subject_B_mask, prompt,
                    result_gallery],
        )


def set_image_seg_unfinished():
    return gr.update(
        visible=True,
        value="<h3>(Unfinished) Extracting Subject Mask...</h3>",
    )


def set_image_seg_finished():
    return gr.update(visible=True, value="<h3>Subject mask ready!</h3>")


def set_image_generate_unfinished():
    return gr.update(
        visible=True,
        value="<h3>(Unfinished) Generating images...</h3>",
    )


def set_image_generate_finished():
    return gr.update(visible=True, value="<h3>Image generation is completed!</h3>")


if __name__ == "__main__":
    title = r"""
            <div style="text-align: center;">
                <h1> AnyStory: Towards Unified Single and Multiple Subject Personalization in Text-to-Image Generation </h1>
                <h1> V2.0.0 </h1>
                <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
                    <a href="https://arxiv.org/pdf/2501.09503"><img src="https://img.shields.io/badge/arXiv-2501.09503-red"></a>
                    &nbsp;
                    <a href='https://aigcdesigngroup.github.io/AnyStory/'><img src='https://img.shields.io/badge/Project_Page-AnyStory-green' alt='Project Page'></a>
                    &nbsp;
                    <a href='https://modelscope.cn/studios/iic/AnyStory'><img src='https://img.shields.io/badge/Demo-ModelScope-blue'></a>
                </div>
                </br>
            </div>
        """

    title_description = r"""
        Official demo of <b>AnyStory 2</b> 🤗. We will continuously update this demo.
        For technical details, please refer to our tech report: <a href='https://arxiv.org/pdf/2501.09503' target='_blank'><b>AnyStory: Towards Unified Single and Multiple Subject Personalization in Text-to-Image Generation</b></a>. 😊
        """

    description = r"""🚀🚀🚀 Quick Start:<br>
        1. Upload subject reference images (clean background; real human IDs unsupported for now), Add prompts (only EN supported), and Click "<b>RUN</b>".<br>
        2. (Recommended) Click "<b>Segment Subject</b>" to create masks (or upload your own B&W masks) for subjects. This helps the model better reference the subject you specify (otherwise, we will perform automatic detection). 🤗<br>
        """

    tips = r"""💡💡💡 Tips:<br>
    If the subject doesn't appear, try adding a detailed description of the subject in the prompt that matches the reference image, and avoid conflicting details (e.g., significantly altering the subject's appearance). Multi-subject referencing in AnyStory2 is still being optimized. 🤗<br>
    """

    citation = r"""
        ---
        📝 **Citation**
        <br>
        If our work is helpful for your research or applications, please cite us via:
        ```bibtex
        @article{he2025anystory,
          title={AnyStory: Towards Unified Single and Multiple Subject Personalization in Text-to-Image Generation},
          author={He, Junjie and Tuo, Yuxiang and Chen, Binghui and Zhong, Chongyang and Geng, Yifeng and Bo, Liefeng},
          journal={arXiv preprint arXiv:2501.09503},
          year={2025}
        }
        ```
        If you have any questions, feel free to open an issue or contact us directly at <b>hejunjie1103@gmail.com</b>.
        """

    js = """
    function createGradioAnimation() {
        var container = document.createElement('div');
        container.id = 'gradio-animation';
        container.style.fontSize = '2em';
        container.style.fontWeight = 'bold';
        container.style.textAlign = 'center';
        container.style.marginBottom = '20px';

        var text = 'Welcome to AnyStory!';
        for (var i = 0; i < text.length; i++) {
            (function(i){
                setTimeout(function(){
                    var letter = document.createElement('span');
                    letter.style.opacity = '0';
                    letter.style.transition = 'opacity 0.5s';
                    letter.innerText = text[i];

                    container.appendChild(letter);

                    setTimeout(function() {
                        letter.style.opacity = '1';
                    }, 50);
                }, i * 250);
            })(i);
        }

        var gradioContainer = document.querySelector('.gradio-container');
        gradioContainer.insertBefore(container, gradioContainer.firstChild);

        return 'Animation created';
    }
    """
    block = gr.Blocks(title="AnyStory2", js=js, theme=gr.themes.Ocean()).queue()
    with block:
        gr.HTML(title)
        gr.HTML(title_description)

        interface()

        gr.HTML("<br>More examples: Intelligent creation of AI story pictures integrated with Qwen Agent")
        gr.Gallery(value=["assets/storyboard_en.png"], columns=1, object_fit="contain", show_label=False)

        gr.Markdown(citation)

    block.launch(share=True, max_threads=10)
    # block.launch(server_name='0.0.0.0', share=False, server_port=9999, max_threads=10)