import gradio as gr
import spaces  # ZERO GPU

# ref: https://qiita.com/tregu148/items/fccccbbc47d966dd2fc2
def gradio_copy_text(_text: None):
    gr.Info("Copied!")

COPY_ACTION_JS = """\
(inputs, _outputs) => {
  // inputs is the string value of the input_text
  if (inputs.trim() !== "") {
    navigator.clipboard.writeText(inputs);
  }
}"""

def _people_tag(noun: str, minimum: int = 1, maximum: int = 5):
    return (
        [f"1{noun}"]
        + [f"{num}{noun}s" for num in range(minimum + 1, maximum + 1)]
        + [f"{maximum+1}+{noun}s"]
    )


PEOPLE_TAGS = (
    _people_tag("girl") + _people_tag("boy") + _people_tag("other") + ["no humans"]
)
RATING_MAP = {
    "safe": "rating_safe",
    "sensitive": "rating_safe",
    "nsfw": "rating_questionable",
    "explicit, nsfw": "rating_explicit",
    "explicit": "rating_explicit",
    "rating:safe": "rating_safe",
    "rating:general": "rating_safe",
    "rating:sensitive": "rating_safe",
    "rating:questionable, nsfw": "rating_explicit",
    "rating:explicit, nsfw": "rating_explicit",
}

DESCRIPTION_MD = """
# Convert general Danbooru tags to Pony e621 tags
""".strip()

DESCRIPTION_MD2 = """
The dictionary was generated using the following repository: [ponapon280/danbooru-e621-converter](https://github.com/ponapon280/danbooru-e621-converter)
""".strip()


def character_list_to_series_list(character_list):
    def get_series_dict():
        import re

        with open('characterfull.txt', 'r') as f:
            lines = f.readlines()

        series_dict = {}
        for line in lines:
            parts = line.strip().split(', ')
            if len(parts) >= 3:
                name = parts[-2].replace("\\", "")
            if name.endswith(")"):
                names = name.split("(")
                character_name = "(".join(names[:-1])
                if character_name.endswith(" "):
                    name = character_name[:-1]
            series = re.sub(r'\\[()]', '', parts[-1])
            series_dict[name] = series

        return series_dict
    
    output_series_tag = []
    series_tag = ""
    series_dict = get_series_dict()
    for tag in character_list:
        series_tag = series_dict.get(tag, "")
        if tag.endswith(")"):
            tags = tag.split("(")
            character_tag = "(".join(tags[:-1])
            if character_tag.endswith(" "):
                character_tag = character_tag[:-1]
            series_tag = tags[-1].replace(")", "")

    if series_tag:
        output_series_tag.append(series_tag)

    return output_series_tag


def get_e621_dict():
    with open('danbooru_e621.csv', 'r', encoding="utf-8") as f:
        lines = f.readlines()

    e621_dict = {}
    for line in lines:
        parts = line.strip().split(',')
        e621_dict[parts[0]] = parts[1]

    return e621_dict


def danbooru_to_e621(dtag, e621_dict):
    def d_to_e(match, e621_dict):
        dtag = match.group(0)
        etag = e621_dict.get(dtag.strip().replace("_", " "), "")
        if etag:
            return etag
        else:
            return dtag
    
    import re
    tag = re.sub(r'[\w ]+', lambda wrapper: d_to_e(wrapper, e621_dict), dtag, 2)

    return tag


def nai_to_webui(ntag):
    def n_to_w(match):
        import re
        ntag = match.group(0)
        power = 1.0
        for i in range(ntag.count("{")):
            power *= 1.05
        for i in range(ntag.count("[")):
            power *= 0.952
        wtag_word = re.sub(r'(?:{+([\w ,_]+)}+)|(?:\[+([\w ,_]+)\]+)', r'\1\2', ntag, 2)
        wtag = f"({wtag_word}:{power:.2f})"
        if wtag:
            return wtag
        else:
            return ntag
    
    import re
    tag = re.sub(r'(?:{+[\w ,_]+}+)|(?:\[+[\w ,_]+\]+)', lambda wrapper: n_to_w(wrapper), ntag)

    return tag


def animagine_prompt(character: list[str], general: list[str], tag_type):
    people_tags: list[str] = []
    other_tags: list[str] = []
    rating_tags: list[str] = []

    e621_dict = get_e621_dict()
    for tag in general:
        tag = danbooru_to_e621(tag, e621_dict)
        if tag in PEOPLE_TAGS:        
            people_tags.append(tag)
        elif tag in RATING_MAP.keys():
            rating_tags.append(RATING_MAP.get(tag.replace(" ",""), ""))            
        else:
            other_tags.append(tag)

    rating_tags = sorted(set(rating_tags), key=rating_tags.index)
    rating_tags = [rating_tags[0]] if rating_tags else []

    output_series_tag = character_list_to_series_list(character)

    all_tag = ", ".join(people_tags + character + output_series_tag + other_tags + rating_tags)

    if tag_type == "NovelAI":
        all_tag = nai_to_webui(all_tag) 
    
    return all_tag


def convert_tags(
    input_copyright = "", input_character = "", input_general = "", tag_type = "WebUI",
):
    character = []
    general = []
    character.append(input_character) if input_character else []
    general = input_general.split(",") if input_general else []
    
    prompt = animagine_prompt(
        character, general, tag_type
    )
    
    return prompt, gr.update(interactive=True,)

    
def demo():
    with gr.Blocks() as ui:
        gr.Markdown(DESCRIPTION_MD)
        
        with gr.Row():
            with gr.Column():
                with gr.Group():
                    input_copyright = gr.Textbox(
                        label="Copyright tags",
                        placeholder="vocaloid",
                    )
                    input_character = gr.Textbox(
                        label="Character tags",
                        placeholder="hatsune miku",
                    )
                    input_general = gr.TextArea(
                        label="General tags",
                        lines=6,
                        placeholder="1girl, solo, ...",
                    )
                    tag_type = gr.Radio(
                        label="Style of tags",
                        choices=["WebUI", "NovelAI"],
                        value="WebUI",
                    )
                    start_btn = gr.Button(value="Convert", variant="primary")

            with gr.Column():
                with gr.Group():
                    prompt_text = gr.TextArea(label="Prompt", lines=6, interactive=False)
                    copy_btn = gr.Button(value="Copy to clipboard", interactive=False)

        gr.Markdown(DESCRIPTION_MD2)

        start_btn.click(
            convert_tags,
            inputs=[input_copyright, input_character, input_general, tag_type],
            outputs=[
                prompt_text,
                copy_btn,
            ],
        )
        copy_btn.click(gradio_copy_text, inputs=[prompt_text], js=COPY_ACTION_JS)
    
    return ui

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