from __future__ import annotations

from dataclasses import dataclass, field

from relationalai.util.structures import OrderedSet

from ...metamodel.builtins import (
    builtins as builtin_registry,
    is_entity_type,
    is_metatype,
    is_primitive,
    is_value_type,
)
from ...metamodel.metamodel import (
    Aggregate,
    Construct,
    Exists,
    Literal,
    Logical,
    Lookup,
    Loop,
    Match,
    Not,
    Require,
    Sequence,
    Task,
    TypeNode,
    Union,
    Update,
    Value,
    Var,
)
from ...metamodel.metamodel_analyzer import VarFinder
from .model_analysis import ModelAnalysis, TableKind
from ...metamodel.rewriter import Rewriter
from .coalesce import SQL2CoalescedArtifacts
from .task_semantics import bindable_var_ids, produced_var_ids, task_var_ids
from .table_shape import SQL2TableShapeSidecar


@dataclass(frozen=True)
class SQL2AggregateShape:
    aggregation_relation_id: int
    group_shape: tuple[str, ...]
    projection_shape: tuple[str, ...]
    non_output_arg_shape: tuple[str, ...]


@dataclass(frozen=True)
class SQL2AggregateInstance:
    shape: SQL2AggregateShape
    group_id: str
    channel_id: int
    aggregate_task_id: int
    output_var_id: int | None


@dataclass(frozen=True)
class SQL2NormalizedChannel:
    channel_id: int
    node_id: int
    target_relation_id: int
    target_relation_name: str
    output_columns: tuple[str, ...]
    key_var_ids: tuple[int, ...]
    rule: Task
    update_task: Update
    body_rule: Logical
    optional_projection_var_ids_by_task_id: tuple[tuple[int, tuple[int, ...]], ...]
    branch_projection_var_ids_by_task_id: tuple[tuple[int, tuple[int, ...]], ...]
    union_mode_by_task_id: tuple[tuple[int, str], ...]
    future_var_ids_by_task_id: tuple[tuple[int, tuple[int, ...]], ...]
    read_table_names: tuple[str, ...] = field(default_factory=tuple)
    aggregate_instances: tuple[SQL2AggregateInstance, ...] = field(default_factory=tuple)


@dataclass(frozen=True)
class SQL2NormalizedGroup:
    group_id: str
    stratum_id: int
    table_name: str
    key_var_ids: tuple[int, ...]
    output_columns: tuple[str, ...]
    key_columns: tuple[str, ...]
    channel_combine_plan: "SQL2ChannelCombinePlan"
    read_table_names: tuple[str, ...]
    node_ids: tuple[int, ...]
    channels: tuple[SQL2NormalizedChannel, ...]
    normalized_rule: Task | None


@dataclass(frozen=True)
class SQL2ChannelCombinePlan:
    strategy: str
    output_columns: tuple[str, ...]
    key_columns: tuple[str, ...]
    aggregate_columns: tuple[str, ...]
    collapsed_channel_groups: tuple[tuple[int, ...], ...]


@dataclass(frozen=True)
class SQL2NormalizedArtifacts:
    groups: tuple[SQL2NormalizedGroup, ...]
    groups_by_stratum: dict[int, tuple[SQL2NormalizedGroup, ...]]
    aggregate_instances_by_shape: dict[SQL2AggregateShape, tuple[SQL2AggregateInstance, ...]]
    relation_key_preserving_by_relation_id: tuple[tuple[int, bool], ...]
    relation_direct_projection_safe_by_relation_id: tuple[tuple[int, bool], ...]
    relation_output_anchor_safe_by_relation_id: tuple[tuple[int, bool], ...]
    table_plans: tuple["SQL2NormalizedTablePlan", ...]
    preseed_tables: tuple["SQL2PreseedTablePlan", ...]
    program_by_stratum: dict[int, Task]
    program_root: Task | None


@dataclass(frozen=True)
class SQL2NormalizedTablePlan:
    table_name: str
    output_columns: tuple[str, ...]
    key_columns: tuple[str, ...]


@dataclass(frozen=True)
class SQL2PreseedTablePlan:
    table_name: str
    output_columns: tuple[str, ...]


class _VarSubstituter(Rewriter):
    def __init__(self, mapping: dict[Var, Var]):
        super().__init__()
        self._mapping = mapping

    def var(self, node: Var):
        return self._mapping.get(node, node)

    def rewrite(self, node):
        return self(node)


