Source code for marivo.semantic._expression_binding

"""Compile and evaluate task-local semantic expression bindings."""

from __future__ import annotations

import ast
import copy
import hashlib
import inspect
import json
import math
import textwrap
from collections.abc import Callable, Mapping, Sequence
from contextvars import ContextVar
from dataclasses import dataclass, field
from types import CellType, FunctionType, MappingProxyType, ModuleType
from typing import Literal, TypeAlias, cast

import ibis
import ibis.expr.types as ir

from marivo.refs import (
    EntityKind,
    FieldKind,
    MetricKind,
    Ref,
    RefPayloadV1,
    SemanticKind,
    SemanticKindTag,
    _decode_ref_payload,
)
from marivo.semantic._definition_expression import describe_expression
from marivo.semantic.definition import ExpressionDescription, _UnsupportedExpression
from marivo.semantic.errors import (
    ErrorKind,
    SemanticError,
    SemanticLoadError,
    SemanticRuntimeError,
)
from marivo.semantic.validator import (
    validate_entity_table_body_ast,
    validate_event_body_ast,
    validate_metric_body_ast,
)

_FIELD_KINDS = frozenset(
    {
        SemanticKind.DIMENSION,
        SemanticKind.TIME_DIMENSION,
        SemanticKind.MEASURE,
    }
)
EventConstant: TypeAlias = str | int | float | bool


@dataclass(frozen=True, slots=True)
class ExpressionBindingV1:
    """One declared field-ref application and its entity argument position."""

    field_ref: RefPayloadV1
    entity_position: int

    def __post_init__(self) -> None:
        if self.field_ref.kind not in _FIELD_KINDS:
            raise ValueError(
                "expression binding field_ref must be a dimension, time_dimension, or measure"
            )
        if type(self.entity_position) is not int or self.entity_position < 0:
            raise ValueError("expression binding entity_position must be a non-negative int")

    def to_ref(self) -> Ref[FieldKind]:
        return cast("Ref[FieldKind]", _decode_ref_payload(self.field_ref))


@dataclass(frozen=True, slots=True)
class ExpressionBody:
    """Process-local callable plus deterministic definition identity."""

    callable: Callable[..., object]
    body_ast_hash: str
    parameter_count: int
    bindings: tuple[ExpressionBindingV1, ...]
    source_column: str | None = None
    source_columns: tuple[str, ...] = ()
    description: ExpressionDescription = field(
        default_factory=lambda: _UnsupportedExpression("description_unavailable")
    )

    def __post_init__(self) -> None:
        if not callable(self.callable):
            raise TypeError("expression body callable must be callable")
        if type(self.body_ast_hash) is not str or not self.body_ast_hash:
            raise ValueError("expression body hash must be a non-empty string")
        if type(self.parameter_count) is not int or self.parameter_count < 0:
            raise ValueError("expression body parameter_count must be a non-negative int")
        if self.source_column is not None and (
            type(self.source_column) is not str or not self.source_column
        ):
            raise ValueError("expression body source_column must be non-empty when provided")
        if any(type(column) is not str or not column for column in self.source_columns):
            raise ValueError("expression body source_columns must contain non-empty strings")
        if len(set(self.source_columns)) != len(self.source_columns):
            raise ValueError("expression body source_columns must be unique")
        for binding in self.bindings:
            if binding.entity_position >= self.parameter_count:
                raise ValueError("expression binding entity_position must be below parameter_count")

    @classmethod
    def for_column(cls, column: str) -> ExpressionBody:
        """Create one deterministic source-free direct-column body."""
        if type(column) is not str or not column:
            raise ValueError("direct-column expression requires a non-empty column name")

        def accessor(entity_alias: ir.Table) -> ir.Value:
            return entity_alias[column]

        encoded = json.dumps(
            {
                "schema": "marivo.expression_body.column/v1",
                "column": column,
            },
            ensure_ascii=True,
            separators=(",", ":"),
            sort_keys=True,
        ).encode()
        return cls(
            callable=accessor,
            body_ast_hash=f"sha256:{hashlib.sha256(encoded).hexdigest()}",
            parameter_count=1,
            bindings=(),
            source_column=column,
            source_columns=(column,),
        )


@dataclass(frozen=True, slots=True)
class CompiledExpressionSidecar:
    """Immutable expression bodies and ownership facts for one catalog."""

    bodies: Mapping[Ref[SemanticKindTag], ExpressionBody]
    field_owners: Mapping[Ref[FieldKind], Ref[EntityKind]]
    catalog_refs: frozenset[Ref[SemanticKindTag]]
    default_cumulative_axes: frozenset[Ref[MetricKind]] = frozenset()

    def __post_init__(self) -> None:
        bodies = dict(self.bodies)
        owners = dict(self.field_owners)
        refs = frozenset(self.catalog_refs)
        for ref in bodies:
            _require_exact_ref(ref, parameter="bodies key")
        for field_ref, entity_ref in owners.items():
            _require_exact_ref(field_ref, allowed=_FIELD_KINDS, parameter="field owner key")
            _require_exact_ref(
                entity_ref,
                allowed=frozenset({SemanticKind.ENTITY}),
                parameter="field owner value",
            )
        for ref in refs:
            _require_exact_ref(ref, parameter="catalog ref")
        object.__setattr__(self, "bodies", MappingProxyType(bodies))
        object.__setattr__(self, "field_owners", MappingProxyType(owners))
        object.__setattr__(self, "catalog_refs", refs)


@dataclass(frozen=True, slots=True)
class ExpressionBodyFrame:
    """One active expression body and its exact entity aliases."""

    owning_ref: Ref[SemanticKindTag]
    ordered_entity_refs: tuple[Ref[EntityKind], ...]
    ordered_entity_aliases: tuple[ir.Table, ...]
    declared_bindings: tuple[ExpressionBindingV1, ...]

    def __post_init__(self) -> None:
        _require_exact_ref(self.owning_ref, parameter="owning_ref")
        if len(self.ordered_entity_refs) != len(self.ordered_entity_aliases):
            raise ValueError("expression frame entity refs and aliases must have equal arity")


