"""Contains :class:`Arbitrary` class and the machinery it needs to work.
You shouldn't need anything other than this class from here.
"""
from abc import ABCMeta, abstractmethod
import re
import inspect
import six
from itertools import islice, count
from functools import wraps

registry = {}
registry[None] = []


class combomethod(object):
    def __init__(self, instance_method, class_method=None):
        self.instance_method = instance_method

        if class_method is None:
            @wraps(instance_method)
            def class_method(cls, *args, **kwargs):
                return instance_method(cls._inst, *args, **kwargs)

        self.class_method = class_method

    def __get__(self, obj=None, objtype=None):
        @wraps(self.instance_method)
        def _wrapper(*args, **kwargs):
            if obj is not None:
                return self.instance_method(obj, *args, **kwargs)
            elif args and isinstance(args[0], Arbitrary):
                return self.instance_method(args[0], *args[1:], **kwargs)
            else:
                return self.class_method(objtype, *args, **kwargs)

        return _wrapper


class ArbitraryMeta(ABCMeta):
    arbitrary_cls = None

    def __new__(meta, name, bases, attrs):
        cls = super(ArbitraryMeta, meta).__new__(meta, name, bases, attrs)

        if name == 'Arbitrary':
            # set the Arbitrary class to an ArbitraryMeta class variable
            # we can't just access Arbitrary class in this __new__ method
            meta.arbitrary_cls = cls
            return cls
        if not meta.arbitrary_cls:
            # Arbitrary class isn't created yet, now creating its superclass
            return cls

        # Arbitrary class already created, now creating some its subclass

        # add hard limiting to gen_serial method
        gen_serial_old = getattr(cls, 'gen_serial')

        @wraps(gen_serial_old)
        def gen_serial_new(self, amount, gen_serial=gen_serial_old):
            return islice(gen_serial(self, amount), amount)

        setattr(cls, 'gen_serial', gen_serial_new)

        for meth_name in ['next_random', 'gen_serial', 'could_generate',
                          'shrink']:
            meth_src = getattr(cls, meth_name)
            meth_new = combomethod(meth_src)
            setattr(cls, meth_name, meth_new)

        if not cls.__dict__.get('__noinst__', False):
            cls._inst = cls()
            cls.__doc__ += '\n%s' % cls.sample_doc()

        registry[None].append(cls)

        def cls_call(self, *args, **kwargs):
            self_args = self._._args['args']
            self_kwargs = self._._args['kwargs']
            args = args + self_args[len(args):]
            self_kwargs.update(kwargs)
            kwargs = self_kwargs

            return type(self)(*args, **kwargs)

        cls.__call__ = cls_call

        return cls

    def __call__(cls, *args, **kwargs):
        inst = super(ArbitraryMeta, cls).__call__(*args, **kwargs)
        inst._args = {'args': args, 'kwargs': kwargs}
        return inst


class Arbitrary(six.with_metaclass(ArbitraryMeta)):
    """Base class for all arbitraries, both built-in and custom.

    To create a custom arbitrary, you need to subclass from :class:`Arbitrary`
    and override (some or all) instance methods listed below."""

    @abstractmethod
    def next_random(self):
        """Get next random value."""
        raise StopIteration

    @abstractmethod
    def gen_serial(self, amount):
        """Generate or return a sequence of serial values.

        Only first ``amount`` elements will be used anyway, no matter how many
        are actually generated (even infinite stream of value is possible).

        :param amount: Number of requested elements, generated value can
                       depend on it (see :class:`.list_` implementation).
        """
        return []

    @abstractmethod
    def could_generate(self, x):
        """Check if x could've been generated by self.

        Required iff you use decorators like :func:`.filter_possible` or
        :func:`.until_possible`."""
        raise NotImplementedError

    @abstractmethod
    def shrink(self, x):
        """Get possible shrinkings of ``x``.

        :param x: Value previously generated by this arbitrary.
        """
        return []

    @classmethod
    def sample(cls):
        try:
            args = cls.__init__.args_for_sample['args']
            kwargs = cls.__init__.args_for_sample['kwargs']
            argsdef = cls.__init__.args_for_sample['def']
        except AttributeError:
            args = []
            kwargs = {}
            argsdef = ''

        inst = cls(*args, **kwargs)

        dct = {
            'args': {
                'args': args,
                'kwargs': kwargs,
                'def': argsdef,
            },
            'random': (inst.next_random() for _ in count()),
            'serial': (
                (2 ** p2, inst.gen_serial(2 ** p2))
                for p2 in count()
            )
        }
        return dct

    @classmethod
    def sample_doc(cls):
        AMOUNT = 30
        sample = cls.sample()

        args_formatted = sample['args']['def']

        def append(line, indent=' ' * 4):
            res.append(indent + line)

        def append_elems(gen):
            for i, elem in enumerate(islice(gen, AMOUNT)):
                append(('>>>' if i == 0 else '...') + ' ' + repr(elem) + ',')

        res = []
        append('', indent='')
        append('.. rst-class:: html-toggle', indent='')
        append('', indent='')
        append('**Samples:**', indent='')
        append('', indent='')
        append('.. code-block:: python', indent='')
        append('')
        append('>>> # example initialization:')
        append('>>> arb = %s(%s)' % (cls.__name__, args_formatted))
        append('>>> # random elements:')
        append_elems(sample['random'])
        append('')
        append('>>> ' + '# serial elements:')
        gen = next(islice(sample['serial'], AMOUNT.bit_length(), None))[1]
        append_elems(gen)
        append('')
        return '\n'.join(res)

    @staticmethod
    def args_for_sample(*args, **kwargs):
        def wrapper(method):
            method_src = inspect.getsource(method)
            m = re.match(r'\s*@Arbitrary.args_for_sample\((.*)\)', method_src)
            argsdef = m.group(1)
            method.args_for_sample = {
                'args': args,
                'kwargs': kwargs,
                'def': argsdef
            }
            return method

        return wrapper

    @staticmethod
    def set_for(type, arb=None):
        def wrapper(arb):
            if type in registry:
                raise LookupError('Arbitrary class already set for %s' % type)
            registry[type] = arb
            return arb

        if arb is None:
            return wrapper
        else:
            return wrapper(arb)

    @staticmethod
    def get_for(type):
        try:
            return registry[type]
        except KeyError:
            raise LookupError('No Arbitrary class is set for %s' % type)

    @staticmethod
    def get_all():
        return registry[None]


class MappedArbitrary(Arbitrary):
    __noinst__ = True

    @abstractmethod
    def pack(self, x):
        pass

    @abstractmethod
    def unpack(self, x):
        pass

    def next_random(self):
        return self.pack(self.inner_arb.next_random())

    def gen_serial(self, amount):
        for el in self.inner_arb.gen_serial(amount):
            yield self.pack(el)

    def could_generate(self, x):
        return self.inner_arb.could_generate(self.unpack(x))

    def shrink(self, x):
        return map(self.pack, self.inner_arb.shrink(self.unpack(x)))
