import gradio as gr
from modules import shared

from sd_dynamic_prompts.helpers import load_magicprompt_models


def on_ui_settings():
    section = "dynamicprompts", "Dynamic Prompts"
    shared.opts.add_option(
        key="dp_ignore_whitespace",
        info=shared.OptionInfo(
            False,
            label="Ignore whitespace in prompts: All newlines, tabs, and multiple spaces are replaced by a single space",
            section=section,
        ),
    )
    shared.opts.add_option(
        key="dp_write_raw_template",
        info=shared.OptionInfo(
            False,
            label="Save template to metadata: Write prompt template into the PNG metadata",
            section=section,
        ),
    )
    shared.opts.add_option(
        key="dp_write_prompts_to_file",
        info=shared.OptionInfo(
            False,
            label="Write prompts to file: Create a new .txt file for every batch containing the prompt template as well as the generated prompts.",
            section=section,
        ),
    )
    shared.opts.add_option(
        key="dp_parser_variant_start",
        info=shared.OptionInfo(
            "{",
            label="String to use as left bracket for parser variants, .e.g {variant1|variant2|variant3}",
            section=section,
        ),
    )
    shared.opts.add_option(
        key="dp_parser_variant_end",
        info=shared.OptionInfo(
            "}",
            label="String to use as right bracket for parser variants, .e.g {variant1|variant2|variant3}",
            section=section,
        ),
    )
    shared.opts.add_option(
        key="dp_parser_wildcard_wrap",
        info=shared.OptionInfo(
            "__",
            label="String to use as wrap for parser wildcard, .e.g __wildcard__",
            section=section,
        ),
    )
    shared.opts.add_option(
        key="dp_limit_jinja_prompts",
        info=shared.OptionInfo(
            False,
            label="Limit Jinja prompts: Limit the number of prompts to batch_count * batch_size. The default is to generate batch_count * batch_size * number of prompts generated by Jinja",
            section=section,
        ),
    )

    shared.opts.add_option(
        key="dp_auto_purge_cache",
        info=shared.OptionInfo(
            False,
            label="Automatically purge wildcard cache on every generation.",
            section=section,
        ),
    )

    shared.opts.add_option(
        key="dp_wildcard_manager_no_dedupe",
        info=shared.OptionInfo(
            False,
            label="Disable de-duplication of wildcards before processing.",
            section=section,
        ),
    )

    shared.opts.add_option(
        key="dp_wildcard_manager_no_sort",
        info=shared.OptionInfo(
            False,
            label="Disable sorting of wildcards before processing.",
            section=section,
        ),
    )

    shared.opts.add_option(
        key="dp_wildcard_manager_shuffle",
        info=shared.OptionInfo(
            False,
            label="Shuffle wildcards before use for more random outputs",
            section=section,
        ),
    )

    magic_models = load_magicprompt_models()
    shared.opts.add_option(
        key="dp_magicprompt_default_model",
        info=shared.OptionInfo(
            magic_models[0] if magic_models else "",
            label="Default magic prompt model",
            component=gr.Dropdown,
            component_args={"choices": magic_models},
            section=section,
        ),
    )

    shared.opts.add_option(
        key="dp_magicprompt_batch_size",
        info=shared.OptionInfo(
            1,
            label="Magic Prompt batch size (higher is faster but uses more memory)",
            component=gr.Slider,
            component_args={"minimum": 1, "maximum": 64, "step": 1},
            section=section,
        ),
    )