@dataclass(frozen=True, slots=True)
class ExpressionBindingContext:
    """Task-local catalog and nested expression-body stack."""

    catalog_definition_fingerprint: str
    expression_sidecar: CompiledExpressionSidecar
    body_frames: tuple[ExpressionBodyFrame, ...]

    def __post_init__(self) -> None:
        if not self.catalog_definition_fingerprint:
            raise ValueError("catalog definition fingerprint must be non-empty")
        if not self.body_frames:
            raise ValueError("expression binding context requires one body frame")


_EXPRESSION_BINDING_CONTEXT: ContextVar[ExpressionBindingContext | None] = ContextVar(
    "_EXPRESSION_BINDING_CONTEXT",
    default=None,
)


def _require_exact_ref(
    value: object,
    *,
    allowed: frozenset[SemanticKind] | None = None,
    parameter: str,
) -> Ref[SemanticKindTag]:
    if type(value) is not Ref:
        raise TypeError(f"{parameter} must be an exact Ref; received {type(value).__name__}")
    ref = cast("Ref[SemanticKindTag]", value)
    if allowed is not None and ref.kind not in allowed:
        expected = ", ".join(sorted(kind.value for kind in allowed))
        raise TypeError(f"{parameter} must have kind in {{{expected}}}; received {ref.kind.value}")
    return ref


def _binding_runtime_error(
    kind: ErrorKind,
    message: str,
    *,
    refs: Sequence[Ref[SemanticKindTag]] = (),
    expected: str | None = None,
    received: str | None = None,
    details: dict[str, object] | None = None,
) -> SemanticRuntimeError:
    return SemanticRuntimeError(
        kind=kind,
        message=message,
        refs=tuple(ref.key for ref in refs),
        hint=(
            "Use ms.bind(field_ref, entity_alias) with one declared field-kind ref inside "
            "the active loaded semantic expression body."
        ),
        expected=expected,
        received=received,
        details=dict(details or {}),
    )


def _find_function_node(tree: ast.Module, fn_name: str) -> ast.FunctionDef:
    functions = [node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef)]
    for node in functions:
        if node.name == fn_name:
            return node
    if functions:
        return functions[0]
    raise SemanticLoadError(
        kind=ErrorKind.COMPILE_ERROR,
        message=f"Expression body {fn_name!r} source does not contain a function definition.",
        expected="one inspectable Python function definition",
        received="no function definition",
    )


def _load_function_ast(fn: Callable[..., object]) -> tuple[ast.Module, ast.FunctionDef]:
    try:
        source = textwrap.dedent(inspect.getsource(fn))
        tree = ast.parse(source)
    except (OSError, TypeError, IndentationError, SyntaxError) as exc:
        raise SemanticLoadError(
            kind=ErrorKind.COMPILE_ERROR,
            message=f"Expression body {fn.__name__!r} source is unavailable or invalid.",
            expected="inspectable Python source for the decorated function",
            received=type(exc).__name__,
        ) from exc
    return tree, _find_function_node(tree, fn.__name__)


def _resolved_symbols(fn: Callable[..., object]) -> dict[str, object]:
    closure = inspect.getclosurevars(fn)
    symbols: dict[str, object] = dict(fn.__globals__)
    symbols.update(closure.globals)
    symbols.update(closure.nonlocals)
    return symbols


def _is_bind_target(node: ast.expr, symbols: Mapping[str, object]) -> bool:
    if isinstance(node, ast.Name):
        return symbols.get(node.id) is bind
    if isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name):
        owner = symbols.get(node.value.id)
        return getattr(owner, node.attr, None) is bind
    return False


def _event_constant_value(
    name: str,
    value: object,
    *,
    owning_ref: Ref[SemanticKindTag],
) -> EventConstant:
    if type(value) not in {str, int, float, bool}:
        raise SemanticLoadError(
            kind=ErrorKind.INVALID_EVENT_PREDICATE,
            message=(
                f"Event predicate {owning_ref.key!r} uses unsupported external name {name!r}."
            ),
            refs=(owning_ref.key,),
            expected="a closed immutable str, int, float, or bool constant",
            received=f"{name}={type(value).__name__}",
            hint="Replace the external value with an immutable scalar constant or a literal.",
        )
    if type(value) is float and not math.isfinite(value):
        raise SemanticLoadError(
            kind=ErrorKind.INVALID_EVENT_PREDICATE,
            message=(
                f"Event predicate {owning_ref.key!r} uses non-finite external constant {name!r}."
            ),
            refs=(owning_ref.key,),
            expected="a finite float constant",
            received=repr(value),
            hint="Replace the value with one finite immutable scalar constant.",
        )
    return cast("EventConstant", value)


def _event_closed_constants(
    function: ast.FunctionDef,
    *,
    symbols: Mapping[str, object],
    owning_ref: Ref[SemanticKindTag],
) -> dict[str, EventConstant]:
    returns = [node for node in ast.walk(function) if isinstance(node, ast.Return)]
    if len(returns) != 1 or returns[0].value is None:
        return {}
    value = returns[0].value
    if (
        isinstance(value, ast.Call)
        and isinstance(value.func, ast.Attribute)
        and isinstance(value.func.value, ast.Name)
        and value.func.value.id == "ms"
        and value.func.attr == "all_rows"
    ):
        return {}
    allowed_name_nodes: set[int] = set()
    for node in ast.walk(value):
        if not isinstance(node, ast.Call) or not (
            isinstance(node.func, ast.Attribute)
            and isinstance(node.func.value, ast.Name)
            and node.func.value.id == "ms"
            and node.func.attr == "bind"
            and len(node.args) == 2
            and isinstance(node.args[0], ast.Name)
            and isinstance(node.args[1], ast.Name)
        ):
            continue
        allowed_name_nodes.update((id(node.func.value), id(node.args[0]), id(node.args[1])))
    constants: dict[str, EventConstant] = {}
    for node in ast.walk(value):
        if not isinstance(node, ast.Name) or id(node) in allowed_name_nodes:
            continue
        if node.id not in symbols:
            raise SemanticLoadError(
                kind=ErrorKind.INVALID_EVENT_PREDICATE,
                message=(
                    f"Event predicate {owning_ref.key!r} uses unresolved external name {node.id!r}."
                ),
                refs=(owning_ref.key,),
                expected="a closed immutable str, int, float, or bool constant",
                received=node.id,
                hint="Define one immutable scalar constant in the Event module or use a literal.",
            )
        constants[node.id] = _event_constant_value(
            node.id,
            symbols[node.id],
            owning_ref=owning_ref,
        )
    return constants