class _SQL2TaskNormalizer(Rewriter):
    def __init__(self):
        super().__init__()
        self._finder = VarFinder()
        self._aggregate_depth = 0
        self._container_depth = 0

    def rewrite(self, node):
        return self(node)

    def enter_aggregate(self, node: Aggregate):
        self._aggregate_depth += 1
        return None

    def aggregate(self, node: Aggregate):
        self._aggregate_depth -= 1
        return self._rewrite_aggregate_body(node)

    def enter_union(self, node: Union):
        self._container_depth += 1
        return None

    def union(self, node: Union):
        self._container_depth -= 1
        return self._normalize_branch_union(node=node, label="Union")

    def _normalize_branch_union(self, *, node: Union, label: str) -> Task:
        tasks = tuple(
            task
            if isinstance(task, Logical)
            else Logical(body=(task,), scope=True)
            for task in node.tasks
        )
        output_vars = tuple(value for value in node.outputs if isinstance(value, Var))
        if len(output_vars) != len(node.outputs):
            raise ValueError(f"{label} outputs must be variables")
        if output_vars:
            output_var_ids = {value.id for value in output_vars}
            for branch in tasks:
                branch_var_ids = task_var_ids(branch)
                missing = sorted(output_var_ids.difference(branch_var_ids))
                if missing:
                    raise ValueError(
                        f"{label} branch must reference all {label.lower()} outputs; "
                        f"missing var ids: {missing}"
                    )
        if not node.outputs and len(tasks) == 1:
            return tasks[0]
        return node.mut(tasks=tasks, outputs=output_vars)

    def enter_match(self, node: Match):
        self._container_depth += 1
        return None

    def match(self, node: Match):
        self._container_depth -= 1
        tasks = tuple(
            task
            if isinstance(task, Logical)
            else Logical(body=(task,), scope=True)
            for task in node.tasks
        )
        output_vars = tuple(value for value in node.outputs if isinstance(value, Var))
        if len(output_vars) != len(node.outputs):
            raise ValueError("Match outputs must be variables")
        # Match with no outputs and no updates is predicate-only semantics,
        # which is equivalent to a disjunction over branches.
        if not output_vars and all(not self._task_contains_update(task) for task in tasks):
            return self._normalize_branch_union(
                node=Union(
                    tasks=tasks,
                    outputs=(),
                    optional=node.optional,
                    scope=node.scope,
                ),
                label="Match",
            )
        guarded_tasks: list[Task] = []
        prior_tasks: list[Task] = []
        for task in tasks:
            if not prior_tasks:
                guarded_tasks.append(task)
            else:
                guarded_prior_tasks: tuple[Task, ...] = tuple(prior_tasks)
                if output_vars:
                    output_var_ids = {value.id for value in output_vars}
                    fresh_by_id: dict[int, Var] = {}
                    substitution: dict[Var, Var] = {}
                    for prior in prior_tasks:
                        for var in self._task_vars(prior):
                            if var.id not in output_var_ids:
                                continue
                            fresh = fresh_by_id.get(var.id)
                            if fresh is None:
                                fresh = Var(type=var.type, name=var.name)
                                fresh_by_id[var.id] = fresh
                            substitution[var] = fresh
                    if substitution:
                        substituter = _VarSubstituter(substitution)
                        guarded_prior_tasks = tuple(substituter.rewrite(prior) for prior in prior_tasks)

                prior_union = Union(tasks=guarded_prior_tasks, outputs=(), scope=True)
                guard = Not(task=Exists(task=prior_union, scope=True), scope=True)
                guarded_tasks.append(Logical(body=(guard, task), scope=True))
            prior_tasks.append(task)
        return self._normalize_branch_union(
            node=Union(
                tasks=tuple(guarded_tasks),
                outputs=output_vars,
                optional=node.optional,
                scope=node.scope,
            ),
            label="Match",
        )

    def _task_contains_update(self, task: Task) -> bool:
        if isinstance(task, Update):
            return True
        if isinstance(task, Logical):
            return any(self._task_contains_update(child) for child in task.body)
        if isinstance(task, (Not, Exists)):
            return self._task_contains_update(task.task)
        if isinstance(task, Aggregate):
            return self._task_contains_update(task.body)
        if isinstance(task, (Union, Match, Sequence)):
            return any(self._task_contains_update(child) for child in task.tasks)
        if isinstance(task, Loop):
            return self._task_contains_update(task.body)
        if isinstance(task, Require):
            if self._task_contains_update(task.domain):
                return True
            if self._task_contains_update(task.check):
                return True
            return task.error is not None and self._task_contains_update(task.error)
        return False

    def enter_not(self, node: Not):
        self._container_depth += 1
        return None

    def not_(self, node: Not):
        self._container_depth -= 1
        return node

    def enter_exists(self, node: Exists):
        self._container_depth += 1
        return None

    def exists(self, node: Exists):
        self._container_depth -= 1
        return node

    def _task_vars(self, task: Task) -> OrderedSet[Var]:
        return self._finder.find_vars(task)

    def _aggregate_input_tuple_vars(self, agg: Aggregate) -> set[Var]:
        return {
            var
            for var in [*agg.group, *agg.projection, *agg.args[:-1]]
            if isinstance(var, Var)
        }

    def _bound_before(self, tasks: list[Task]) -> list[set[Var]]:
        bound: set[Var] = set()
        out: list[set[Var]] = []
        for task in tasks:
            out.append(set(bound))
            if isinstance(task, Lookup):
                for arg in task.args:
                    if isinstance(arg, Var):
                        bound.add(arg)
            elif isinstance(task, Aggregate) and task.args and isinstance(task.args[-1], Var):
                bound.add(task.args[-1])
        return out

    def _wrap_dead_lookup_branches_in_exists(
        self,
        tasks: list[Task],
        *,
        live_seed: set[Var] | None = None,
    ) -> tuple[list[Task], bool]:
        if not tasks:
            return tasks, False
        if any(isinstance(task, (Aggregate, Exists, Match, Union, Not)) for task in tasks):
            return tasks, False

        task_vars = [self._task_vars(task) for task in tasks]
        bound_before = self._bound_before(tasks)
        live: OrderedSet[Var] = OrderedSet(live_seed or [])
        for task, vars_used in zip(tasks, task_vars):
            if not isinstance(task, Lookup):
                live.update(vars_used)

        new_rev: list[Task] = []
        pending_lookups: list[Lookup] = []
        pending_quantified: OrderedSet[Var] = OrderedSet()

        def flush_pending():
            if not pending_lookups:
                return
            exists_task = Exists(
                vars=tuple(pending_quantified),
                task=Logical(body=tuple(reversed(pending_lookups)), scope=True),
                scope=True,
            )
            new_rev.append(exists_task)
            pending_lookups.clear()
            pending_quantified.clear()

        for ix in range(len(tasks) - 1, -1, -1):
            task = tasks[ix]
            vars_used = task_vars[ix]
            if isinstance(task, Lookup) and task.relation.fields and all(not field.input for field in task.relation.fields):
                new_vars = OrderedSet(var for var in vars_used if var not in live)
                live_vars = OrderedSet(var for var in vars_used if var in live)
                if new_vars and all(var in bound_before[ix] for var in live_vars):
                    pending_lookups.append(task)
                    pending_quantified.update(new_vars)
                    continue
            flush_pending()
            new_rev.append(task)
            live.update(vars_used)

        flush_pending()
        new_body = list(reversed(new_rev))
        return new_body, tuple(new_body) != tuple(tasks)

    def _rewrite_aggregate_body(self, agg: Aggregate) -> Aggregate:
        if not isinstance(agg.body, Logical) or not agg.body.body or agg.body.optional:
            return agg
        required_vars = self._aggregate_input_tuple_vars(agg)
        rewritten_tasks, changed = self._wrap_dead_lookup_branches_in_exists(
            list(agg.body.body),
            live_seed=required_vars,
        )
        if not changed:
            return agg
        return agg.mut(body=agg.body.mut(body=tuple(rewritten_tasks)))

    def _is_equality_lookup(self, task: Task) -> bool:
        return isinstance(task, Lookup) and (
            task.relation is builtin_registry.core.eq or task.relation.name in {"=", "=="}
        )

    def _is_non_equality_predicate_lookup(self, task: Task) -> bool:
        if not isinstance(task, Lookup):
            return False
        if self._is_equality_lookup(task):
            return False
        return task.relation in {
            builtin_registry.core.neq,
            builtin_registry.core.gt,
            builtin_registry.core.gte,
            builtin_registry.core.lt,
            builtin_registry.core.lte,
        } or task.relation.name in {"!=", ">", ">=", "<", "<="}

    def _lookup_matches_with_local_binding(
        self,
        *,
        inner_lookup: Lookup,
        outer_lookup: Lookup,
        outer_vars: set[Var],
        local_bindings: dict[Var, Value],
    ) -> tuple[bool, dict[Var, Value]]:
        if inner_lookup.relation is not outer_lookup.relation:
            return False, local_bindings
        if len(inner_lookup.args) != len(outer_lookup.args):
            return False, local_bindings

        next_bindings = dict(local_bindings)
        for inner_arg, outer_arg in zip(inner_lookup.args, outer_lookup.args):
            if isinstance(inner_arg, Var):
                if inner_arg in outer_vars:
                    if inner_arg != outer_arg:
                        return False, local_bindings
                    continue
                bound = next_bindings.get(inner_arg)
                if bound is None:
                    next_bindings[inner_arg] = outer_arg
                    continue
                if bound != outer_arg:
                    return False, local_bindings
                continue
            if inner_arg != outer_arg:
                return False, local_bindings
        return True, next_bindings

    def _correlation_substitutions(
        self,
        *,
        tasks: tuple[Task, ...],
        outer_vars: set[Var],
    ) -> dict[Var, Var]:
        mapping: dict[Var, Var] = {}

        def bind(inner_var: Var, outer_var: Var) -> int:
            bound = mapping.get(inner_var)
            if bound is None:
                mapping[inner_var] = outer_var
                return 1
            if bound != outer_var:
                return -1
            return 0

        changed = True
        while changed:
            changed = False
            for task in tasks:
                if not self._is_equality_lookup(task):
                    continue
                assert isinstance(task, Lookup)
                if len(task.args) != 2:
                    continue
                left, right = task.args
                if not isinstance(left, Var) or not isinstance(right, Var):
                    continue
                left_outer = left in outer_vars
                right_outer = right in outer_vars
                if left_outer and not right_outer:
                    result = bind(right, left)
                    if result < 0:
                        return {}
                    changed = changed or result > 0
                    continue
                if right_outer and not left_outer:
                    result = bind(left, right)
                    if result < 0:
                        return {}
                    changed = changed or result > 0
                    continue
                left_bound = mapping.get(left)
                right_bound = mapping.get(right)
                if left_bound is not None and right_outer:
                    result = bind(left, right)
                    if result < 0:
                        return {}
                    changed = changed or result > 0
                elif right_bound is not None and left_outer:
                    result = bind(right, left)
                    if result < 0:
                        return {}
                    changed = changed or result > 0
                elif left_bound is not None and right_bound is None and right not in outer_vars:
                    result = bind(right, left_bound)
                    if result < 0:
                        return {}
                    changed = changed or result > 0
                elif right_bound is not None and left_bound is None and left not in outer_vars:
                    result = bind(left, right_bound)
                    if result < 0:
                        return {}
                    changed = changed or result > 0
                elif left_bound is not None and right_bound is not None and left_bound != right_bound:
                    return {}
        return mapping

    def _derive_protected_local_vars(
        self,
        *,
        tasks: tuple[Task, ...],
        outer_vars: set[Var],
    ) -> set[Var]:
        protected: set[Var] = set()
        for task in tasks:
            if not isinstance(task, Lookup):
                continue
            if not self._is_non_equality_predicate_lookup(task):
                continue
            for arg in task.args:
                if isinstance(arg, Var) and arg not in outer_vars:
                    protected.add(arg)

        changed = True
        while changed:
            changed = False
            for task in tasks:
                if not isinstance(task, Lookup):
                    continue
                local_vars = [
                    arg for arg in task.args
                    if isinstance(arg, Var) and arg not in outer_vars
                ]
                if not local_vars:
                    continue
                if not any(var in protected for var in local_vars):
                    continue
                for var in local_vars:
                    if var in protected:
                        continue
                    protected.add(var)
                    changed = True
        return protected

    def _outer_var_canonical_map(self, *, outer_tasks: tuple[Task, ...]) -> dict[Var, Var]:
        outer_vars: set[Var] = set()
        for task in outer_tasks:
            outer_vars.update(self._task_vars(task))
        if not outer_vars:
            return {}

        parent: dict[Var, Var] = {var: var for var in outer_vars}

        def find(var: Var) -> Var:
            cur = parent.get(var, var)
            while parent.get(cur, cur) != cur:
                cur = parent[cur]
            root = cur
            cur = var
            while parent.get(cur, cur) != root:
                nxt = parent[cur]
                parent[cur] = root
                cur = nxt
            return root

        def union(left: Var, right: Var) -> None:
            left_root = find(left)
            right_root = find(right)
            if left_root == right_root:
                return
            if left_root.id <= right_root.id:
                parent[right_root] = left_root
            else:
                parent[left_root] = right_root

        for task in outer_tasks:
            if not self._is_equality_lookup(task):
                continue
            assert isinstance(task, Lookup)
            if len(task.args) != 2:
                continue
            left, right = task.args
            if isinstance(left, Var) and isinstance(right, Var):
                union(left, right)

        return {var: find(var) for var in outer_vars}

    def _rewrite_correlated_lookup_logical(
        self,
        *,
        logical: Logical,
        outer_tasks: tuple[Task, ...],
    ) -> tuple[Logical, bool]:
        if logical.optional or not logical.body:
            return logical, False
        if any(not isinstance(task, Lookup) for task in logical.body):
            return logical, False

        outer_vars = set()
        outer_lookups: list[Lookup] = []
        for task in outer_tasks:
            outer_vars.update(var for var in self._task_vars(task) if isinstance(var, Var))
            if isinstance(task, Lookup):
                outer_lookups.append(task)
        if not outer_vars or not outer_lookups:
            return logical, False

        canonical_outer_by_var = self._outer_var_canonical_map(outer_tasks=outer_tasks)
        outer_substitution: dict[Var, Var] = {
            var: canonical
            for var, canonical in canonical_outer_by_var.items()
            if var != canonical
        }
        outer_vars_for_correlation: set[Var] = set(outer_vars)
        if outer_substitution:
            outer_lookups = [
                _VarSubstituter(outer_substitution).rewrite(lookup)
                for lookup in outer_lookups
            ]
            outer_vars = {
                outer_substitution[var] if var in outer_substitution else var
                for var in outer_vars
            }

        substitutions = self._correlation_substitutions(
            tasks=logical.body,
            outer_vars=outer_vars_for_correlation,
        )
        if not substitutions:
            return logical, False

        inner_tasks = tuple(_VarSubstituter(substitutions).rewrite(task) for task in logical.body)
        if outer_substitution:
            inner_tasks = tuple(_VarSubstituter(outer_substitution).rewrite(task) for task in inner_tasks)
        protected_local_vars = self._derive_protected_local_vars(tasks=inner_tasks, outer_vars=outer_vars)
        local_bindings: dict[Var, Value] = {}

        rewritten: list[Task] = []
        changed = False
        for task in inner_tasks:
            assert isinstance(task, Lookup)
            local_vars = [
                arg for arg in task.args
                if isinstance(arg, Var) and arg not in outer_vars
            ]
            if self._is_equality_lookup(task) and len(task.args) == 2 and task.args[0] == task.args[1]:
                changed = True
                continue
            if any(var in protected_local_vars for var in local_vars):
                rewritten.append(task)
                continue

            matched = False
            next_bindings = local_bindings
            for outer_lookup in outer_lookups:
                ok, candidate_bindings = self._lookup_matches_with_local_binding(
                    inner_lookup=task,
                    outer_lookup=outer_lookup,
                    outer_vars=outer_vars,
                    local_bindings=local_bindings,
                )
                if ok:
                    matched = True
                    next_bindings = candidate_bindings
                    break
            if matched:
                local_bindings = next_bindings
                changed = True
                continue
            rewritten.append(task)

        deduped: list[Task] = []
        for task in rewritten:
            if any(task == prior for prior in deduped):
                changed = True
                continue
            deduped.append(task)

        pruned: list[Task] = []
        for task in deduped:
            if not isinstance(task, Lookup):
                pruned.append(task)
                continue
            if len(task.args) != 1 or not isinstance(task.args[0], Var):
                pruned.append(task)
                continue
            local_var = task.args[0]
            if local_var in outer_vars:
                pruned.append(task)
                continue
            bound = local_bindings.get(local_var)
            if isinstance(bound, Var) and bound in outer_vars:
                changed = True
                continue
            pruned.append(task)
        deduped = pruned

        if not changed:
            return logical, False
        return logical.mut(body=tuple(deduped)), True

    def _simplify_correlated_quantifier_task(
        self,
        *,
        task: Task,
        outer_tasks: tuple[Task, ...],
    ) -> tuple[Task, bool]:
        if isinstance(task, Exists) and isinstance(task.task, Logical):
            inner_logical = task.task
            changed_any = False
            if (
                len(inner_logical.body) == 1
                and isinstance(inner_logical.body[0], Not)
                and isinstance(inner_logical.body[0].task, Logical)
            ):
                negated = inner_logical.body[0]
                assert isinstance(negated.task, Logical)
                rewritten_neg_body, changed = self._rewrite_correlated_lookup_logical(
                    logical=negated.task,
                    outer_tasks=outer_tasks,
                )
                if changed:
                    changed_any = True
                    negated = negated.mut(task=rewritten_neg_body)
                    inner_logical = inner_logical.mut(body=(negated,))
            else:
                rewritten_inner, changed = self._rewrite_correlated_lookup_logical(
                    logical=inner_logical,
                    outer_tasks=outer_tasks,
                )
                if changed:
                    changed_any = True
                    inner_logical = rewritten_inner
            if changed_any:
                return task.mut(task=inner_logical), True
            return task, False

        if isinstance(task, Not) and isinstance(task.task, Logical):
            rewritten_inner, changed = self._rewrite_correlated_lookup_logical(
                logical=task.task,
                outer_tasks=outer_tasks,
            )
            if changed:
                return task.mut(task=rewritten_inner), True
        return task, False

    def _simplify_correlated_quantifier_tasks(self, tasks: list[Task]) -> tuple[list[Task], bool]:
        rewritten: list[Task] = []
        changed = False
        for task in tasks:
            updated, task_changed = self._simplify_correlated_quantifier_task(
                task=task,
                outer_tasks=tuple(rewritten),
            )
            rewritten.append(updated)
            changed = changed or task_changed
        return rewritten, changed

    def _coalesce_duplicate_aggregates(self, tasks: list[Task]) -> tuple[list[Task], bool]:
        seen: dict[tuple[object, ...], Var] = {}
        substitutions: dict[Var, Var] = {}
        rewritten: list[Task] = []
        changed = False
        for task in tasks:
            current = _VarSubstituter(substitutions).rewrite(task) if substitutions else task
            if not isinstance(current, Aggregate) or not current.args or not isinstance(current.args[-1], Var):
                rewritten.append(current)
                continue
            output_var = current.args[-1]
            signature = (
                current.aggregation,
                tuple(current.projection),
                tuple(current.group),
                tuple(current.args[:-1]),
                current.body,
            )
            existing = seen.get(signature)
            if existing is not None:
                substitutions[output_var] = existing
                changed = True
                continue
            seen[signature] = output_var
            rewritten.append(current)
        return rewritten, changed

    def _schedule_tasks(self, tasks: tuple[Task, ...]) -> tuple[Task, ...]:
        if not tasks:
            return tuple()

        updates = tuple(task for task in tasks if isinstance(task, Update))
        schedulable = tuple(task for task in tasks if not isinstance(task, Update))
        if not schedulable:
            return updates

        # Repeated references to the same task id are redundant in a
        # conjunction. Keep the last occurrence so filters run in the most
        # informed context.
        last_index_by_id: dict[int, int] = {task.id: ix for ix, task in enumerate(schedulable)}
        schedulable = tuple(task for ix, task in enumerate(schedulable) if last_index_by_id[task.id] == ix)
        if not schedulable:
            return updates

        ordered: list[Task] = []
        bound: set[int] = set()
        pending = list(schedulable)
        while pending:
            progressed = False
            bindable_by_remaining: list[set[int]] = [set() for _ in pending]
            non_lookup_producible_by_remaining: list[set[int]] = [set() for _ in pending]
            has_lookup_by_remaining: list[bool] = [False for _ in pending]
            running: set[int] = set()
            running_non_lookup: set[int] = set()
            running_has_lookup = False
            for ix in range(len(pending) - 1, -1, -1):
                bindable_by_remaining[ix] = set(running)
                non_lookup_producible_by_remaining[ix] = set(running_non_lookup)
                has_lookup_by_remaining[ix] = running_has_lookup
                running.update(bindable_var_ids(pending[ix]))
                if not isinstance(pending[ix], Lookup):
                    if not isinstance(pending[ix], Union):
                        running_non_lookup.update(produced_var_ids(pending[ix]))
                else:
                    running_has_lookup = True

            deferred: list[Task] = []
            deferred_bindable: set[int] = set()
            for ix, task in enumerate(pending):
                available_later = bindable_by_remaining[ix].union(deferred_bindable)
                if self._should_defer_filter_task(
                    task=task,
                    bound=bound,
                    bindable_by_remaining=available_later,
                    has_lookup_by_remaining=has_lookup_by_remaining[ix],
                ):
                    deferred.append(task)
                    deferred_bindable.update(bindable_var_ids(task))
                    continue
                if isinstance(task, Construct):
                    missing_inputs = {
                        value.id
                        for value in task.values
                        if isinstance(value, Var) and value.id not in bound
                    }
                    if missing_inputs and missing_inputs.intersection(available_later):
                        deferred.append(task)
                        deferred_bindable.update(bindable_var_ids(task))
                        continue
                if isinstance(task, Union) and task.outputs:
                    missing_outputs = {
                        value.id
                        for value in task.outputs
                        if isinstance(value, Var) and value.id not in bound
                    }
                    if missing_outputs and missing_outputs.intersection(available_later):
                        deferred.append(task)
                        deferred_bindable.update(bindable_var_ids(task))
                        continue
                if isinstance(task, Lookup):
                    missing_input_vars = {
                        arg.id
                        for field, arg in zip(task.relation.fields, task.args)
                        if field.input and isinstance(arg, Var) and arg.id not in bound
                    }
                    if missing_input_vars and missing_input_vars.intersection(available_later):
                        deferred.append(task)
                        deferred_bindable.update(bindable_var_ids(task))
                        continue
                    unresolved_output_vars = {
                        arg.id
                        for field, arg in zip(task.relation.fields, task.args)
                        if not field.input and isinstance(arg, Var) and arg.id not in bound
                    }
                    if unresolved_output_vars and unresolved_output_vars.intersection(non_lookup_producible_by_remaining[ix]):
                        deferred.append(task)
                        deferred_bindable.update(bindable_var_ids(task))
                        continue
                ordered.append(task)
                bound.update(bindable_var_ids(task))
                progressed = True

            if not deferred:
                break
            if not progressed:
                ordered.extend(deferred)
                break
            pending = deferred

        ordered = self._push_filters_after_binding_lookups(ordered)
        ordered = self._push_filter_lookups_before_aggregates(ordered)
        ordered = self._push_correlated_quantifiers_before_aggregates(ordered)
        return tuple(ordered) + updates

    def _push_filters_after_binding_lookups(self, ordered: list[Task]) -> list[Task]:
        tasks = list(ordered)
        ix = 0
        while ix < len(tasks):
            task = tasks[ix]
            if not self._is_filter_task(task):
                ix += 1
                continue
            bound_before = set()
            for prior in tasks[:ix]:
                bound_before.update(bindable_var_ids(prior))
            missing = self._filter_free_var_ids(task).difference(bound_before)
            if not missing:
                ix += 1
                continue
            move_to: int | None = None
            for look_ix in range(ix + 1, len(tasks)):
                later = tasks[look_ix]
                if not isinstance(later, Lookup):
                    continue
                if missing.intersection(bindable_var_ids(later)):
                    move_to = look_ix + 1
                    break
            if move_to is None:
                ix += 1
                continue
            moved = tasks.pop(ix)
            if move_to > ix:
                move_to -= 1
            tasks.insert(move_to, moved)
        return tasks

    def _push_filter_lookups_before_aggregates(self, ordered: list[Task]) -> list[Task]:
        tasks = list(ordered)
        while True:
            moved_any = False
            bound_before: set[int] = set()
            for ix, task in enumerate(tasks):
                if not isinstance(task, Aggregate):
                    bound_before.update(bindable_var_ids(task))
                    continue

                move_indices: list[int] = []
                for look_ix in range(ix + 1, len(tasks)):
                    candidate = tasks[look_ix]
                    if not self._is_pure_filter_lookup(candidate):
                        continue
                    required_var_ids = task_var_ids(candidate)
                    if required_var_ids.issubset(bound_before):
                        move_indices.append(look_ix)
                if not move_indices:
                    bound_before.update(bindable_var_ids(task))
                    continue

                moved_tasks = [tasks[pos] for pos in move_indices]
                for pos in reversed(move_indices):
                    tasks.pop(pos)
                insert_at = ix
                for moved_task in moved_tasks:
                    tasks.insert(insert_at, moved_task)
                    insert_at += 1
                moved_any = True
                break

            if not moved_any:
                return tasks

    def _push_correlated_quantifiers_before_aggregates(self, ordered: list[Task]) -> list[Task]:
        tasks = list(ordered)
        while True:
            moved_any = False
            bound_before: set[int] = set()
            for ix, task in enumerate(tasks):
                if not isinstance(task, Aggregate):
                    bound_before.update(bindable_var_ids(task))
                    continue

                aggregate_output_var_ids: set[int] = set()
                if task.args and isinstance(task.args[-1], Var):
                    aggregate_output_var_ids.add(task.args[-1].id)

                move_indices: list[int] = []
                for look_ix in range(ix + 1, len(tasks)):
                    candidate = tasks[look_ix]
                    if not isinstance(candidate, (Exists, Not)):
                        continue
                    candidate_var_ids = task_var_ids(candidate)
                    if candidate_var_ids.intersection(aggregate_output_var_ids):
                        continue
                    # Move only correlated quantifiers. Uncorrelated ones should
                    # stay in place to preserve scope semantics.
                    if not candidate_var_ids.intersection(bound_before):
                        continue
                    move_indices.append(look_ix)
                if not move_indices:
                    bound_before.update(bindable_var_ids(task))
                    continue

                moved_tasks = [tasks[pos] for pos in move_indices]
                for pos in reversed(move_indices):
                    tasks.pop(pos)
                insert_at = ix
                for moved_task in moved_tasks:
                    tasks.insert(insert_at, moved_task)
                    insert_at += 1
                moved_any = True
                break

            if not moved_any:
                return tasks

    def _is_pure_filter_lookup(self, task: Task) -> bool:
        return (
            isinstance(task, Lookup)
            and bool(task.relation.fields)
            and all(field.input for field in task.relation.fields)
        )

    def _is_filter_task(self, task: Task) -> bool:
        return isinstance(task, (Not, Exists)) or (isinstance(task, Union) and not task.outputs)

    def _filter_free_var_ids(self, task: Task) -> set[int]:
        if isinstance(task, (Not, Exists)):
            return task_var_ids(task.task)
        if isinstance(task, Union) and not task.outputs:
            return task_var_ids(task)
        return set()

    def _should_defer_filter_task(
        self,
        *,
        task: Task,
        bound: set[int],
        bindable_by_remaining: set[int],
        has_lookup_by_remaining: bool,
    ) -> bool:
        if isinstance(task, (Not, Exists)):
            if has_lookup_by_remaining:
                return True
            missing = task_var_ids(task.task) - bound
            if missing and missing.intersection(bindable_by_remaining):
                return True
            if not bound and bindable_by_remaining:
                return True
            return False
        if isinstance(task, Union) and not task.outputs:
            if has_lookup_by_remaining:
                return True
            missing = task_var_ids(task) - bound
            if missing and missing.intersection(bindable_by_remaining):
                return True
            if not bound and bindable_by_remaining:
                return True
            return False
        return False

    def _canonicalize_bound_output_unions(self, tasks: list[Task]) -> tuple[list[Task], bool]:
        rewritten: list[Task] = []
        changed = False
        produced: set[int] = set()
        for task in tasks:
            current = task
            if isinstance(task, Union) and task.outputs:
                output_var_ids = {value.id for value in task.outputs if isinstance(value, Var)}
                if output_var_ids and output_var_ids.issubset(produced):
                    current = task.mut(outputs=())
                    changed = True
            rewritten.append(current)
            produced.update(produced_var_ids(current))
        return rewritten, changed

    def _canonicalize_optional_forms(self, tasks: list[Task]) -> tuple[list[Task], bool]:
        rewritten: list[Task] = []
        changed = False
        bound: set[int] = set()
        for task in tasks:
            current = task
            if isinstance(task, Logical) and task.optional:
                newly_bound = bindable_var_ids(task).difference(bound)
                if not newly_bound:
                    changed = True
                    continue
                if not bound:
                    current = task.mut(optional=False)
                    changed = True
            rewritten.append(current)
            bound.update(bindable_var_ids(current))
        return rewritten, changed

    def logical(self, node: Logical):
        if not node.body:
            return node
        tasks = list(node.body)
        changed = False

        scheduled = self._schedule_tasks(tuple(tasks))
        if tuple(tasks) != scheduled:
            tasks = list(scheduled)
            changed = True
        tasks, quantifier_simplified = self._simplify_correlated_quantifier_tasks(tasks)
        changed = changed or quantifier_simplified
        tasks, union_canonicalized = self._canonicalize_bound_output_unions(tasks)
        changed = changed or union_canonicalized
        tasks, optional_canonicalized = self._canonicalize_optional_forms(tasks)
        changed = changed or optional_canonicalized

        if self._aggregate_depth == 0 and self._container_depth == 0 and not node.optional:
            tasks, coalesced = self._coalesce_duplicate_aggregates(tasks)
            changed = changed or coalesced

        if changed:
            return node.mut(body=tuple(tasks))
        return node


