# Evaluation machinery:
#   - it flattens (evaluates) the context, then interate over a cartesian
#     product of yacl classes (or one class) given, and the context set items.
#
# The actual magic of evaluation happens in yacl_base.get* methods.
# The logic in this file is still tightly integrated with logic in the yacl_base.
# It is provided in a separate file, to reduce the source code file sizes,
# and to allow using yacl_base without needing to import yacl_eval.go.
# Technically there is a circular dependency between the two, but it works.

import yacl_base
import yacl_context
import sys  # for stderr
from typing import Iterable, Any, List, Type, Union, Optional

export = yacl_base.export


def _evaluate(cls: Type[yacl_base.YaclBaseClass], current_context_item, execute: bool = True) -> Optional[dict[str, Any]]:
    assert yacl_base._current_full_eval_context.current_evaluated_class is None, "Some kind of weird recurssion! That is not allowed"
    assert yacl_base._current_full_eval_context.current_evaluated_context is None, "Some kind of weird recurssion! That is not allowed"
    assert len(yacl_base._current_full_eval_context.memoized_attrib_values) == 0, "Error. attrib cache should be empty. Unclean previous call to evaluate_?"

    yacl_base._current_full_eval_context.current_most_derived_class = cls
    yacl_base._current_full_eval_context.current_evaluated_class = cls
    yacl_base._current_full_eval_context.current_evaluated_context = current_context_item
    yacl_base._current_full_eval_context.current_evaluated_attrib_name = []
    yacl_base._current_full_eval_context.memoized_attrib_values = {}  # cache

    # TODO(baryluk): Assign all the above to local variables, to avoid repeated key lookups in yacl_base module and in FullEvalContext object.

    _enabled = yacl_base.get('_enable')

    if _enabled:
        result = {}
        for attrib_name in dir(cls):
            if attrib_name.startswith("_") or attrib_name == "_abc_impl":
                continue

            yacl_base._current_full_eval_context.current_evaluated_attrib_name = []
            # yacl_base._current_full_eval_context.current_evaluated_attrib_name.append(attrib_name)

            attrib_value = yacl_base.get(attrib_name)

            # assert len(yacl_base._current_full_eval_context.current_evaluated_attrib_name) == 1, f"Implementation error! attrib eval stack not properly undone, while evaluating attrib '{attrib_name}'. Got: {yacl_base._current_full_eval_context.current_evaluated_attrib_name}"
            # assert yacl_base._current_full_eval_context.current_evaluated_attrib_name[0] == attrib_name, f"Implementation error! attrib eval stack not properly undone, while evaluating attrib '{attrib_name}'. Got: {yacl_base._current_full_eval_context.current_evaluated_attrib_name}"

            assert len(yacl_base._current_full_eval_context.current_evaluated_attrib_name) == 0, f"Implementation error! attrib eval stack not properly undone, while evaluating attrib '{attrib_name}'. Expected: [], Got: {yacl_base._current_full_eval_context.current_evaluated_attrib_name}."

            assert attrib_value is not yacl_base.External, f"attribute '{attrib_name}' is external and must come from external context, yet get() did not process it"
            assert attrib_value is not yacl_base.ExternalPrimary, f"attribute '{attrib_name}' is external_primary and must come from external context, yet get() did not process it"

            result[attrib_name] = attrib_value

            assert yacl_base._current_full_eval_context.current_evaluated_class is cls, f"Implementation error! _current_evaluated_class not restored properly, while evaluating attrib '{attrib_name}'. Expected: {cls}, Got: {yacl_base._current_full_eval_context.current_evaluated_class}'"
            assert attrib_name in yacl_base._current_full_eval_context.memoized_attrib_values, f"Implementation error! Caching broken? Expected attrib '{attrib_name}' to be in cache after evaluation of get('{attrib_name}'), but found nothing."

        # _output(result)

    # Poison all the values in the current eval context.
    yacl_base._current_full_eval_context.current_most_derived_class = None
    yacl_base._current_full_eval_context.current_evaluated_class = None
    yacl_base._current_full_eval_context.current_evaluated_context = None
    yacl_base._current_full_eval_context.current_evaluated_attrib_name = []
    yacl_base._current_full_eval_context.memoized_attrib_values = {}

    if _enabled:
        result['__type'] = cls

        if hasattr(cls, '_execute'):
            # We use standard Python mro for looking up _execute.
            result['__execute'] = cls._execute  # For debugging.
            if execute:
                cls._execute(result)

        return result

    return None


import abc


@export
def go(cls_or_list: Union[Type[yacl_base.YaclBaseClass], Iterable[Type[yacl_base.YaclBaseClass]]], /, *, context_set: Iterable[dict[str, Any]] = [{}], execute: bool = True) -> List[dict[str, Any]]:
    """
    Evaluate each element of the context_set (after expanding any encountered Expand and Cross references),
    against cls_or_list (which can be a single class, or a list of classes).
    Evaluation against list of classes, is equivalent to multiple individual calls to `go`
    with each class in turn, just faster due to sharing of some precomputed state.
    Evaluation order is not guaranteed and implementation defined.

    Additionally go will return all the results. This can be useful for defining
    nested data structures using other yacl classes.
    """
    list_of_classes = cls_or_list if isinstance(cls_or_list, list) else [cls_or_list]

    # Iterate over flattened (expanded) context and eval all objects.
    flat_context = list(yacl_context.flatten_context(context_set))

    all_results = []

    try:
        yacl_base._current_full_eval_context.current_full_context = flat_context

        for cls in list_of_classes:
            # assert type(cls) is type, f"Can only pass classes, not objects or instances of the class. Got: {cls} of type {type(cls)}"
            # <class 'abc.ABCMeta'>
            assert type(cls) is abc.ABCMeta, f"Can only pass classes, derived from abc.ABCMeta. Got: {cls} of type {type(cls)}"
            mro = cls.mro()
            assert yacl_base.YaclBaseClass in mro, f"Can only passes classes that derive from YaclBaseClass: Got: {cls}, with mro: {mro}"
            assert len(mro) >= 4, f"Expected at least 4 elements in mro: [..., YourClass, YaclBaseClass, ABC, object]. Got: {mro}"
            assert yacl_base.YaclBaseClass is mro[-3], f"YaclBaseClass should be a true base class of {cls} (ignoring 'object'). Got mro: {mro}."

            for context_item in flat_context:
                # print(context_item)
                # print()
                try:
                    # Evaluate and possibly execute. Also gather results to return later.
                    # yield _evaluate(cls, context_item)  # Problem with using yield is that a lot of error checking will be done lazily, which is not great.
                    result = _evaluate(cls, context_item, execute=execute)
                    if result is not None:  # We allow result to be 'empty' dict. It will never be really empty, because of '__type' and possible '__execute 'key.
                        # But result will be None, if _enable attribute did not return True.
                        all_results.append(result)
                except:
                    print(f"\n\nError evaluating class {cls.__module__}.{cls.__qualname__} with context {context_item}", file=sys.stderr)
                    print(f"Current eval stack:")
                    for i, stack_element in enumerate(yacl_base._total_eval_stack):
                        print(f"    {i:3}  {stack_element}")
                    # if flags.keep_going: ...
                    raise
    finally:
        yacl_base._current_full_eval_context.current_full_context = None

    return all_results