def _closure_cell(value: object) -> CellType:
    def capture() -> object:
        return value

    closure = capture.__closure__
    if closure is None:
        raise AssertionError("captured value did not create a closure cell")
    return closure[0]


def _freeze_event_callable(
    fn: Callable[..., object],
    constants: Mapping[str, EventConstant],
) -> Callable[..., object]:
    if not constants:
        return fn
    frozen_globals = dict(fn.__globals__)
    frozen_globals.update(
        (name, value) for name, value in constants.items() if name in frozen_globals
    )
    closure = fn.__closure__
    if closure is not None:
        frozen_cells = tuple(
            _closure_cell(constants[name]) if name in constants else cell
            for name, cell in zip(fn.__code__.co_freevars, closure, strict=True)
        )
    else:
        frozen_cells = None
    frozen = FunctionType(
        fn.__code__,
        frozen_globals,
        fn.__name__,
        fn.__defaults__,
        frozen_cells,
    )
    frozen.__kwdefaults__ = fn.__kwdefaults__
    frozen.__annotations__ = dict(fn.__annotations__)
    frozen.__module__ = fn.__module__
    frozen.__qualname__ = fn.__qualname__
    return frozen


class _BindingCollector(ast.NodeVisitor):
    def __init__(
        self,
        *,
        symbols: Mapping[str, object],
        parameter_positions: Mapping[str, int],
        owning_ref: Ref[SemanticKindTag],
    ) -> None:
        self._symbols = symbols
        self._parameter_positions = parameter_positions
        self._owning_ref = owning_ref
        self.bindings: list[ExpressionBindingV1] = []
        self.index_by_key: dict[tuple[SemanticKind, str, int], int] = {}

    def visit_Call(self, node: ast.Call) -> None:
        if not _is_bind_target(node.func, self._symbols):
            if isinstance(node.func, ast.Name) and type(self._symbols.get(node.func.id)) is Ref:
                legacy_ref = cast("Ref[SemanticKindTag]", self._symbols[node.func.id])
                raise SemanticLoadError(
                    kind=ErrorKind.INVALID_BINDING_REF,
                    message=(
                        f"Expression body {self._owning_ref.key!r} calls Ref "
                        f"{legacy_ref.key!r} directly. Ref values are immutable identities, "
                        "not expression callables."
                    ),
                    refs=(self._owning_ref.key, legacy_ref.key),
                    expected="ms.bind(field_ref, entity_parameter)",
                    received="field_ref(entity_parameter)",
                )
            self.generic_visit(node)
            return
        if len(node.args) != 2 or node.keywords or not isinstance(node.args[0], ast.Name):
            raise SemanticLoadError(
                kind=ErrorKind.BINDING_ALIAS_NOT_DIRECT,
                message=(
                    f"Expression body {self._owning_ref.key!r} must call "
                    "ms.bind(field_ref, entity_parameter) with two direct arguments."
                ),
                refs=(self._owning_ref.key,),
                expected="ms.bind(field_ref, entity_parameter)",
                received=ast.dump(node, include_attributes=False),
            )
        value = self._symbols.get(node.args[0].id)
        if type(value) is not Ref:
            raise SemanticLoadError(
                kind=ErrorKind.INVALID_BINDING_REF,
                message=(
                    f"Expression body {self._owning_ref.key!r} binds a value that is not "
                    "an exact field ref."
                ),
                refs=(self._owning_ref.key,),
                expected="a dimension, time_dimension, or measure Ref",
                received=type(value).__name__,
            )
        field_ref = cast("Ref[SemanticKindTag]", value)
        if field_ref.kind not in _FIELD_KINDS:
            raise SemanticLoadError(
                kind=ErrorKind.INVALID_BINDING_REF,
                message=(
                    f"Expression body {self._owning_ref.key!r} calls non-field ref "
                    f"{field_ref.key!r}."
                ),
                refs=(self._owning_ref.key, field_ref.key),
                expected="a dimension, time_dimension, or measure Ref",
                received=field_ref.kind.value,
            )
        if (
            not isinstance(node.args[1], ast.Name)
            or node.args[1].id not in self._parameter_positions
        ):
            raise SemanticLoadError(
                kind=ErrorKind.BINDING_ALIAS_NOT_DIRECT,
                message=(
                    f"Field ref {field_ref.key!r} must bind to one direct "
                    "expression-body entity parameter."
                ),
                refs=(self._owning_ref.key, field_ref.key),
                expected="ms.bind(field_ref, entity_parameter)",
                received=ast.dump(node, include_attributes=False),
            )
        entity_position = self._parameter_positions[node.args[1].id]
        key = (field_ref.kind, field_ref.path, entity_position)
        if key not in self.index_by_key:
            self.index_by_key[key] = len(self.bindings)
            self.bindings.append(
                ExpressionBindingV1(
                    field_ref=RefPayloadV1.from_ref(field_ref),
                    entity_position=entity_position,
                )
            )
        self.generic_visit(node)


