from modules import shared
from typing import List
from gradio import inputs as gr

# kaomoji from WD 1.4 tagger csv. thanks, Meow-San#5400!
DEFAULT_KAMOJIS = '0_0, (o)_(o), +_+, +_-, ._., <o>_<o>, <|>_<|>, =_=, >_<, 3_3, 6_9, >_o, @_@, ^_^, o_o, u_u, x_x, |_|, ||_||'

DEFAULT_OFF = '[name].[output_extension]'


def slider_wrapper(value, elem_id, **kwargs):
    # required or else gradio will throw errors
    return gr.Slider(**kwargs)


def on_ui_settings():
    Its = InterrogatorSettings
    section = 'tagger', 'Tagger'
    shared.opts.add_option(
        key='tagger_out_filename_fmt',
        info=shared.OptionInfo(
            DEFAULT_OFF,
            label='Tag file output format. Leave blank to use same filename or'
            ' e.g. "[name].[hash:sha1].[output_extension]". Also allowed are '
            '[extension] or any other [hash:<algorithm>] supported by hashlib',
            section=section,
        ),
    )
    shared.opts.onchange(
        key='tagger_out_filename_fmt',
        func=Its.set_output_filename_format
    )

    shared.opts.add_option(
        key='tagger_count_threshold',
        info=shared.OptionInfo(
            100.0,
            label="Maximum number of tags to be shown in the UI",
            section=section,
            component=slider_wrapper,
            component_args={"minimum": 1.0, "maximum": 500.0, "step": 1.0},
        ),
    )

    shared.opts.add_option(
        key='tagger_batch_recursive',
        info=shared.OptionInfo(
            True,
            label='Glob recursively with input directory pattern',
            section=section,
        ),
    )
    shared.opts.add_option(
        key='tagger_auto_serde_json',
        info=shared.OptionInfo(
            True,
            label='Auto load and save JSON database',
            section=section,
        ),
    )
    shared.opts.add_option(
        key='tagger_store_images',
        info=shared.OptionInfo(
            False,
            label='Store images in database',
            section=section,
        ),
    )
    shared.opts.add_option(
        key='tagger_verbose',
        info=shared.OptionInfo(
            False,
            label='Console log tag counts per file, no progress bar',
            section=section,
        ),
    )
    shared.opts.add_option(
        key='tagger_repl_us',
        info=shared.OptionInfo(
            True,
            label='Use spaces instead of underscore',
            section=section,
        ),
    )
    shared.opts.add_option(
        key='tagger_repl_us_excl',
        info=shared.OptionInfo(
            DEFAULT_KAMOJIS,
            label='Excudes (split by comma)',
            section=section,
        ),
    )
    shared.opts.onchange(
        key='tagger_repl_us_excl',
        func=Its.set_us_excl
    )
    shared.opts.add_option(
        key='tagger_escape',
        info=shared.OptionInfo(
            False,
            label='Escape brackets',
            section=section,
        ),
    )
    shared.opts.add_option(
        key='tagger_batch_size',
        info=shared.OptionInfo(
            1024,
            label='batch size for large queries',
            section=section,
        ),
    )


def split_str(string: str, separator=',') -> List[str]:
    return [x.strip() for x in string.split(separator) if x]


class InterrogatorSettings:
    kamojis = set(split_str(DEFAULT_KAMOJIS))
    output_filename_format = DEFAULT_OFF

    @classmethod
    def set_us_excl(cls):
        ruxs = getattr(shared.opts, 'tagger_repl_us_excl', DEFAULT_KAMOJIS)
        cls.kamojis = set(split_str(ruxs))

    @classmethod
    def set_output_filename_format(cls):
        fnfmt = getattr(shared.opts, 'tagger_out_filename_fmt', DEFAULT_OFF)
        if fnfmt[-12:] == '.[extension]':
            print("refused to write an image extension")
            fnfmt = fnfmt[:-12] + '.[output_extension]'

        cls.output_filename_format = fnfmt.strip()