import functools
import textwrap

import pyperf


IMPLEMENTATIONS = {
    'dict': """\
    cls = dict
    """,
    'multidict_c': """\
    from multidict._multidict import MultiDict as cls, istr
    """,
    'cimultidict_c': """\
    from multidict._multidict import CIMultiDict as cls, istr
    """,
    'multidict_py': """\
    from multidict._multidict_py import MultiDict as cls, istr
    """,
    'cimultidict_py': """\
    from multidict._multidict_py import CIMultiDict as cls, istr
    """
}

INIT = """\
dct = cls()
"""

FILL = """\
for i in range(20):
    dct['key'+str(i)] = str(i)

key = 'key10'
"""


FILL_ISTR = """\
for i in range(20):
    key = istr('key'+str(i))
    dct[key] = str(i)

key = istr('key10')
"""


SET_ITEM = """\
dct[key] = '1'
dct[key] = '2'
dct[key] = '3'
dct[key] = '4'
dct[key] = '5'
dct[key] = '6'
dct[key] = '7'
dct[key] = '8'
dct[key] = '9'
dct[key] = '10'
""" * 10


GET_ITEM = """\
dct[key]
dct[key]
dct[key]
dct[key]
dct[key]
dct[key]
dct[key]
dct[key]
dct[key]
dct[key]
""" * 10


ADD = """\
add(key, '1')
add(key, '2')
add(key, '3')
add(key, '4')
add(key, '5')
add(key, '6')
add(key, '7')
add(key, '8')
add(key, '9')
add(key, '10')
""" * 10

SETUP_ADD = """\
add = dct.add
"""

def benchmark_name(name, ctx, prefix=None, use_prefix=False):
    if use_prefix:
        return '%s%s' % (prefix % ctx, name)

    return name


def add_impl_option(cmd, args):
    if args.impl:
        cmd.extend(['--impl', args.impl])


if __name__ == '__main__':
    runner = pyperf.Runner(add_cmdline_args=add_impl_option)

    parser = runner.argparser
    parser.description = ('Allows to measure performance of MultiMapping and '
                          'MutableMultiMapping implementations')
    parser.add_argument('--impl', choices=sorted(IMPLEMENTATIONS),
                        help='specific implementation to benchmark')

    options = parser.parse_args()
    implementations = (options.impl,) if options.impl else IMPLEMENTATIONS
    inner_loops = 50

    for impl in implementations:
        # print("=======================", impl, "======================")
        imports = textwrap.dedent(IMPLEMENTATIONS[impl])
        name = functools.partial(benchmark_name, ctx=dict(impl=impl),
                                 prefix='(impl = %(impl)s) ',
                                 use_prefix=len(implementations) > 1)

        runner.timeit(name('setitem str'),
                      SET_ITEM, imports + INIT + FILL,
                      inner_loops=inner_loops)
        runner.timeit(name('getitem str'),
                      GET_ITEM, imports + INIT + FILL,
                      inner_loops=inner_loops)

        # MultiDict specific
        if impl == 'dict':
            continue

        runner.timeit(name('setitem istr'),
                      SET_ITEM, imports + INIT + FILL_ISTR,
                      inner_loops=inner_loops)
        runner.timeit(name('getitem istr'),
                      GET_ITEM, imports + INIT + FILL_ISTR,
                      inner_loops=inner_loops)
        runner.timeit(name('add str'),
                      ADD, imports + INIT + FILL + SETUP_ADD,
                      inner_loops=inner_loops)
        runner.timeit(name('add istr'),
                      ADD, imports + INIT + FILL_ISTR + SETUP_ADD,
                      inner_loops=inner_loops)