class _NormalizedBody(ast.NodeTransformer):
    def __init__(
        self,
        *,
        symbols: Mapping[str, object],
        parameter_positions: Mapping[str, int],
        binding_indexes: Mapping[tuple[SemanticKind, str, int], int],
        constant_bindings: Mapping[str, EventConstant],
        symbol_tokens: Mapping[str, str] | None = None,
    ) -> None:
        self._symbols = symbols
        self._parameter_positions = parameter_positions
        self._binding_indexes = binding_indexes
        self._constant_bindings = constant_bindings
        self._symbol_tokens = symbol_tokens or {}

    def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.AST:
        node.name = "expression_body"
        node.decorator_list = []
        node.returns = None
        node.type_comment = None
        if node.body and isinstance(node.body[0], ast.Expr):
            value = node.body[0].value
            if isinstance(value, ast.Constant) and isinstance(value.value, str):
                node.body = node.body[1:]
        node.args.defaults = []
        node.args.kw_defaults = []
        self.generic_visit(node)
        return node

    def visit_arg(self, node: ast.arg) -> ast.AST:
        position = self._parameter_positions.get(node.arg)
        if position is not None:
            node.arg = f"entity_{position}"
        node.annotation = None
        node.type_comment = None
        return node

    def visit_Name(self, node: ast.Name) -> ast.AST:
        position = self._parameter_positions.get(node.id)
        if position is not None:
            node.id = f"entity_{position}"
        elif node.id in self._symbol_tokens:
            # Encode the resolved Ibis symbol identity, not the alias spelling:
            # two aliases for one symbol normalize alike, one alias rebound to a
            # different Ibis symbol produces a different normalized body.
            return ast.copy_location(
                ast.Name(id=self._symbol_tokens[node.id], ctx=ast.Load()),
                node,
            )
        elif node.id in self._constant_bindings:
            return ast.copy_location(
                ast.Constant(value=self._constant_bindings[node.id]),
                node,
            )
        return node

    def visit_Call(self, node: ast.Call) -> ast.AST:
        if (
            _is_bind_target(node.func, self._symbols)
            and len(node.args) == 2
            and isinstance(node.args[0], ast.Name)
        ):
            value = self._symbols.get(node.args[0].id)
            argument = node.args[1]
            if type(value) is Ref and isinstance(argument, ast.Name):
                ref = cast("Ref[SemanticKindTag]", value)
                position = self._parameter_positions.get(argument.id)
                if position is not None:
                    index = self._binding_indexes.get((ref.kind, ref.path, position))
                    if index is not None:
                        return ast.copy_location(
                            ast.Call(
                                func=ast.Name(id=f"field_{index}", ctx=ast.Load()),
                                args=[ast.Name(id=f"entity_{position}", ctx=ast.Load())],
                                keywords=[],
                            ),
                            node,
                        )
        self.generic_visit(node)
        return node


def _normalized_body_hash(
    function: ast.FunctionDef,
    *,
    symbols: Mapping[str, object],
    parameter_positions: Mapping[str, int],
    binding_indexes: Mapping[tuple[SemanticKind, str, int], int],
    constant_bindings: Mapping[str, EventConstant],
    symbol_tokens: Mapping[str, str] | None = None,
) -> str:
    normalized = copy.deepcopy(function)
    transformed = _NormalizedBody(
        symbols=symbols,
        parameter_positions=parameter_positions,
        binding_indexes=binding_indexes,
        constant_bindings=constant_bindings,
        symbol_tokens=symbol_tokens,
    ).visit(normalized)
    ast.fix_missing_locations(transformed)
    encoded = ast.dump(transformed, include_attributes=False).encode()
    return f"sha256:{hashlib.sha256(encoded).hexdigest()}"


def _physical_source_columns(
    function: ast.FunctionDef,
    *,
    parameter_positions: Mapping[str, int],
) -> tuple[str, ...]:
    columns: list[str] = []
    parameter_names = frozenset(parameter_positions)
    for node in ast.walk(function):
        column: str | None = None
        if isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name):
            if node.value.id in parameter_names:
                column = node.attr
        elif (
            isinstance(node, ast.Subscript)
            and isinstance(node.value, ast.Name)
            and node.value.id in parameter_names
            and isinstance(node.slice, ast.Constant)
            and isinstance(node.slice.value, str)
        ):
            column = node.slice.value
        if column is not None and column not in columns:
            columns.append(column)
    return tuple(columns)


def _entity_table_source_columns(
    function: ast.FunctionDef,
    *,
    parameter_positions: Mapping[str, int],
) -> tuple[str, ...]:
    """Capture column reads on the injected Entity parameter.

    Method calls such as ``raw.filter(...)`` attribute the method name to the
    parameter; only value reads (attribute/subscript outside a call-func
    position) are source columns.
    """
    parameter_names = frozenset(parameter_positions)
    call_func_nodes: set[int] = set()
    for node in ast.walk(function):
        if isinstance(node, ast.Call):
            call_func_nodes.add(id(node.func))
    columns: list[str] = []
    for node in ast.walk(function):
        if not (
            isinstance(node, (ast.Attribute, ast.Subscript))
            and isinstance(node.value, ast.Name)
            and node.value.id in parameter_names
        ):
            continue
        if id(node) in call_func_nodes:
            continue
        column: str | None = None
        if isinstance(node, ast.Attribute):
            column = node.attr
        elif isinstance(node.slice, ast.Constant) and isinstance(node.slice.value, str):
            column = node.slice.value
        if column is not None and column not in columns:
            columns.append(column)
    return tuple(columns)


# Attribute access on the injected Table that materializes or compiles the
# expression instead of building one. Rejected as an execution escape hatch.
_ENTITY_EXECUTION_ATTRS: frozenset[str] = frozenset(
    {
        "compile",
        "execute",
        "sql",
        "raw_sql",
        "to_arrow",
        "to_csv",
        "to_delta",
        "to_json",
        "to_pandas",
        "to_parquet",
        "to_polars",
        "to_pyarrow",
        "to_torch",
    }
)

# Ibis symbols whose value depends on execution time or backend state. The
# design bans clock reads, random values, and hidden external data from an
# Entity body; aliasing these names does not bypass the ban because the
# check resolves symbol identity, not spelling.
_ENTITY_NONDETERMINISTIC_IBIS_NAMES: frozenset[str] = frozenset(
    {
        "now",
        "today",
        "random",
        "uuid",
        "integer_range",
    }
)

# Deferred instance names on the ibis module admitted in an Entity body.
_ENTITY_DEFERRED_IBIS_NAMES: frozenset[str] = frozenset({"_", "t"})

