import gradio as gr
from functions import CAT_ARTIST, CAT_CHARACTER, CAT_COPYRIGHT, CAT_GENERAL, CAT_LORE, CAT_META, CAT_SPECIES, PromptBuilder, parse_tag, parse_tags, related_tags

def query_tag(tag: str, category: int):
    if category == -1:
        category = None 
    return related_tags(parse_tag(tag), category=category)

def generate_prompt(include: str, focus: str, exclude: str, avoid: str, skip: str, rating: str, general: int, artist: int, species: int, copyright: int, character: int, meta: int) -> str:
    try:
        builder = PromptBuilder(skip=list(parse_tags(skip)), min_posts=50, rating=rating)
        for tag in parse_tags(include):
            builder = builder.include(tag)
        for tag in parse_tags(focus):
            builder = builder.focus(tag)
        for tag in parse_tags(exclude):
            builder = builder.exclude(tag)
        for tag in parse_tags(avoid):
            builder = builder.avoid(tag)
        if artist > 0:
            builder = builder.pick(CAT_ARTIST, artist, 10)
        if species > 0:
            builder = builder.pick(CAT_SPECIES, species, 10)
        if copyright > 0:
            builder = builder.pick(CAT_COPYRIGHT, copyright, 10)
        if character > 0:
            builder = builder.pick(CAT_CHARACTER, character, 10)
        if meta > 0:
            builder = builder.pick(CAT_META, meta, 10)
        if general > 0:
            builder = builder.pick(CAT_GENERAL, general, 50)
        return builder.get_one()
    except Exception as e:
        return str(e)

with gr.Blocks() as demo:
    with gr.Tab("Tag Explorer"):
        tag = gr.Textbox(label="Tag")
        category = gr.Dropdown(label="Category", choices=[("All", -1), ("General", CAT_GENERAL), ("Artist", CAT_ARTIST), ("Copyright", CAT_COPYRIGHT), ("Character", CAT_CHARACTER), ("Species", CAT_SPECIES), ("Meta", CAT_META), ("Lore", CAT_LORE)], value=-1)
        query = gr.Button("Query")
        output = gr.Dataframe()
        query.click(fn=query_tag, inputs=[tag, category], outputs=output)
    with gr.Tab("Prompt Expander"):
        include = gr.Textbox(label="Positive Prompt - Start with these tags and weight picks toward them")
        focus = gr.Textbox(label="Focus - Used for picks but not necessarily added to the prompt")
        exclude = gr.Textbox(label="Negative Prompt - Put these in the negative prompt and weight picks against them")
        avoid = gr.Textbox(label="Other Negatives - Weighted against picks but not put in the negatives")
        skip = gr.Textbox(label="Skip - Never pick these tags")
        rating = gr.Dropdown(label="Rating Limit (not 100% reliable)", choices=[("Safe", "s"), ("Questionable", "q"), ("Explicit", "e")], value="s")
        with gr.Accordion(label="Tag Counts"):
            with gr.Row():
                general = gr.Number(5, label="General", precision=0, minimum=0, maximum=20)
                artist = gr.Number(0, label="Artist", precision=0, minimum=0, maximum=5)
                species = gr.Number(0, label="Species", precision=0, minimum=0, maximum=5)
                copyright = gr.Number(0, label="Copyright", precision=0, minimum=0, maximum=5)
                character = gr.Number(0, label="Character", precision=0, minimum=0, maximum=5)
                meta = gr.Number(0, label="Meta", precision=0, minimum=0, maximum=5)

        generate = gr.Button("Generate")
        output = gr.Textbox(label="Prompt")
        generate.click(fn=generate_prompt, inputs=[include, focus, exclude, avoid, skip, rating, general, artist, species, copyright, character, meta], outputs=output)

demo.launch()