"""Evaluation engine — recursive evaluation with simplification."""

from __future__ import annotations
from typing import Any, Dict, Set
from .expr import Expr, Symbol
from .rules import SIMPLIFICATION_RULES, replace_all

# Lazy load extended rules to avoid circular imports
_extended_loaded = False
_extended_rules = []

def _ensure_extended():
    global _extended_loaded, _extended_rules
    if not _extended_loaded:
        try:
            from .extended import EXTENDED_RULES
            from .extended2 import EXTENDED2_RULES
            from .extended3 import EXTENDED3_RULES
            from .extended4 import EXTENDED4_RULES
            from .extended5 import EXTENDED5_RULES
            from .extended6 import register as ext6
            from .extended7 import register as ext7
            from .extended8 import register as ext8
            from .extended9 import register as ext9
            from .extended10 import register as ext10
            from .extended11 import register as ext11
            from .extended12 import register as ext12
            from .extended13 import register as ext13
            from .extended14 import register as ext14
            from .extended15 import register as ext15
            from .extended16 import register as ext16
            from .extended17 import register as ext17
            from .extended18 import register as ext18
            from .extended19 import register as ext19
            from .extended20 import register as ext20
            from .extended21 import register as ext21
            from .extended22 import register as ext22
            from .extended23 import register as ext23
            from .extended24 import register as ext24
            from .extended25 import register as ext25
            from .extended26 import register as ext26
            from .extended27 import register as ext27
            from .extended28 import register as ext28
            from .extended29 import register as ext29
            from .extended30 import register as ext30
            from .extended31 import register as ext31
            from .extended32 import register as ext32
            from .extended33 import register as ext33
            from .visualization import VISUALIZATION_RULES
            _extended_rules = (EXTENDED_RULES + EXTENDED2_RULES + EXTENDED3_RULES + 
                             EXTENDED4_RULES + EXTENDED5_RULES + ext6() + ext7() + 
                             ext8() + ext9() + ext10() + ext11() + ext12() + ext13() + 
                             ext14() + ext15() + ext16() + ext17() + ext18() + ext19() + 
                             ext20() + ext21() + ext22() + ext23() + ext24() + ext25() + 
                             ext26() + ext27() + ext28() + ext29() + ext30() + ext31() + 
                             ext32() + ext33() + VISUALIZATION_RULES)
        except Exception as e:
            # Fallback: try loading what we can
            try:
                from .extended import EXTENDED_RULES
                from .extended2 import EXTENDED2_RULES
                from .extended3 import EXTENDED3_RULES
                from .extended4 import EXTENDED4_RULES
                from .extended5 import EXTENDED5_RULES
                from .visualization import VISUALIZATION_RULES
                _extended_rules = EXTENDED_RULES + EXTENDED2_RULES + EXTENDED3_RULES + EXTENDED4_RULES + EXTENDED5_RULES + VISUALIZATION_RULES
            except Exception:
                try:
                    from .extended import EXTENDED_RULES
                    from .extended2 import EXTENDED2_RULES
                    from .extended3 import EXTENDED3_RULES
                    _extended_rules = EXTENDED_RULES + EXTENDED2_RULES + EXTENDED3_RULES
                except Exception:
                    try:
                        from .extended import EXTENDED_RULES
                        from .extended2 import EXTENDED2_RULES
                        _extended_rules = EXTENDED_RULES + EXTENDED2_RULES
                    except Exception:
                        try:
                            from .extended import EXTENDED_RULES
                            _extended_rules = EXTENDED_RULES
                        except Exception:
                            _extended_rules = []
        _extended_loaded = True

# Hold attributes
_HOLD_ALL: Set[str] = set()
_HOLD_FIRST: Set[str] = set()

_MAX_DEPTH = 256
_memo: Dict[int, Any] = {}


def set_hold_all(head: str):
    _HOLD_ALL.add(head)


def set_hold_first(head: str):
    _HOLD_FIRST.add(head)


def clear_memo():
    _memo.clear()


def evaluate(expr, depth: int = 0) -> Any:
    """Recursively evaluate an expression, applying simplification rules."""
    if depth > _MAX_DEPTH:
        raise RecursionError(f"Evaluation depth exceeded {_MAX_DEPTH}")

    # Atoms evaluate to themselves
    if isinstance(expr, (int, float, str, bool)):
        return expr
    if isinstance(expr, Symbol):
        return expr

    if not isinstance(expr, Expr):
        return expr

    # Memoization
    key = hash(expr)
    if key in _memo:
        cached = _memo[key]
        if cached == expr:
            return cached

    head = expr.head

    # Evaluate arguments (respecting Hold attributes)
    if head in _HOLD_ALL:
        new_args = expr.args
    elif head in _HOLD_FIRST:
        new_args = (expr.args[0],) + tuple(evaluate(a, depth + 1) for a in expr.args[1:]) if expr.args else ()
    else:
        new_args = tuple(evaluate(a, depth + 1) for a in expr.args)

    new_expr = Expr(head, *new_args)

    # Apply simplification rules
    result = replace_all(new_expr, SIMPLIFICATION_RULES)

    # If simplification didn't change anything, try extended rules
    if isinstance(result, Expr) and result == new_expr:
        _ensure_extended()
        if _extended_rules:
            ext_result = replace_all(new_expr, _extended_rules)
            if ext_result != new_expr:
                result = ext_result

    # If result changed, evaluate again
    if isinstance(result, Expr) and result != new_expr:
        result = evaluate(result, depth + 1)

    _memo[key] = result
    return result