# Deterministic public Ibis expression symbols admitted in an Entity body:
# the root module's __all__ minus the nondeterministic clock/random symbols.
# Attributes outside __all__ (submodules, options, backends) are rejected so
# body spelling cannot reach execution or connection surfaces.
_ENTITY_IBIS_ADMITTED_NAMES: frozenset[str] = (
    frozenset(getattr(ibis, "__all__", ())) - _ENTITY_NONDETERMINISTIC_IBIS_NAMES
)


def _is_admitted_ibis_symbol(value: object) -> bool:
    """Check one resolved body name against the ibis module by identity.

    A name is admitted when it is the ``ibis`` module itself, the deferred
    ``_``/``t`` instances, or the exact attribute for one admitted public
    Ibis symbol (``ibis.__all__`` minus clock/random names). This covers
    author import aliases regardless of local spelling. Module
    ``__getattr__`` is bypassed so lazy backend imports are never triggered
    by validation. Captured Tables, helper callables, nondeterministic
    clock/random symbols, and unrelated objects are rejected even when
    their names or modules coincide with Ibis types.
    """
    if value is ibis:
        return True
    for name, attr in vars(ibis).items():
        if attr is value:
            return name in _ENTITY_IBIS_ADMITTED_NAMES or name in _ENTITY_DEFERRED_IBIS_NAMES
    return False


def _entity_body_symbol_name(value: object) -> str | None:
    """Return the canonical ``ibis.<name>`` for one admitted symbol, else None.

    Names resolve by identity against the ``ibis`` module attributes, so
    author import aliases and ``ibis.<attr>`` spellings map to one canonical
    identity. Deferred instances and the module itself admit no hash token.
    """
    if value is ibis:
        return None
    for name, attr in vars(ibis).items():
        if attr is value:
            return name if name in _ENTITY_IBIS_ADMITTED_NAMES else None
    return None


def _rejected_ibis_attribute_name(value: object) -> str | None:
    """Return the ``ibis.<name>`` for one resolved-but-rejected ibis attribute.

    Distinguishes a rejected Ibis module attribute (clock, random, internal
    surfaces) from a captured non-Ibis value so the error can name the
    banned symbol identity.
    """
    if value is ibis:
        return None
    for name, attr in vars(ibis).items():
        if attr is value and name not in _ENTITY_DEFERRED_IBIS_NAMES:
            return str(name)
    return None


def _validate_entity_table_closed_symbols(
    function: ast.FunctionDef,
    *,
    symbols: Mapping[str, object],
    parameter_names: frozenset[str],
    owning_key: str,
) -> dict[str, str]:
    """Require every body ``Name`` to be the parameter or one admitted ibis symbol.

    Parameter annotations are type metadata, not runtime expression values,
    and are excluded from the closed-name check. Returns the mapping of body
    local names to canonical ``ibis.<name>`` identity tokens for admitted
    Ibis symbols, used to normalize the body hash by resolved identity.
    """
    symbol_tokens: dict[str, str] = {}
    annotation_nodes: set[int] = {
        id(parameter.annotation)
        for parameter in (
            *function.args.posonlyargs,
            *function.args.args,
            *function.args.kwonlyargs,
        )
        if parameter.annotation is not None
    }
    if function.returns is not None:
        annotation_nodes.add(id(function.returns))
    for node in ast.walk(_entity_table_body_nodes(function)):
        if isinstance(node, ast.Call) and _is_bind_target(node.func, symbols):
            raise SemanticLoadError(
                kind=ErrorKind.INVALID_COMPONENT_BODY,
                message=(
                    f"Entity body {owning_key!r} calls ms.bind, which requires fields "
                    "that do not exist as inputs to their own Entity relation."
                ),
                refs=(owning_key,),
                expected="one Ibis Table expression over the injected Table parameter",
                received=ast.dump(node.func, include_attributes=False),
                hint="Read source columns from the injected parameter instead of ms.bind.",
            )
    import ibis

    body_root = _entity_table_body_nodes(function)
    for node in ast.walk(body_root):
        # Attribute access on the ibis module (ibis.now(), ibis.random()) must
        # resolve to one admitted expression symbol; alias spellings are checked
        # at their Name nodes below.
        if (
            isinstance(node, ast.Attribute)
            and isinstance(node.value, ast.Name)
            and type(symbols.get(node.value.id)) is ModuleType
            and getattr(symbols.get(node.value.id), "__name__", None) == "ibis"
        ):
            attr_value = getattr(symbols[node.value.id], node.attr, None)
            canonical = _entity_body_symbol_name(attr_value)
            if canonical is None:
                rejected = _rejected_ibis_attribute_name(attr_value)
                raise SemanticLoadError(
                    kind=ErrorKind.COMPILE_ERROR,
                    message=(
                        f"Entity body {owning_key!r} uses Ibis symbol "
                        f"ibis.{rejected or node.attr}, which is not admitted. Clock "
                        "reads, random values, generated identifiers, and non-expression "
                        "module attributes are not allowed in an Entity body."
                    ),
                    refs=(owning_key,),
                    expected=(
                        "deterministic expression-building symbols such as ibis.literal, "
                        "ibis.window, or ibis.row_number"
                    ),
                    received=f"ibis.{node.attr}",
                    hint=(
                        "Move the nondeterministic value into an explicit authored column "
                        "or pass it as data, not as a body symbol."
                    ),
                )
            symbol_tokens.setdefault(f"{node.value.id}.{node.attr}", f"ibis:{canonical}")
    for node in ast.walk(body_root):
        if not isinstance(node, ast.Name):
            continue
        if node.id in parameter_names or id(node) in annotation_nodes:
            continue
        value = symbols.get(node.id)
        if _is_admitted_ibis_symbol(value):
            # The ibis module may only be used for admitted attribute access;
            # each attribute spelling is checked at its Attribute node above.
            if value is ibis:
                continue
            canonical = _entity_body_symbol_name(value)
            if canonical is not None:
                symbol_tokens[node.id] = f"ibis:{canonical}"
            continue
        if value is None and node.id not in symbols:
            raise SemanticLoadError(
                kind=ErrorKind.COMPILE_ERROR,
                message=(f"Entity body {owning_key!r} uses unresolved external name {node.id!r}."),
                refs=(owning_key,),
                expected=("the injected Table parameter, ibis module symbols, or inline literals"),
                received=node.id,
                hint=(
                    "Derive the expression from the injected parameter and import "
                    "ibis symbols in the defining module."
                ),
            )
        rejected_ibis_name = _rejected_ibis_attribute_name(value)
        if rejected_ibis_name is not None:
            raise SemanticLoadError(
                kind=ErrorKind.COMPILE_ERROR,
                message=(
                    f"Entity body {owning_key!r} uses Ibis symbol ibis.{rejected_ibis_name} "
                    f"(bound as {node.id!r}), which is not admitted. Clock reads, random "
                    "values, generated identifiers, and non-expression module attributes "
                    "are not allowed in an Entity body."
                ),
                refs=(owning_key,),
                expected=(
                    "deterministic expression-building symbols such as ibis.literal, "
                    "ibis.window, or ibis.row_number"
                ),
                received=f"ibis.{rejected_ibis_name}",
                hint=(
                    "Move the nondeterministic value into an explicit authored column or "
                    "pass it as data, not as a body symbol."
                ),
            )
        raise SemanticLoadError(
            kind=ErrorKind.COMPILE_ERROR,
            message=(
                f"Entity body {owning_key!r} uses external value {node.id!r} of type "
                f"{type(value).__name__}, which is not a supported Ibis expression symbol."
            ),
            refs=(owning_key,),
            expected=("the injected Table parameter, ibis module symbols, or inline literals"),
            received=f"{node.id}={type(value).__name__}",
            hint=(
                "Replace captured values or helper callables with an inline Ibis "
                "expression over the injected parameter."
            ),
        )
    return symbol_tokens


