import gradio as gr
import json
from random import sample


with open("meta.json") as f:
    meta = json.load(f)

meta_s = [x for x in meta if x["rating"] == "s"]
meta_g = [x for x in meta if x["rating"] == "g"]
with open("tag2cat.json") as f:
    tag2cat = json.load(f)
f = lambda cat, e: ", ".join(
    t
    for t in (tag.replace("_", " ") for tag in e["tag_string_general"].split())
    if tag2cat.get(t, "") == cat
)

cats = "Character Outfit Action Scene".split()


def do(ratings):
    if not ratings:
        raise gr.Error("Please select at least one of the ratings")
    m = (
        meta_g
        if ratings[0] == "General"
        else meta_s
        if ratings[0] == "Sensitive"
        else meta
    )
    idxs = sample(range(len(m)), 4)
    parts = [f(cat, m[idx]) for cat, idx in zip(cats, idxs)]
    return (",\n".join(parts),) + tuple(
        _
        for p, idx in zip(parts, idxs)
        for _ in (
            f"{p} ([danbooru post](https://danbooru.donmai.us/posts/{m[idx]['id']}))",
            f'![]({m[idx]["large_file_url"]})',
        )
    )


with gr.Blocks() as demo:
    with gr.Row():
        ratings = gr.CheckboxGroup(
            ["General", "Sensitive"],
            label="Allowed ratings",
            value=["General", "Sensitive"],
        )
        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()
        text_outfit = gr.Markdown()
        text_action = gr.Markdown()
        text_scene = gr.Markdown()
    with gr.Row(equal_height=True):
        img_char = gr.Markdown()
        img_outfit = gr.Markdown()
        img_action = gr.Markdown()
        img_scene = gr.Markdown()
    button.click(
        do,
        [ratings],
        [
            joined,
            text_char,
            img_char,
            text_outfit,
            img_outfit,
            text_action,
            img_action,
            text_scene,
            img_scene,
        ],
    )

demo.launch()