import inspect
from typing import Dict, Any
from eth2spec.utils.ssz.ssz_typing import View
from eth2spec.utils.ssz.ssz_impl import serialize


def vector_test(description: str = None):
    """
    vector_test decorator: Allow a caller to pass "generator_mode=True" to make the test yield data,
     but behave like a normal test (ignoring the yield, but fully processing) a test when not in "generator_mode"
    This should always be the most outer decorator around functions that yield data.
    This is to deal with silent iteration through yielding function when in a pytest
     context (i.e. not in generator mode).
    :param description: Optional description for the test to add to the metadata.
    :return: Decorator.
    """
    def runner(fn):
        # this wraps the function, to yield type-annotated entries of data.
        # Valid types are:
        #   - "meta": all key-values with this type can be collected by the generator, to put somewhere together.
        #   - "ssz": raw SSZ bytes
        #   - "data": a python structure to be encoded by the user.
        def entry(*args, **kw):

            def generator_mode():
                if description is not None:
                    # description can be explicit
                    yield 'description', 'meta', description

                # transform the yielded data, and add type annotations
                for data in fn(*args, **kw):
                    # if not 2 items, then it is assumed to be already formatted with a type:
                    # e.g. ("bls_setting", "meta", 1)
                    if len(data) != 2:
                        yield data
                        continue
                    # Try to infer the type, but keep it as-is if it's not a SSZ type or bytes.
                    (key, value) = data
                    if value is None:
                        continue
                    if isinstance(value, View):
                        yield key, 'ssz', serialize(value)
                    elif isinstance(value, bytes):
                        yield key, 'ssz', value
                    elif isinstance(value, list) and all([isinstance(el, (View, bytes)) for el in value]):
                        for i, el in enumerate(value):
                            if isinstance(el, View):
                                yield f'{key}_{i}', 'ssz', serialize(el)
                            elif isinstance(el, bytes):
                                yield f'{key}_{i}', 'ssz', el
                        yield f'{key}_count', 'meta', len(value)
                    else:
                        # Not a ssz value.
                        # The data will now just be yielded as any python data,
                        #  something that should be encodable by the generator runner.
                        yield key, 'data', value

            # check generator mode, may be None/else.
            # "pop" removes it, so it is not passed to the inner function.
            if kw.pop('generator_mode', False) is True:
                # return the yielding function as a generator object.
                # Don't yield in this function itself, that would make pytest skip over it.
                return generator_mode()
            else:
                # Just complete the function, ignore all yielded data,
                # we are not using it (or processing it, i.e. nearly zero efficiency loss)
                # Pytest does not support yielded data in the outer function, so we need to wrap it like this.
                for _ in fn(*args, **kw):
                    continue
                return None

        return entry

    return runner


def with_meta_tags(tags: Dict[str, Any]):
    """
    Decorator factory, yields meta tags (key, value) pairs to the output of the function.
    Useful to build test-vector annotations with.
    :param tags: dict of tags
    :return: Decorator.
    """
    def runner(fn):
        def entry(*args, **kw):
            yielded_any = False
            for part in fn(*args, **kw):
                yield part
                yielded_any = True
            # Do not add tags if the function is not returning a dict at all (i.e. not in generator mode).
            # As a pytest, we do not want to be yielding anything (unsupported by pytest)
            if yielded_any:
                for k, v in tags.items():
                    yield k, 'meta', v
        return entry
    return runner


def build_transition_test(fn, pre_fork_name, post_fork_name, fork_epoch=None):
    """
    Handles the inner plumbing to generate `transition_test`s.
    See that decorator in `context.py` for more information.
    """
    def _adapter(*args, **kwargs):
        post_spec = kwargs["phases"][post_fork_name]

        pre_fork_counter = 0

        def pre_tag(obj):
            nonlocal pre_fork_counter
            pre_fork_counter += 1
            return obj

        def post_tag(obj):
            return obj

        yield "post_fork", "meta", post_fork_name

        has_fork_epoch = False
        if fork_epoch:
            kwargs["fork_epoch"] = fork_epoch
            has_fork_epoch = True
            yield "fork_epoch", "meta", fork_epoch

        # massage args to handle an optional custom state using
        # `with_custom_state` decorator
        expected_args = inspect.getfullargspec(fn)
        if "phases" not in expected_args.kwonlyargs:
            kwargs.pop("phases", None)

        for part in fn(*args,
                       post_spec=post_spec,
                       pre_tag=pre_tag,
                       post_tag=post_tag,
                       **kwargs):
            if part[0] == "fork_epoch":
                has_fork_epoch = True
            yield part
        assert has_fork_epoch

        if pre_fork_counter > 0:
            yield "fork_block", "meta", pre_fork_counter - 1
    return _adapter