def _entity_table_body_nodes(function: ast.FunctionDef) -> ast.AST:
    """Return the single return expression with decorators and docstring excluded.

    Falls back to the whole function when the return-count contract is not yet
    established; shape validation then reports the structural failure.
    """
    returns = [node for node in ast.walk(function) if isinstance(node, ast.Return)]
    if len(returns) == 1 and returns[0].value is not None:
        return returns[0].value
    return function


def _validate_entity_table_no_execution(
    function: ast.FunctionDef,
    *,
    owning_key: str,
) -> None:
    """Reject materialization or SQL escape attributes inside the Entity body."""
    for node in ast.walk(_entity_table_body_nodes(function)):
        if isinstance(node, ast.Attribute) and node.attr in _ENTITY_EXECUTION_ATTRS:
            raise SemanticLoadError(
                kind=ErrorKind.SQL_ESCAPE_HATCH,
                message=(
                    f"Entity body {owning_key!r} uses .{node.attr}, which materializes "
                    "or rewrites the expression. Return the Ibis Table expression itself."
                ),
                refs=(owning_key,),
                expected="one returned Ibis Table expression over the injected parameter",
                received=f".{node.attr}",
                hint=(
                    "Marivo materializes the Entity relation through its resolver; "
                    "remove the execution or SQL call from the body."
                ),
            )


