import gradio as gr
import json
from random import sample


with open("meta.json") as f:
    meta = json.load(f)
with open("tag2cat.json") as f:
    tag2cat = json.load(f)
f = lambda cat: [
    ", ".join(
        t
        for t in (tag.replace("_", " ") for tag in e["tag_string_general"].split())
        if tag2cat.get(t, "") == cat
    )
    for e in meta
]
l = [f(c) for c in "Character Outfit Action Scene".split()]


def do():
    idxs = sample(range(len(meta)), 4)
    parts = [ll[idx] for ll, idx in zip(l, idxs)]
    return (",\n".join(parts),) + tuple(
        _
        for p, idx in zip(parts, idxs)
        for _ in (
            f"{p} ([danbooru post](https://danbooru.donmai.us/posts/{meta[idx]['id']}))",
            meta[idx]["large_file_url"],
        )
    )


with gr.Blocks() as demo:
    with gr.Row():
        button = gr.Button("Get Random Prompt", variant="primary")
    with gr.Row():
        joined = gr.Text(label="generated prompt", show_copy_button=True)
    with gr.Row(equal_height=True):
        text_char = gr.Markdown(label="character")
        text_outfit = gr.Markdown(label="outfit")
        text_action = gr.Markdown(label="action")
        text_scene = gr.Markdown(label="scene")
    with gr.Row(equal_height=True):
        img_char = gr.Image(label="character ref")
        img_outfit = gr.Image(label="outfit ref")
        img_action = gr.Image(label="action ref")
        img_scene = gr.Image(label="scene ref")
    button.click(
        do,
        [],
        [
            joined,
            text_char,
            img_char,
            text_outfit,
            img_outfit,
            text_action,
            img_action,
            text_scene,
            img_scene,
        ],
    )

demo.launch()