class SQL2Normalizer:
    """Post-coalescing aggregate normalization for SQL2 metamodel artifacts."""

    def __init__(self):
        self._rewriter = _SQL2TaskNormalizer()
        self._max_rewrite_passes = 1

    def _normalize_task_to_fixpoint(self, task: Task) -> Task:
        current = task
        for _ in range(self._max_rewrite_passes):
            rewritten = self._rewriter.rewrite(current)
            if rewritten == current:
                return rewritten
            current = rewritten
        return current

    def normalize(
        self,
        coalesced: SQL2CoalescedArtifacts,
        *,
        model_analysis: ModelAnalysis,
        table_shape: SQL2TableShapeSidecar,
        model_key_preserving_table_names: tuple[str, ...] = (),
    ) -> SQL2NormalizedArtifacts:
        groups: list[SQL2NormalizedGroup] = []
        groups_by_stratum: dict[int, list[SQL2NormalizedGroup]] = {}
        instances_by_shape: dict[SQL2AggregateShape, list[SQL2AggregateInstance]] = {}

        for _stratum_id, stratum_groups in sorted(coalesced.groups_by_stratum.items()):
            for group in stratum_groups:
                normalized_channels: list[SQL2NormalizedChannel] = []
                for channel in group.channels:
                    normalized_rule = self._normalize_task_to_fixpoint(channel.rule)
                    update_task, body_rule = self._extract_update_and_body_rule(normalized_rule)
                    required_output_var_ids = {
                        arg.id
                        for field, arg in zip(update_task.relation.fields, update_task.args)
                        if not field.input and isinstance(arg, Var)
                    }
                    optional_projection_plan = self._optional_projection_plan(body_rule)
                    future_var_plan = self._future_var_plan(
                        body_rule,
                        required_output_var_ids=required_output_var_ids,
                    )
                    branch_projection_plan = self._branch_projection_plan(body_rule)
                    union_mode_plan = self._union_mode_plan(
                        body_rule,
                        required_output_var_ids=required_output_var_ids,
                        branch_projection_plan=branch_projection_plan,
                        future_var_plan=future_var_plan,
                    )
                    channel_read_table_names = self._read_table_names(
                        task=body_rule,
                        model_analysis=model_analysis,
                        table_shape=table_shape,
                    )
                    instances = self._collect_aggregate_instances(
                        rule=normalized_rule,
                        group_id=group.group_id,
                        channel_id=channel.channel_id,
                    )
                    normalized_channel = SQL2NormalizedChannel(
                        channel_id=channel.channel_id,
                        node_id=channel.node_id,
                        target_relation_id=channel.target_relation_id,
                        target_relation_name=channel.target_relation_name,
                        output_columns=channel.output_columns,
                        key_var_ids=channel.key_var_ids,
                        rule=normalized_rule,
                        update_task=update_task,
                        body_rule=body_rule,
                        optional_projection_var_ids_by_task_id=tuple(sorted(optional_projection_plan.items())),
                        branch_projection_var_ids_by_task_id=tuple(sorted(branch_projection_plan.items())),
                        union_mode_by_task_id=tuple(sorted(union_mode_plan.items())),
                        future_var_ids_by_task_id=tuple(sorted(future_var_plan.items())),
                        read_table_names=channel_read_table_names,
                        aggregate_instances=instances,
                    )
                    normalized_channels.append(normalized_channel)
                    for instance in instances:
                        instances_by_shape.setdefault(instance.shape, []).append(instance)

                normalized_rule = (
                    Logical(body=tuple(channel.rule for channel in normalized_channels), scope=True)
                    if normalized_channels
                    else None
                )
                normalized_group = SQL2NormalizedGroup(
                    group_id=group.group_id,
                    stratum_id=group.stratum_id,
                    table_name=group.table_name,
                    key_var_ids=group.key_var_ids,
                    output_columns=group.output_columns,
                    key_columns=group.key_columns,
                    channel_combine_plan=self._channel_combine_plan_for_group(
                        group=group,
                        channels=tuple(normalized_channels),
                        table_shape=table_shape,
                    ),
                    read_table_names=self._read_table_names(
                        task=normalized_rule,
                        model_analysis=model_analysis,
                        table_shape=table_shape,
                    ),
                    node_ids=group.node_ids,
                    channels=tuple(normalized_channels),
                    normalized_rule=normalized_rule,
                )
                groups.append(normalized_group)
                groups_by_stratum.setdefault(group.stratum_id, []).append(normalized_group)

        ordered_groups = tuple(groups)
        ordered_groups_by_stratum = {
            stratum_id: tuple(stratum_groups)
            for stratum_id, stratum_groups in sorted(groups_by_stratum.items())
        }
        aggregate_instances_by_shape = {
            shape: tuple(sorted(instances, key=lambda inst: (inst.group_id, inst.channel_id, inst.aggregate_task_id)))
            for shape, instances in sorted(
                instances_by_shape.items(),
                key=lambda item: (
                    item[0].aggregation_relation_id,
                    item[0].group_shape,
                    item[0].projection_shape,
                    item[0].non_output_arg_shape,
                ),
            )
        }
        table_plans = self._table_plans(
            groups=ordered_groups,
            table_shape=table_shape,
        )
        relation_key_preserving_by_relation_id = self._relation_key_preserving_by_relation_id(
            groups=ordered_groups,
            table_shape=table_shape,
        )
        relation_direct_projection_safe_by_relation_id = self._relation_direct_projection_safe_by_relation_id(
            groups=ordered_groups,
            relation_key_preserving_by_relation_id=relation_key_preserving_by_relation_id,
            model_key_preserving_table_names=model_key_preserving_table_names,
        )
        relation_output_anchor_safe_by_relation_id = self._relation_output_anchor_safe_by_relation_id(
            groups=ordered_groups,
            relation_direct_projection_safe_by_relation_id=relation_direct_projection_safe_by_relation_id,
        )
        preseed_tables = self._preseed_table_plan(
            groups=ordered_groups,
            model_analysis=model_analysis,
            table_shape=table_shape,
        )
        program_by_stratum: dict[int, Task] = {}
        for stratum_id, stratum_groups in ordered_groups_by_stratum.items():
            rules = tuple(
                group.normalized_rule
                for group in stratum_groups
                if group.normalized_rule is not None
            )
            program_by_stratum[stratum_id] = Logical(body=rules, scope=True)
        program_root = (
            Logical(body=tuple(program_by_stratum[stratum_id] for stratum_id in sorted(program_by_stratum)), scope=True)
            if program_by_stratum
            else None
        )
        return SQL2NormalizedArtifacts(
            groups=ordered_groups,
            groups_by_stratum=ordered_groups_by_stratum,
            aggregate_instances_by_shape=aggregate_instances_by_shape,
            relation_key_preserving_by_relation_id=relation_key_preserving_by_relation_id,
            relation_direct_projection_safe_by_relation_id=relation_direct_projection_safe_by_relation_id,
            relation_output_anchor_safe_by_relation_id=relation_output_anchor_safe_by_relation_id,
            table_plans=table_plans,
            preseed_tables=preseed_tables,
            program_by_stratum=program_by_stratum,
            program_root=program_root,
        )

    def _read_table_names(
        self,
        *,
        task: Task | None,
        model_analysis: ModelAnalysis,
        table_shape: SQL2TableShapeSidecar,
    ) -> tuple[str, ...]:
        if task is None:
            return ()
        out: set[str] = set()

        def table_name_for_lookup(lookup: Lookup) -> str | None:
            table = table_shape.table_for_relation_id(lookup.relation.id)
            if table is not None:
                return table.table_name
            try:
                table_info = model_analysis.get_table_info(lookup.relation)
            except Exception:
                table_info = None
            if table_info is not None:
                return table_info.name
            try:
                column_info = model_analysis.get_column_info(lookup.relation)
            except Exception:
                column_info = None
            if column_info is not None and column_info.table is not None:
                return column_info.table.name
            return None

        def visit(node: Task) -> None:
            if isinstance(node, Lookup):
                table_name = table_name_for_lookup(node)
                if table_name is not None:
                    out.add(table_name)
                return
            for child in self._iter_child_tasks(node):
                visit(child)

        visit(task)
        return tuple(sorted(out))

    def _channel_combine_plan_for_group(
        self,
        *,
        group,
        channels: tuple[SQL2NormalizedChannel, ...],
        table_shape: SQL2TableShapeSidecar,
    ) -> SQL2ChannelCombinePlan:
        output_columns = tuple(group.output_columns)
        key_column_set = set(group.key_columns)
        key_columns = tuple(column for column in output_columns if column in key_column_set)
        if not output_columns or not key_columns:
            channel_groups = tuple((channel.channel_id,) for channel in channels)
            return SQL2ChannelCombinePlan(
                strategy="set_union",
                output_columns=output_columns,
                key_columns=key_columns,
                aggregate_columns=(),
                collapsed_channel_groups=channel_groups,
            )

        assignments_by_channel_id: dict[int, dict[str, Value | None]] = {
            channel.channel_id: self._channel_output_assignments(
                channel=channel,
                output_columns=output_columns,
                table_shape=table_shape,
            )
            for channel in channels
        }

        buckets: dict[tuple[object, tuple[Value | None, ...]], list[list[int]]] = {}
        for channel in channels:
            assignments = assignments_by_channel_id[channel.channel_id]
            key_signature = tuple(assignments.get(column) for column in key_columns)
            bucket_key = (self._channel_body_signature(channel.body_rule), key_signature)
            cluster_pool = buckets.setdefault(bucket_key, [])
            placed = False
            for cluster in cluster_pool:
                existing = [assignments_by_channel_id[channel_id] for channel_id in cluster]
                if self._channel_assignments_compatible(
                    existing=existing,
                    candidate=assignments,
                    output_columns=output_columns,
                ):
                    cluster.append(channel.channel_id)
                    placed = True
                    break
            if not placed:
                cluster_pool.append([channel.channel_id])

        collapsed_channel_groups: list[tuple[int, ...]] = []
        for channel in channels:
            for cluster_pool in buckets.values():
                for cluster in cluster_pool:
                    if cluster and cluster[0] == channel.channel_id:
                        collapsed_channel_groups.append(tuple(cluster))
                        break
                else:
                    continue
                break

        aggregate_columns = tuple(column for column in output_columns if column not in key_column_set)
        return SQL2ChannelCombinePlan(
            strategy="by_key_max",
            output_columns=output_columns,
            key_columns=key_columns,
            aggregate_columns=aggregate_columns,
            collapsed_channel_groups=tuple(collapsed_channel_groups),
        )

    def _channel_output_assignments(
        self,
        *,
        channel: SQL2NormalizedChannel,
        output_columns: tuple[str, ...],
        table_shape: SQL2TableShapeSidecar,
    ) -> dict[str, Value | None]:
        output: dict[str, Value | None] = {column_name: None for column_name in output_columns}
        update = channel.update_task
        property_column = table_shape.property_column_by_relation_id.get(update.relation.id)
        if property_column is not None:
            if "_id" in output and update.args:
                output["_id"] = update.args[0]
            if property_column in output and update.args:
                output[property_column] = update.args[-1]
            return output
        for ix, rel_field in enumerate(update.relation.fields):
            column_name = "_id" if ix == 0 and "_id" in output else rel_field.name
            if column_name not in output:
                continue
            if ix >= len(update.args):
                continue
            output[column_name] = update.args[ix]
        return output

    def _channel_assignments_compatible(
        self,
        *,
        existing: list[dict[str, Value | None]],
        candidate: dict[str, Value | None],
        output_columns: tuple[str, ...],
    ) -> bool:
        for column in output_columns:
            existing_non_null: Value | None = None
            for assignment in existing:
                value = assignment.get(column)
                if self._value_is_null(value):
                    continue
                if existing_non_null is None:
                    existing_non_null = value
                    continue
                if value != existing_non_null:
                    return False
            candidate_value = candidate.get(column)
            if self._value_is_null(candidate_value):
                continue
            if existing_non_null is not None and candidate_value != existing_non_null:
                return False
        return True

    def _value_is_null(self, value: Value | None) -> bool:
        return value is None or (isinstance(value, Literal) and value.value is None)

    def _relation_key_preserving_by_relation_id(
        self,
        *,
        groups: tuple[SQL2NormalizedGroup, ...],
        table_shape: SQL2TableShapeSidecar,
    ) -> tuple[tuple[int, bool], ...]:
        channels_by_relation_id: dict[int, list[SQL2NormalizedChannel]] = {}
        groups_by_relation_id: dict[int, set[str]] = {}
        group_by_id = {group.group_id: group for group in groups}

        for group in groups:
            for channel in group.channels:
                channels_by_relation_id.setdefault(channel.target_relation_id, []).append(channel)
                groups_by_relation_id.setdefault(channel.target_relation_id, set()).add(group.group_id)

        plan: list[tuple[int, bool]] = []
        for relation_id, channels in sorted(channels_by_relation_id.items()):
            key_preserving = True
            for channel in channels:
                key_preserving = key_preserving and self._channel_is_key_preserving(
                    channel,
                    table_shape=table_shape,
                )
                if not key_preserving:
                    break
            if key_preserving and len(channels) > 1:
                for group_id in groups_by_relation_id.get(relation_id, set()):
                    group = group_by_id[group_id]
                    combine_plan = group.channel_combine_plan
                    if combine_plan.strategy != "by_key_max" or len(combine_plan.key_columns) != 1:
                        key_preserving = False
                        break
            plan.append((relation_id, key_preserving))
        return tuple(plan)

    def _relation_direct_projection_safe_by_relation_id(
        self,
        *,
        groups: tuple[SQL2NormalizedGroup, ...],
        relation_key_preserving_by_relation_id: tuple[tuple[int, bool], ...],
        model_key_preserving_table_names: tuple[str, ...],
    ) -> tuple[tuple[int, bool], ...]:
        read_table_names_by_relation_id: dict[int, set[str]] = {}
        for group in groups:
            for channel in group.channels:
                read_table_names_by_relation_id.setdefault(channel.target_relation_id, set()).update(
                    channel.read_table_names
                )

        key_preserving_by_relation_id = dict(relation_key_preserving_by_relation_id)
        model_key_preserving_table_name_set = set(model_key_preserving_table_names)

        relation_ids = sorted(
            set(key_preserving_by_relation_id).union(read_table_names_by_relation_id)
        )
        plan: list[tuple[int, bool]] = []
        for relation_id in relation_ids:
            if not key_preserving_by_relation_id.get(relation_id, False):
                plan.append((relation_id, False))
                continue
            read_table_names = read_table_names_by_relation_id.get(relation_id, set())
            if not read_table_names:
                plan.append((relation_id, True))
                continue
            if len(read_table_names) != 1:
                plan.append((relation_id, False))
                continue
            safe = bool(model_key_preserving_table_name_set) and read_table_names.issubset(
                model_key_preserving_table_name_set
            )
            plan.append((relation_id, safe))
        return tuple(plan)

    def _relation_output_anchor_safe_by_relation_id(
        self,
        *,
        groups: tuple[SQL2NormalizedGroup, ...],
        relation_direct_projection_safe_by_relation_id: tuple[tuple[int, bool], ...],
    ) -> tuple[tuple[int, bool], ...]:
        channel_read_table_counts_by_relation_id: dict[int, list[int]] = {}
        for group in groups:
            for channel in group.channels:
                channel_read_table_counts_by_relation_id.setdefault(channel.target_relation_id, []).append(
                    len(channel.read_table_names)
                )

        direct_safe_by_relation_id = dict(relation_direct_projection_safe_by_relation_id)
        relation_ids = sorted(set(direct_safe_by_relation_id).union(channel_read_table_counts_by_relation_id))
        plan: list[tuple[int, bool]] = []
        for relation_id in relation_ids:
            if not direct_safe_by_relation_id.get(relation_id, False):
                plan.append((relation_id, False))
                continue
            read_counts = channel_read_table_counts_by_relation_id.get(relation_id, [])
            # Output-anchor projection is only safe when the relation has a
            # concrete source rowset channel (exactly one read-table source).
            safe = any(count == 1 for count in read_counts)
            plan.append((relation_id, safe))
        return tuple(plan)

    def _channel_is_key_preserving(
        self,
        channel: SQL2NormalizedChannel,
        *,
        table_shape: SQL2TableShapeSidecar,
    ) -> bool:
        if len(channel.key_var_ids) != 1:
            return False
        if not channel.update_task.args or not isinstance(channel.update_task.args[0], Var):
            return False
        if channel.update_task.args[0].id != channel.key_var_ids[0]:
            return False

        bindable_ids = bindable_var_ids(channel.body_rule)
        update_arg_var_ids = {
            arg.id
            for arg in channel.update_task.args
            if isinstance(arg, Var)
        }
        if not update_arg_var_ids:
            return False
        if self._body_has_non_key_preserving_ops(channel.body_rule):
            return False
        if not self._key_constructs_are_identity_safe(channel):
            return False
        if not self._bindable_vars_determined_by_key(
            channel=channel,
            key_var_id=channel.key_var_ids[0],
            bindable_ids=bindable_ids,
            update_arg_var_ids=update_arg_var_ids,
            table_shape=table_shape,
        ):
            return False
        return True

    def _key_constructs_are_identity_safe(self, channel: SQL2NormalizedChannel) -> bool:
        key_var_id = channel.key_var_ids[0]
        key_constructs = [
            task
            for task in self._walk_tasks(channel.body_rule)
            if isinstance(task, Construct) and task.id_var.id == key_var_id
        ]
        if not key_constructs:
            return True
        return all(self._construct_inverse_is_key_safe(construct) for construct in key_constructs)

    def _body_has_non_key_preserving_ops(self, task: Task) -> bool:
        if isinstance(task, Logical):
            return any(self._body_has_non_key_preserving_ops(child) for child in task.body)
        if isinstance(task, (Lookup, Construct)):
            return False
        return True

    def _bindable_vars_determined_by_key(
        self,
        *,
        channel: SQL2NormalizedChannel,
        key_var_id: int,
        bindable_ids: set[int],
        update_arg_var_ids: set[int],
        table_shape: SQL2TableShapeSidecar,
    ) -> bool:
        closure = set(update_arg_var_ids)
        tasks = tuple(self._walk_tasks(channel.body_rule))

        changed = True
        while changed:
            changed = False
            for task in tasks:
                if isinstance(task, Construct):
                    value_var_ids = self._value_var_ids(task.values)
                    if (
                        task.id_var.id in closure
                        and task.id_var.id == key_var_id
                        and self._construct_inverse_is_key_safe(task)
                        and not value_var_ids.issubset(closure)
                    ):
                        before = len(closure)
                        closure.update(value_var_ids)
                        changed = changed or len(closure) != before
                    if value_var_ids.issubset(closure) and task.id_var.id not in closure:
                        closure.add(task.id_var.id)
                        changed = True
                elif isinstance(task, Lookup):
                    if not self._lookup_is_functional(task.relation, table_shape=table_shape):
                        continue
                    var_ids = self._value_var_ids(task.args)
                    unresolved_var_ids = var_ids.difference(closure)
                    resolved_var_ids = var_ids.intersection(closure)
                    if len(unresolved_var_ids) != 1 or not resolved_var_ids:
                        continue
                    unresolved_var_id = next(iter(unresolved_var_ids))
                    if unresolved_var_id not in closure:
                        closure.add(unresolved_var_id)
                        changed = True

        return bindable_ids.issubset(closure)

    def _construct_inverse_is_key_safe(self, construct: Construct) -> bool:
        value_vars = self._collect_value_vars(construct.values)
        if not value_vars:
            return True
        return all(is_entity_type(var.type) for var in value_vars)

    def _collect_value_vars(self, values: tuple[Value, ...] | list[Value]) -> tuple[Var, ...]:
        out: list[Var] = []
        for value in values:
            if isinstance(value, Var):
                out.append(value)
            elif isinstance(value, (tuple, list)):
                out.extend(self._collect_value_vars(value))
        return tuple(out)

    def _lookup_is_functional(self, relation, *, table_shape: SQL2TableShapeSidecar) -> bool:
        if self._is_builtin_relation(relation):
            return True
        return relation.id in table_shape.property_column_by_relation_id

    def _value_var_ids(self, values: tuple[Value, ...] | list[Value]) -> set[int]:
        out: set[int] = set()
        for value in values:
            if isinstance(value, Var):
                out.add(value.id)
            elif isinstance(value, (tuple, list)):
                out.update(self._value_var_ids(value))
        return out

    def _walk_tasks(self, task: Task):
        yield task
        for child in self._iter_child_tasks(task):
            yield from self._walk_tasks(child)

    def _channel_body_signature(self, task: Task) -> Task:
        # Use canonicalized task structure (ids already excluded from equality).
        return self._canonicalize_noop_logicals(task)

    def _canonicalize_noop_logicals(self, task: Task) -> Task:
        if isinstance(task, Logical):
            normalized = task.mut(body=tuple(self._canonicalize_noop_logicals(child) for child in task.body))
            while (
                not normalized.optional
                and len(normalized.body) == 1
                and isinstance(normalized.body[0], Logical)
                and not normalized.body[0].optional
            ):
                normalized = normalized.body[0]
            return normalized
        if isinstance(task, Not):
            return task.mut(task=self._canonicalize_noop_logicals(task.task))
        if isinstance(task, Exists):
            return task.mut(task=self._canonicalize_noop_logicals(task.task))
        if isinstance(task, Aggregate):
            return task.mut(body=self._canonicalize_noop_logicals(task.body))
        if isinstance(task, Union):
            return task.mut(tasks=tuple(self._canonicalize_noop_logicals(child) for child in task.tasks))
        if isinstance(task, Match):
            return task.mut(tasks=tuple(self._canonicalize_noop_logicals(child) for child in task.tasks))
        if isinstance(task, Sequence):
            return task.mut(tasks=tuple(self._canonicalize_noop_logicals(child) for child in task.tasks))
        if isinstance(task, Loop):
            return task.mut(body=self._canonicalize_noop_logicals(task.body))
        if isinstance(task, Require):
            return task.mut(
                domain=self._canonicalize_noop_logicals(task.domain),
                check=self._canonicalize_noop_logicals(task.check),
                error=self._canonicalize_noop_logicals(task.error) if task.error is not None else None,
            )
        return task

    def _preseed_table_plan(
        self,
        *,
        groups: tuple[SQL2NormalizedGroup, ...],
        model_analysis: ModelAnalysis,
        table_shape: SQL2TableShapeSidecar,
    ) -> tuple[SQL2PreseedTablePlan, ...]:
        table_columns: dict[str, tuple[str, ...]] = {}
        logical_tables = sorted(
            (
                table
                for table in model_analysis.table_info.values()
                if table.kind not in {TableKind.DATA, TableKind.EXTERNAL}
            ),
            key=lambda table: (table.name, table.relation.id),
        )
        for table in logical_tables:
            relation = table.relation
            if self._is_builtin_relation(relation):
                continue
            if isinstance(relation, TypeNode) and (is_primitive(relation) or is_value_type(relation) or is_metatype(relation)):
                continue
            table_name = table_shape.table_name_by_relation_id.get(relation.id, table.name)
            columns = tuple(column.name for column in table.columns)
            if not columns:
                columns = self._inherited_empty_columns(
                    relation=relation,
                    table_shape=table_shape,
                )
            if not columns:
                continue
            existing = table_columns.get(table_name, ())
            table_columns[table_name] = tuple(dict.fromkeys((*existing, *columns)))

        first_write_rank_by_table: dict[str, int] = {}
        first_read_rank_by_table: dict[str, int] = {}
        for rank, group in enumerate(groups):
            first_write_rank_by_table.setdefault(group.table_name, rank)
            read_table_names = set(group.read_table_names)
            for candidate_table in table_columns:
                if candidate_table in first_read_rank_by_table:
                    continue
                if candidate_table in read_table_names:
                    first_read_rank_by_table[candidate_table] = rank

        plans: list[SQL2PreseedTablePlan] = []
        for table_name, columns in sorted(table_columns.items()):
            first_write_rank = first_write_rank_by_table.get(table_name)
            first_read_rank = first_read_rank_by_table.get(table_name)
            needs_preseed = first_write_rank is None or (
                first_read_rank is not None and first_read_rank <= first_write_rank
            )
            if not needs_preseed:
                continue
            plans.append(SQL2PreseedTablePlan(table_name=table_name, output_columns=columns))
        return tuple(plans)

    def _table_plans(
        self,
        *,
        groups: tuple[SQL2NormalizedGroup, ...],
        table_shape: SQL2TableShapeSidecar,
    ) -> tuple[SQL2NormalizedTablePlan, ...]:
        output_by_table: dict[str, tuple[str, ...]] = {}
        key_by_table: dict[str, set[str]] = {}
        for table in table_shape.tables:
            output_by_table[table.table_name] = tuple(dict.fromkeys(column.column_name for column in table.columns))
            key_by_table[table.table_name] = set(table.key_column_names)
        for group in groups:
            existing_output = output_by_table.get(group.table_name, ())
            merged_output = tuple(dict.fromkeys((*existing_output, *group.output_columns)))
            output_by_table[group.table_name] = merged_output
            merged_keys = set(key_by_table.get(group.table_name, set()))
            merged_keys.update(group.key_columns)
            key_by_table[group.table_name] = merged_keys

        plans: list[SQL2NormalizedTablePlan] = []
        for table_name in sorted(output_by_table):
            output_columns = output_by_table[table_name]
            key_set = key_by_table.get(table_name, set())
            key_columns = tuple(column for column in output_columns if column in key_set)
            remaining_keys = tuple(sorted(column for column in key_set if column not in set(output_columns)))
            plans.append(
                SQL2NormalizedTablePlan(
                    table_name=table_name,
                    output_columns=output_columns,
                    key_columns=(*key_columns, *remaining_keys),
                )
            )
        return tuple(plans)

    def _is_builtin_relation(self, relation) -> bool:
        for library in builtin_registry._libraries.values():
            if library.data.get(relation.name) is not None:
                return True
        return False

    def _inherited_empty_columns(
        self,
        *,
        relation,
        table_shape: SQL2TableShapeSidecar,
    ) -> tuple[str, ...]:
        for super_type in getattr(relation, "super_types", ()):
            super_table = table_shape.table_for_relation_id(super_type.id)
            if super_table is None or not super_table.columns:
                continue
            return tuple(column.column_name for column in super_table.columns)
        return ()

    def _extract_update_and_body_rule(self, task: Task) -> tuple[Update, Logical]:
        update_task = self._find_update_task(task)
        if update_task is None:
            raise ValueError("SQL2 normalize expects each channel rule to contain exactly one Update task.")
        pruned = self._remove_update_from_task(task, update_task_id=update_task.id)
        if not isinstance(pruned, Logical):
            raise ValueError("SQL2 normalize expects channel rule without Update to remain Logical.")
        return update_task, pruned

    def _find_update_task(self, task: Task) -> Update | None:
        if isinstance(task, Update):
            return task
        for child in self._iter_child_tasks(task):
            found = self._find_update_task(child)
            if found is not None:
                return found
        return None

    def _iter_child_tasks(self, task: Task) -> tuple[Task, ...]:
        if isinstance(task, Logical):
            return tuple(task.body)
        if isinstance(task, Sequence):
            return tuple(task.tasks)
        if isinstance(task, Loop):
            return (task.body,)
        if isinstance(task, Require):
            out: list[Task] = [task.domain, task.check]
            if task.error is not None:
                out.append(task.error)
            return tuple(out)
        if isinstance(task, (Not, Exists)):
            return (task.task,)
        if isinstance(task, Aggregate):
            return (task.body,)
        if isinstance(task, (Union, Match)):
            return tuple(task.tasks)
        return ()

    def _remove_update_from_task(self, task: Task, *, update_task_id: int) -> Task | None:
        if isinstance(task, Update):
            return None if task.id == update_task_id else task
        if isinstance(task, Logical):
            new_body: list[Task] = []
            for child in task.body:
                rewritten = self._remove_update_from_task(child, update_task_id=update_task_id)
                if rewritten is not None:
                    new_body.append(rewritten)
            return task.mut(body=tuple(new_body))
        if isinstance(task, Not):
            rewritten = self._remove_update_from_task(task.task, update_task_id=update_task_id)
            if rewritten is None:
                return None
            return task.mut(task=rewritten)
        if isinstance(task, Exists):
            rewritten = self._remove_update_from_task(task.task, update_task_id=update_task_id)
            if rewritten is None:
                return None
            return task.mut(task=rewritten)
        if isinstance(task, Aggregate):
            rewritten = self._remove_update_from_task(task.body, update_task_id=update_task_id)
            if rewritten is None:
                return None
            return task.mut(body=rewritten)
        if isinstance(task, Union):
            new_tasks: list[Task] = []
            for child in task.tasks:
                rewritten = self._remove_update_from_task(child, update_task_id=update_task_id)
                if rewritten is not None:
                    new_tasks.append(rewritten)
            return task.mut(tasks=tuple(new_tasks))
        if isinstance(task, Match):
            new_tasks: list[Task] = []
            for child in task.tasks:
                rewritten = self._remove_update_from_task(child, update_task_id=update_task_id)
                if rewritten is not None:
                    new_tasks.append(rewritten)
            return task.mut(tasks=tuple(new_tasks))
        return task

    def _optional_projection_plan(self, task: Task) -> dict[int, tuple[int, ...]]:
        plan: dict[int, tuple[int, ...]] = {}

        def walk(node: Task, bound: set[int]) -> set[int]:
            if isinstance(node, Logical):
                local_bound = set(bound)
                for child in node.body:
                    if isinstance(child, Logical) and child.optional:
                        projected = tuple(sorted(bindable_var_ids(child).difference(local_bound)))
                        plan[child.id] = projected
                    child_bindable = walk(child, set(local_bound))
                    local_bound.update(child_bindable)
                return local_bound.difference(bound)
            if isinstance(node, Aggregate):
                walk(node.body, set(bound))
                return bindable_var_ids(node)
            if isinstance(node, (Not, Exists)):
                walk(node.task, set(bound))
                return bindable_var_ids(node)
            if isinstance(node, (Union, Match)):
                for child in node.tasks:
                    walk(child, set(bound))
                return bindable_var_ids(node)
            return bindable_var_ids(node)

        walk(task, set())
        return plan

    def _branch_projection_plan(self, task: Task) -> dict[int, tuple[int, ...]]:
        plan: dict[int, tuple[int, ...]] = {}

        def walk(node: Task) -> None:
            if isinstance(node, Union):
                explicit_output_var_ids = {value.id for value in node.outputs if isinstance(value, Var)}
                shared_branch_var_ids: set[int] | None = None
                for child in node.tasks:
                    child_bindable = bindable_var_ids(child)
                    if shared_branch_var_ids is None:
                        shared_branch_var_ids = set(child_bindable)
                    else:
                        shared_branch_var_ids.intersection_update(child_bindable)
                    walk(child)
                carry_var_ids = explicit_output_var_ids | (shared_branch_var_ids or set())
                if carry_var_ids:
                    plan[node.id] = tuple(sorted(carry_var_ids))
                return
            for child in self._iter_child_tasks(node):
                walk(child)

        walk(task)
        return plan

    def _future_var_plan(
        self,
        task: Task,
        *,
        required_output_var_ids: set[int],
    ) -> dict[int, tuple[int, ...]]:
        plan: dict[int, tuple[int, ...]] = {}

        def merge(task_id: int, future_var_ids: set[int]) -> None:
            current = plan.get(task_id)
            if current is None:
                plan[task_id] = tuple(sorted(future_var_ids))
                return
            plan[task_id] = tuple(sorted(set(current).union(future_var_ids)))

        def walk(node: Task, *, ambient_future_var_ids: set[int]) -> None:
            if isinstance(node, Logical):
                body = tuple(node.body)
                future_by_index: list[set[int]] = [set() for _ in body]
                future_var_ids: set[int] = set(ambient_future_var_ids)
                for ix in range(len(body) - 1, -1, -1):
                    future_by_index[ix] = set(future_var_ids)
                    future_var_ids.update(task_var_ids(body[ix]))
                for ix, child in enumerate(body):
                    merge(child.id, future_by_index[ix])
                    walk(child, ambient_future_var_ids=set(future_by_index[ix]))
                return
            if isinstance(node, Aggregate):
                body_future_var_ids = set(ambient_future_var_ids)
                body_future_var_ids.update(
                    value.id for value in [*node.group, *node.projection, *node.args[:-1]] if isinstance(value, Var)
                )
                merge(node.body.id, body_future_var_ids)
                walk(node.body, ambient_future_var_ids=body_future_var_ids)
                return
            if isinstance(node, (Union, Match)):
                branch_future_var_ids = set(ambient_future_var_ids)
                branch_future_var_ids.update(value.id for value in node.outputs if isinstance(value, Var))
                for child in node.tasks:
                    merge(child.id, branch_future_var_ids)
                    walk(child, ambient_future_var_ids=set(branch_future_var_ids))
                return
            for child in self._iter_child_tasks(node):
                merge(child.id, ambient_future_var_ids)
                walk(child, ambient_future_var_ids=set(ambient_future_var_ids))

        walk(task, ambient_future_var_ids=set(required_output_var_ids))
        return plan

    def _union_mode_plan(
        self,
        task: Task,
        *,
        required_output_var_ids: set[int],
        branch_projection_plan: dict[int, tuple[int, ...]],
        future_var_plan: dict[int, tuple[int, ...]],
    ) -> dict[int, str]:
        plan: dict[int, str] = {}

        def walk(
            node: Task,
            *,
            optional_scope: bool = False,
            bound_var_ids: set[int] | None = None,
        ) -> set[int]:
            if bound_var_ids is None:
                bound_var_ids = set()
            child_optional_scope = optional_scope
            if isinstance(node, Logical) and node.optional:
                child_optional_scope = True
            if isinstance(node, Logical):
                local_bound = set(bound_var_ids)
                for child in node.body:
                    child_bindable = walk(
                        child,
                        optional_scope=child_optional_scope,
                        bound_var_ids=set(local_bound),
                    )
                    local_bound.update(child_bindable)
                return local_bound.difference(bound_var_ids)
            if isinstance(node, Union):
                projection_var_ids = set(branch_projection_plan.get(node.id, ()))
                if node.outputs:
                    explicit_output_var_ids = {
                        value.id for value in node.outputs if isinstance(value, Var)
                    }
                    carry_var_ids = projection_var_ids.difference(explicit_output_var_ids)
                    allow_optional_by_signature = carry_var_ids.issubset(required_output_var_ids)
                    plan[node.id] = (
                        "bind_optional"
                        if (node.optional or (child_optional_scope and allow_optional_by_signature))
                        else "bind_required"
                    )
                    bound_by_node = set(explicit_output_var_ids)
                else:
                    # Union without explicit outputs is purely predicate semantics:
                    # it contributes branch disjunction but does not bind new vars.
                    plan[node.id] = "filter"
                    bound_by_node = set()
                for child in node.tasks:
                    walk(
                        child,
                        optional_scope=child_optional_scope,
                        bound_var_ids=set(bound_var_ids),
                    )
                return bound_by_node
            if isinstance(node, Aggregate):
                walk(
                    node.body,
                    optional_scope=child_optional_scope,
                    bound_var_ids=set(bound_var_ids),
                )
                return bindable_var_ids(node)
            for child in self._iter_child_tasks(node):
                walk(
                    child,
                    optional_scope=child_optional_scope,
                    bound_var_ids=set(bound_var_ids),
                )
            return bindable_var_ids(node)

        walk(task, bound_var_ids=set())
        return plan

    def _collect_aggregate_instances(
        self,
        *,
        rule: Task,
        group_id: str,
        channel_id: int,
    ) -> tuple[SQL2AggregateInstance, ...]:
        instances: list[SQL2AggregateInstance] = []

        def visit(task: Task):
            if isinstance(task, Aggregate):
                output_var_id = task.args[-1].id if task.args and isinstance(task.args[-1], Var) else None
                instances.append(
                    SQL2AggregateInstance(
                        shape=self._shape_for(task),
                        group_id=group_id,
                        channel_id=channel_id,
                        aggregate_task_id=task.id,
                        output_var_id=output_var_id,
                    )
                )
            for child in self._iter_child_tasks(task):
                visit(child)

        visit(rule)
        instances.sort(key=lambda inst: (inst.aggregate_task_id, inst.channel_id))
        return tuple(instances)

    def _shape_for(self, aggregate: Aggregate) -> SQL2AggregateShape:
        return SQL2AggregateShape(
            aggregation_relation_id=aggregate.aggregation.id,
            group_shape=tuple(self._value_shape(value) for value in aggregate.group),
            projection_shape=tuple(self._value_shape(value) for value in aggregate.projection),
            non_output_arg_shape=tuple(self._value_shape(value) for value in aggregate.args[:-1]),
        )

    def _value_shape(self, value: Value) -> str:
        if isinstance(value, Var):
            type_id = getattr(value.type, "id", 0)
            return f"var:{value.name}:{type_id}"
        if isinstance(value, Literal):
            type_id = getattr(value.type, "id", 0)
            return f"lit:{type_id}:{value.value!r}"
        if isinstance(value, tuple):
            return f"tuple:({','.join(self._value_shape(item) for item in value)})"
        if value is None:
            return "none"
        return f"const:{value!r}"