def compile_expression_body(
    fn: Callable[..., object],
    *,
    owning_ref: Ref[SemanticKindTag],
    ordered_entity_refs: tuple[Ref[EntityKind], ...] = (),
    body_kind: Literal[
        "dimension",
        "time_dimension",
        "measure",
        "metric",
        "event",
        "entity_table",
    ]
    | None = None,
) -> ExpressionBody:
    """Validate one decorator body and capture exact semantic field bindings."""
    owning = _require_exact_ref(owning_ref, parameter="owning_ref")
    for entity_ref in ordered_entity_refs:
        _require_exact_ref(
            entity_ref,
            allowed=frozenset({SemanticKind.ENTITY}),
            parameter="ordered_entity_refs item",
        )
    body_kind_by_ref: dict[
        SemanticKind,
        Literal["dimension", "time_dimension", "measure", "metric", "event"],
    ] = {
        SemanticKind.DIMENSION: "dimension",
        SemanticKind.TIME_DIMENSION: "time_dimension",
        SemanticKind.MEASURE: "measure",
        SemanticKind.METRIC: "metric",
        SemanticKind.EVENT: "event",
    }
    resolved_kind = body_kind if body_kind is not None else body_kind_by_ref.get(owning.kind)
    if resolved_kind is None or (
        body_kind == "entity_table" and owning.kind is not SemanticKind.ENTITY
    ):
        raise SemanticLoadError(
            kind=ErrorKind.INVALID_BINDING_REF,
            message=f"Ref {owning.key!r} cannot own an expression body.",
            refs=(owning.key,),
            expected=(
                "dimension, time_dimension, measure, metric, or event"
                if body_kind != "entity_table"
                else "entity"
            ),
            received=owning.kind.value,
        )
    body_kind = resolved_kind
    if body_kind == "entity_table":
        _, function = _load_function_ast(fn)
        parameters = (*function.args.posonlyargs, *function.args.args)
        parameter_positions = {parameter.arg: index for index, parameter in enumerate(parameters)}
        symbols = _resolved_symbols(fn)
        validate_entity_table_body_ast(fn)
        _validate_entity_table_no_execution(function, owning_key=owning.key)
        symbol_tokens = _validate_entity_table_closed_symbols(
            function,
            symbols=symbols,
            parameter_names=frozenset(parameter_positions),
            owning_key=owning.key,
        )
        parameter = parameters[0].arg if parameters else ""
        entities = {parameter: cast("Ref[EntityKind]", owning)} if parameter else {}
        description = describe_expression(
            function,
            entities=entities,
            bindings={},
        )
        return ExpressionBody(
            description=description,
            callable=fn,
            body_ast_hash=_normalized_body_hash(
                function,
                symbols=symbols,
                parameter_positions=parameter_positions,
                binding_indexes={},
                constant_bindings={},
                symbol_tokens=symbol_tokens,
            ),
            parameter_count=1,
            bindings=(),
            source_columns=_entity_table_source_columns(
                function,
                parameter_positions=parameter_positions,
            ),
        )
    if body_kind == "event":
        validate_event_body_ast(fn)
    else:
        validate_metric_body_ast(fn, "base", body_kind=body_kind)
    _, function = _load_function_ast(fn)
    if function.args.vararg is not None or function.args.kwarg is not None:
        raise SemanticLoadError(
            kind=ErrorKind.COMPILE_ERROR,
            message=f"Expression body {owning.key!r} cannot use variadic parameters.",
            refs=(owning.key,),
            expected="only direct entity parameters",
            received="variadic parameter",
        )
    if function.args.kwonlyargs:
        raise SemanticLoadError(
            kind=ErrorKind.COMPILE_ERROR,
            message=f"Expression body {owning.key!r} cannot use keyword-only parameters.",
            refs=(owning.key,),
            expected="only direct positional entity parameters",
            received="keyword-only parameter",
        )
    parameters = (*function.args.posonlyargs, *function.args.args)
    if len(parameters) != len(ordered_entity_refs):
        raise SemanticLoadError(
            kind=ErrorKind.COMPILE_ERROR,
            message=f"Expression body {owning.key!r} has the wrong entity arity.",
            refs=(owning.key,),
            expected=str(len(ordered_entity_refs)),
            received=str(len(parameters)),
        )
    parameter_positions = {parameter.arg: index for index, parameter in enumerate(parameters)}
    symbols = _resolved_symbols(fn)
    event_constants = (
        _event_closed_constants(
            function,
            symbols=symbols,
            owning_ref=owning,
        )
        if body_kind == "event"
        else {}
    )
    collector = _BindingCollector(
        symbols=symbols,
        parameter_positions=parameter_positions,
        owning_ref=owning,
    )
    collector.visit(function)
    description_bindings: dict[int, tuple[Ref[FieldKind], Ref[EntityKind]]] = {}
    for node in ast.walk(function):
        if (
            isinstance(node, ast.Call)
            and _is_bind_target(node.func, symbols)
            and (
                len(node.args) == 2
                and isinstance(node.args[0], ast.Name)
                and isinstance(node.args[1], ast.Name)
            )
        ):
            position = parameter_positions[node.args[1].id]
            for binding in collector.bindings:
                if (
                    binding.to_ref() == symbols[node.args[0].id]
                    and binding.entity_position == position
                ):
                    description_bindings[id(node)] = (
                        binding.to_ref(),
                        ordered_entity_refs[position],
                    )
    description = describe_expression(
        function,
        entities={
            name: ordered_entity_refs[position] for name, position in parameter_positions.items()
        },
        bindings=description_bindings,
    )
    if body_kind == "event":
        invalid = tuple(
            binding.to_ref()
            for binding in collector.bindings
            if binding.field_ref.kind is not SemanticKind.DIMENSION
        )
        if invalid:
            raise SemanticLoadError(
                kind=ErrorKind.INVALID_EVENT_PREDICATE,
                message="Event predicates may bind categorical Dimensions only.",
                refs=(owning.key, *(ref.key for ref in invalid)),
                expected="Ref[dimension] bindings on the Event source Entity",
                received=", ".join(ref.kind.value for ref in invalid),
            )
    return ExpressionBody(
        description=description,
        callable=_freeze_event_callable(fn, event_constants) if body_kind == "event" else fn,
        body_ast_hash=_normalized_body_hash(
            function,
            symbols=symbols,
            parameter_positions=parameter_positions,
            binding_indexes=collector.index_by_key,
            constant_bindings=event_constants,
        ),
        parameter_count=len(parameters),
        bindings=tuple(collector.bindings),
        source_columns=_physical_source_columns(
            function,
            parameter_positions=parameter_positions,
        ),
    )


def _validate_body_result(
    result: object,
    *,
    owning_ref: Ref[SemanticKindTag],
    table_result: bool = False,
) -> ir.Value:
    if not isinstance(result, ir.Value):
        if table_result and isinstance(result, ibis.Table):
            return cast("ir.Value", result)
        raise _binding_runtime_error(
            ErrorKind.BINDING_RESULT_INVALID,
            f"Expression body {owning_ref.key!r} must return one Ibis value.",
            refs=(owning_ref,),
            expected="ibis.expr.types.Table" if table_result else "ibis.expr.types.Value",
            received=type(result).__name__,
        )
    if table_result and not isinstance(result, ir.Table):
        raise _binding_runtime_error(
            ErrorKind.BINDING_RESULT_INVALID,
            f"Expression body {owning_ref.key!r} must return one Ibis Table relation.",
            refs=(owning_ref,),
            expected="ibis.expr.types.Table",
            received=type(result).__name__,
        )
    return result


def _call_body(
    body: ExpressionBody,
    aliases: tuple[ir.Table, ...],
    *,
    owning_ref: Ref[SemanticKindTag],
    table_result: bool = False,
) -> ir.Value:
    try:
        result = body.callable(*aliases)
    except SemanticError:
        raise
    except Exception as exc:
        hint = (
            " Import every name used by the body in its defining module."
            if isinstance(exc, NameError)
            else ""
        )
        raise _binding_runtime_error(
            ErrorKind.MATERIALIZE_FAILED,
            f"Expression body {owning_ref.key!r} failed: {exc}.{hint}",
            refs=(owning_ref,),
            received=type(exc).__name__,
        ) from exc
    return _validate_body_result(result, owning_ref=owning_ref, table_result=table_result)


def evaluate_expression_body(
    *,
    catalog_definition_fingerprint: str,
    expression_sidecar: CompiledExpressionSidecar,
    owning_ref: Ref[SemanticKindTag],
    body: ExpressionBody,
    entity_refs: tuple[Ref[EntityKind], ...],
    aliases: tuple[ir.Table, ...],
    table_result: bool = False,
) -> ir.Value:
    """Evaluate one root body inside a fresh task-local binding context.

    With ``table_result`` the body must return one Ibis Table relation; the
    default accepts any Ibis value (scalars and columns for fields and
    metrics).
    """
    owning = _require_exact_ref(owning_ref, parameter="owning_ref")
    if body.parameter_count != len(entity_refs) or len(entity_refs) != len(aliases):
        raise _binding_runtime_error(
            ErrorKind.BINDING_ENTITY_MISMATCH,
            f"Expression body {owning.key!r} received the wrong entity arity.",
            refs=(owning,),
            expected=str(body.parameter_count),
            received=str(len(aliases)),
        )
    frame = ExpressionBodyFrame(
        owning_ref=owning,
        ordered_entity_refs=entity_refs,
        ordered_entity_aliases=aliases,
        declared_bindings=body.bindings,
    )
    context = ExpressionBindingContext(
        catalog_definition_fingerprint=catalog_definition_fingerprint,
        expression_sidecar=expression_sidecar,
        body_frames=(frame,),
    )
    token = _EXPRESSION_BINDING_CONTEXT.set(context)
    try:
        return _call_body(body, aliases, owning_ref=owning, table_result=table_result)
    finally:
        _EXPRESSION_BINDING_CONTEXT.reset(token)


[docs] def bind(field: Ref[FieldKind], entity_alias: ir.Table, /) -> ir.Value: """Apply a semantic field ref to an entity alias in an expression body. Parameters ---------- field: Exact dimension, time-dimension, or measure ref declared in the loaded semantic project. entity_alias: Direct entity parameter of the active decorated expression body. Returns ------- ibis.expr.types.Value The referenced field expression evaluated on ``entity_alias``. Example ------- >>> @ms.metric(entities=[orders], additivity="additive") ... def revenue(orders): ... return ms.bind(amount, orders).sum() Constraints ----------- Only valid inside a loaded semantic expression body. The field must belong to the bound entity and must be captured as a direct ``ms.bind`` argument. """ ref = field if type(ref) is not Ref: raise _binding_runtime_error( ErrorKind.INVALID_BINDING_REF, "Semantic field application requires an exact Ref value.", expected="Ref[dimension | time_dimension | measure]", received=type(ref).__name__, ) candidate = cast("Ref[SemanticKindTag]", ref) if candidate.kind not in _FIELD_KINDS: raise _binding_runtime_error( ErrorKind.INVALID_BINDING_REF, f"Ref {candidate.key!r} is not a bindable semantic field ref.", refs=(candidate,), expected="dimension, time_dimension, or measure Ref", received=candidate.kind.value, ) field_ref = candidate context = _EXPRESSION_BINDING_CONTEXT.get() if context is None: raise _binding_runtime_error( ErrorKind.BINDING_CONTEXT_MISSING, f"Field ref {field_ref.key!r} has no active expression binding context.", refs=(field_ref,), expected="evaluation through the active catalog materializer", received="no active context", ) frame = context.body_frames[-1] positions = [ index for index, candidate in enumerate(frame.ordered_entity_aliases) if candidate is entity_alias ] if not positions: raise _binding_runtime_error( ErrorKind.BINDING_ALIAS_NOT_DIRECT, f"Field ref {field_ref.key!r} did not receive a direct body entity alias.", refs=(frame.owning_ref, field_ref), expected="the exact entity parameter object", received=type(entity_alias).__name__, ) if len(positions) != 1: raise _binding_runtime_error( ErrorKind.BINDING_ALIAS_AMBIGUOUS, f"Field ref {field_ref.key!r} received an alias used at multiple positions.", refs=(frame.owning_ref, field_ref), expected="one unique entity parameter position", received=str(positions), ) position = positions[0] declared = any( binding.entity_position == position and binding.field_ref.kind == field_ref.kind and binding.field_ref.path == field_ref.path for binding in frame.declared_bindings ) if not declared: raise _binding_runtime_error( ErrorKind.BINDING_NOT_DECLARED, f"Field ref {field_ref.key!r} was not declared for entity position {position}.", refs=(frame.owning_ref, field_ref), expected="a binding captured from the expression body AST", received=f"position {position}", ) sidecar = context.expression_sidecar if field_ref not in sidecar.catalog_refs: raise _binding_runtime_error( ErrorKind.BINDING_TARGET_MISSING, f"Field ref {field_ref.key!r} is absent from the active catalog.", refs=(frame.owning_ref, field_ref), received=context.catalog_definition_fingerprint, ) owner = sidecar.field_owners.get(cast("Ref[FieldKind]", field_ref)) expected_owner = frame.ordered_entity_refs[position] if owner != expected_owner: raise _binding_runtime_error( ErrorKind.BINDING_ENTITY_MISMATCH, f"Field ref {field_ref.key!r} does not belong to entity position {position}.", refs=(frame.owning_ref, field_ref, expected_owner), expected=expected_owner.key, received=owner.key if owner is not None else "missing owner", ) body = sidecar.bodies.get(field_ref) if body is None: raise _binding_runtime_error( ErrorKind.BINDING_TARGET_MISSING, f"Field ref {field_ref.key!r} has no compiled expression body.", refs=(frame.owning_ref, field_ref), ) if any(active.owning_ref == field_ref for active in context.body_frames): raise _binding_runtime_error( ErrorKind.BINDING_CYCLE, f"Expression binding cycle reached {field_ref.key!r}.", refs=(*(active.owning_ref for active in context.body_frames), field_ref), ) if body.parameter_count != 1: raise _binding_runtime_error( ErrorKind.BINDING_ENTITY_MISMATCH, f"Field ref {field_ref.key!r} must have a one-entity expression body.", refs=(field_ref,), expected="1", received=str(body.parameter_count), ) nested_frame = ExpressionBodyFrame( owning_ref=field_ref, ordered_entity_refs=(expected_owner,), ordered_entity_aliases=(entity_alias,), declared_bindings=body.bindings, ) nested_context = ExpressionBindingContext( catalog_definition_fingerprint=context.catalog_definition_fingerprint, expression_sidecar=sidecar, body_frames=(*context.body_frames, nested_frame), ) token = _EXPRESSION_BINDING_CONTEXT.set(nested_context) try: return _call_body(body, (entity_alias,), owning_ref=field_ref) finally: _EXPRESSION_BINDING_CONTEXT.reset(token)