from __future__ import annotations
from collections import defaultdict
from dataclasses import dataclass, field
from pathlib import Path
from typing import Iterable, Iterator, Sequence as TypingSequence, cast


from relationalai.util.error import Part, Source, err, exc

from ...util.structures import OrderedSet

from .metamodel import Aggregate, Construct, Field, Logical, Loop, Lookup, Model, Node, Relation, Require, ScalarType, Sequence, Task, Update, Var, Not, Match, Union, Break, UnionType, Value
from .rewriter import NO_WALK, Walker, Rewriter
from .builtins import builtins as b, get_supertypes_transitive
from ..experimental.loopy.builtins import empty_relation

from contextlib import contextmanager

#------------------------------------------------------
# Constants
#------------------------------------------------------

NORMAL_ORDER = {
    Lookup: 0,
    Aggregate: 0,
    Union: 0,
    Match: 0,
    Logical: 0,
    Require: 0,
    Not: 1,
    Construct: 0,
    Update: 0,
}

#------------------------------------------------------
# Analysis
#------------------------------------------------------

@dataclass
class Analysis:
    fds: dict[Field, tuple[Field, ...]] = field(default_factory=dict)
    mandatory: set[Relation] = field(default_factory=set)

    def is_property(self, relation: Relation) -> bool:
        if not relation.fields:
            return False
        return len(self.fds.get(relation.fields[-1], ())) == len(relation.fields) - 1

#------------------------------------------------------
# Analyzer
#------------------------------------------------------

class Analyzer(Walker):

    def analyze(self, model:Model) -> Analysis:
        self.analysis = Analysis()
        self(model.root)
        object.__setattr__(model, "_analysis", self.analysis)
        return self.analysis

    def require(self, node:Require):
        if not isinstance(node.check, Logical):
            return

        for item in node.check.body:
            if not isinstance(item, Lookup):
                continue
            if item.relation is b.constraints.unique_fields:
                self.track_unique(item)
            # elif item.relation == b.core["mandatory"]:
            #     self.track_mandatory(item)

    def track_unique(self, lookup: Lookup):
        args = cast(list[Field], lookup.args[0])
        if not isinstance(args, (list, tuple)) or any(not isinstance(arg, Field) for arg in args):
            return

        rel = args[0]._relation
        determined = [field for field in rel.fields if field not in args]
        if len(determined) == 1:
            self.analysis.fds[determined[0]] = tuple(args)

    def track_mandatory(self, lookup: Lookup):
        pass

#------------------------------------------------------
# Var finder
#------------------------------------------------------

class VarFinder(Walker):
    def find_vars(self, node:Node|Iterable[Node]|Iterable[Value], positive_only=False) -> OrderedSet[Var]:
        self.vars:OrderedSet[Var] = OrderedSet()
        self.positive_only = positive_only
        self(node)
        return self.vars

    def enter_not(self, node:Node):
        if self.positive_only:
            return NO_WALK

    def var(self, node:Var):
        self.vars.add(node)

    def lookup(self, node:Lookup):
        pass

class ContainerVarFinder(Walker):
    def find_vars(self, node:Node|Iterable[Node]) -> tuple[set[Var],set[Var]]:
        self.noncontainer_vars:set[Var] = set()
        self.container_vars:set[Var] = set()
        self.containers_in_scope = 0
        self(node)
        return self.noncontainer_vars, self.container_vars

    def var(self, node:Var):
        if self.containers_in_scope == 0:
            self.noncontainer_vars.add(node)
        else:
            self.container_vars.add(node)

    def enter_not(self, _:Node):
        return NO_WALK

    def enter_logical(self, _:Logical):
        self.containers_in_scope += 1

    def logical(self, _:Logical):
        self.containers_in_scope -= 1

    def enter_union(self, _:Union):
        self.containers_in_scope += 1

    def union(self, _:Union):
        self.containers_in_scope -= 1

    def enter_match(self, _:Match):
        self.containers_in_scope += 1

    def match(self, _:Match):
        self.containers_in_scope -= 1

#------------------------------------------------------
# Redundant Population Filter
#------------------------------------------------------

class RedundantPopulationFilter(Walker):
    """
    Identifies and removes redundant population lookups during query normalization.

    Population lookups may be redundant if the variable is used elsewhere with a more
    specific or equally specific type constraint. This filter performs a conservative
    analysis to identify which population lookups can safely be removed without changing
    query semantics.

    E.g. Cat(c::Cat), name(c::Cat, n::String) --> Cat(c::Cat) is redundant

    A population lookup CANNOT be removed if:
    1. No other lookup references the variable (would leave the variable unsupported)
    2. The declared field type is a supertype of the population type (population narrows)
       Example: Field type is Animal, population is Cat --> population narrows, keep it
    3. The field type is a heterogeneous union (population narrows the possibilities)
       Example: Field type is Cat|Dog, population is Cat --> population narrows, keep it
    4. The population extends AnyEntity, because that needs to remain as a filter.

    Edge case optimization: If the field type is a union where ALL types are subtypes
    of the population type (e.g., field is Cat|Kitten, population is Cat), we keep the
    population as a candidate since it doesn't actually narrow.
    """
    def find_removable_populations(self, logical: Logical):
        """
        Identify population lookups that are redundant and can be safely removed.

        Args:
            logical: The logical block to analyze for redundant population lookups

        Returns:
            Set of population lookups that are redundant and safe to remove

        Algorithm:
            1. Collect all population lookups (where relation is a ScalarType but it does
               not extend AnyEntity).
            2. Track all vars used in non-population lookups
            3. Walk through the logical body, checking each lookup's field types against
               the population types of its variable arguments
            4. Remove populations from candidates if they provide type narrowing
            5. Return only populations for vars that have other lookup references
        """
        self.candidate_populations: dict[Var, Lookup] = {}
        referenced_vars: set[Var] = set()
        for node in logical.body:
            if isinstance(node, Lookup) and node.relation not in b.core:
                if type(node.relation) is ScalarType:
                    # this is a population lookup for the var with that type; note that this
                    # will not include population lookups of sub or super types, nor lookups
                    # to tables (hence checking type() instead of isinstance())
                    assert isinstance(node.args[0], Var)
                    if node.args[0].type == node.relation:
                        self.candidate_populations[node.args[0]] = node
                else:
                    referenced_vars.update([arg for arg in node.args if isinstance(arg, Var)])

        self(logical)

        return set([pop for pop in self.candidate_populations.values() if pop.args[0] in referenced_vars])

    def lookup(self, node: Lookup):
        """
        Analyze a lookup to determine if associated population lookups can be removed.

        For each variable argument in this lookup, check if its population lookup is
        redundant based on the field's declared type. Remove from candidates if the
        population provides type narrowing that would be lost.
        """
        # Skip builtin lookups and any population lookups themselves
        if (
            node.relation in b.core or
            type(node.relation) is ScalarType
        ):
            return
        for i, arg in enumerate(node.args):
            if not isinstance(arg, Var) or arg not in self.candidate_populations:
                continue
            field_type = node.relation.fields[i].type
            if field_type is arg.type:
                # this lookup is binding the var to exactly the same type as the population,
                # so the population is still a candidate to be removed
                continue
            elif isinstance(field_type, UnionType):
                population_type = self.candidate_populations[arg].relation.fields[0].type
                for t in field_type.types:
                    assert type(t) is ScalarType
                    supertypes = get_supertypes_transitive(t)
                    if (t is not population_type) and (population_type not in supertypes):
                        del self.candidate_populations[arg]
                        break
            elif type(field_type) is ScalarType and isinstance(arg.type, ScalarType):
                if field_type is not b.core.Any and field_type is not b.core.AnyEntity:
                    supertypes = get_supertypes_transitive(arg.type)
                    assert field_type in supertypes, f"{node}[{i}]: {field_type} vs {', '.join(str(x) for x in supertypes)}"
                del self.candidate_populations[arg]


#------------------------------------------------------
# Context
#------------------------------------------------------

class Frame:
    def __init__(self):
        self.nodes:OrderedSet[Task] = OrderedSet()
        self.post_nodes:OrderedSet[Task] = OrderedSet()
        self.indirect_filters:dict[Var, OrderedSet[Task]] = defaultdict(OrderedSet[Task])

class Context:
    def __init__(self):
        self.stack:list[Frame] = [Frame()]

    #------------------------------------------------------
    # subcontext
    #------------------------------------------------------

    def push(self):
        f = Frame()
        self.stack.append(f)
        return f

    def pop(self) -> Frame:
        prev = self.stack.pop()
        return prev

    #------------------------------------------------------
    # Task capture
    #------------------------------------------------------

    def _check_frames(self, node:Task):
        for frame in reversed(self.stack[:-1]):
            if node in frame.nodes:
                return True
        return False

    def add(self, node:Task):
        if self._check_frames(node):
            return
        self.stack[-1].nodes.add(node)

    def extend(self, nodes:TypingSequence[Task]):
        for node in nodes:
            self.add(node)

    def add_post(self, node:Task):
        if self._check_frames(node):
            return
        self.stack[-1].post_nodes.add(node)

    @property
    def nodes(self):
        return self.stack[-1].nodes

    @property
    def stack_nodes(self) -> Iterator[Task]:
        for frame in self.stack:
            yield from frame.nodes

    @property
    def frame(self) -> Frame:
        return self.stack[-1]

#------------------------------------------------------
# RuleBuilder
#------------------------------------------------------

class RuleBuilder(Walker):
    var_finder = VarFinder()

    def __init__(self, unnest: bool = True):
        super().__init__()
        self.unnest = unnest

    def build(self, ctx: Context, node:Node|Iterable[Node], root_frame:Frame) -> list[Task]:
        self.ctx = ctx
        self.node_stack = []
        self.root_frame = root_frame
        self.ctx.push()
        self.rules = []
        self(node)
        final = self.ctx.pop()
        if not self.unnest:
            self.rules.extend(final.nodes)
        return self.rules

    #------------------------------------------------------
    # Context handling
    #------------------------------------------------------

    def push(self, node:Node):
        self.node_stack.append(node)
        return self.ctx.push()

    def pop(self):
        self.node_stack.pop()
        return self.ctx.pop()

    def cur_node(self) -> Node|None:
        if len(self.node_stack) < 1:
            return None
        return self.node_stack[-1]

    def satisfy_indirects(self):
        for node in reversed(self.node_stack):
            # We should never push indirect filters into Nots as it
            # changes the actual logic of the query
            if isinstance(node, (Not)):
                return
        f = self.ctx.frame
        handled = set()
        if f.indirect_filters:
            for var in self.var_finder.find_vars(f.nodes):
                tasks = f.indirect_filters.get(var, ())
                for task in tasks:
                    if task not in handled and task not in self.node_stack:
                        handled.add(task)
                        self(task)

    #------------------------------------------------------
    # Relation ops
    #------------------------------------------------------

    def lookup(self, node:Lookup):
        self.ctx.add(node)

    def construct(self, node:Construct):
        self.ctx.add(node)

    def enter_aggregate(self, node:Aggregate):
        self.push(node)

    def aggregate(self, node:Aggregate):
        collected = self.pop()
        self.ctx.add(node.mut(body=collected.nodes[0]))

    def update(self, update: Update):
        if self.unnest:
            # build a rule for just this update
            body = (*self.ctx.stack_nodes, update)
            self.rules.append(Logical(body=body, scope=True, source=update.source))
        else:
            self.ctx.add(update)

    #------------------------------------------------------
    # Logical
    #------------------------------------------------------

    def enter_logical(self, logical:Logical):
        # Nested logicals that aren't optional and aren't a scope should just be inlined
        if not logical.optional and not logical.scope and isinstance(self.cur_node(), Logical):
            return

        f = self.push(logical)
        f.indirect_filters = self.root_frame.indirect_filters

        # TODO: we should actually toposort this based on var dependencies
        # normalize the body order so that nots come after most ops, followed by constructs
        # and updates
        def sort_key(x):
            order = NORMAL_ORDER.get(type(x), len(NORMAL_ORDER))
            if isinstance(x, Lookup) and isinstance(x.relation, ScalarType):
                order = -1
            if isinstance(x, Logical) and x.optional:
                order = len(NORMAL_ORDER) + 1
            return order

        ordered_body = sorted(logical.body, key=sort_key)
        for item in ordered_body:
            self(item)

        self.satisfy_indirects()
        f = self.pop()
        optional = False
        if not self.unnest:
            optional = logical.optional
        self.ctx.add(Logical(tuple(f.nodes), scope=logical.scope, source=logical.source, optional=optional))
        return NO_WALK

    #------------------------------------------------------
    # Passthrough Containers
    #------------------------------------------------------

    def enter_not(self, node:Not):
        self.push(node)

    def not_(self, node:Not):
        cur = self.pop()
        self.ctx.add(node.mut(task=cur.nodes[0], source=node.source))

    def enter_match(self, match:Match):
        self.push(match)

    def match(self, match:Match):
        cur = self.pop()
        self.ctx.add(match.mut(tasks=tuple(cur.nodes), source=match.source))

    def enter_union(self, union:Union):
        self.push(union)

    def union(self, union:Union):
        cur = self.pop()
        self.ctx.add(union.mut(tasks=tuple(cur.nodes), source=union.source))

    def enter_require(self, require:Require):
        self.push(require)

        self.ctx.push()
        self(require.domain)
        domain = self.ctx.pop().nodes[0]

        self.ctx.push()
        self(require.check)
        check = self.ctx.pop().nodes[0]

        self.pop()
        self.ctx.add(require.mut(domain=domain, check=check, source=require.source))
        return NO_WALK

    # TODO: preserve reordering and unwrapping invariants, but still normalize

    def enter_loop(self, node:Loop):
        self.ctx.add(node)
        return NO_WALK

    def enter_break(self, node:Break):
        self.ctx.add(node)
        return NO_WALK

    def enter_sequence(self, node:Sequence):
        self.ctx.add(node)
        return NO_WALK


#------------------------------------------------------
# Lowering
#------------------------------------------------------

class Lowering(Rewriter):
    def normalize(self, node: Task) -> Task:
        self.concept_to_propagator: dict[ScalarType, Relation] = {}

        new_node = self(node)
        assert isinstance(new_node, Logical)
        new_body = list(new_node.body)

        handled = set()
        worklist = list(self.concept_to_propagator.keys())
        while worklist:
            t = worklist.pop()
            if t in handled:
                continue
            handled.add(t)

            prop = self.concept_to_propagator[t]
            args = (Var(t, "x"),)
            inner_body = [Lookup(prop, args), Update(t, args)]
            for st in t.super_types:
                if st is b.core.AnyEntity:
                    continue
                worklist.append(st)
                st_prop = self.concept_to_propagator.get(st, None)
                if not st_prop:
                    st_prop = self._concept_to_propagator(st)
                    self.concept_to_propagator[st] = st_prop
                inner_body.append(Update(st_prop, args))
            new_body.append(Logical(tuple(inner_body), scope=True))

        return Logical(tuple(new_body), scope=new_node.scope, annotations=new_node.annotations, source=new_node.source)

    def _concept_to_propagator(self, concept: ScalarType):
        prop = self.concept_to_propagator.get(concept, None)
        if not prop:
            prop = Relation(f"new_{concept.name}", (Field("x", concept),), source=concept.source)
            self.concept_to_propagator[concept] = prop
        return prop

    def logical(self, logical: Logical):
        new_body = list(logical.body)
        changed = False

        # Handle flow of values to supertypes
        constructed_vars, pop_updates = set(), set()
        for item in new_body:
            if isinstance(item, Construct):
                constructed_vars.add(item.id_var)
            elif isinstance(item, Update) and type(item.relation) is ScalarType:
                pop_updates.add(item)

        for item in pop_updates:
            assert len(item.args) == 1 and isinstance(item.args[0], Var)
            var = item.args[0]
            if var not in constructed_vars:
                continue
            if not var.type.super_types:
                continue
            new_body.remove(item)
            new_body.append(Update(self._concept_to_propagator(var.type), (var,)))
            changed = True

        # Inline Updates from optional-only-Update child logicals
        inlined_body = []
        for item in new_body:
            if not (isinstance(item, Logical) and all(isinstance(t, Update) for t in item.body)):
                inlined_body.append(item)
                continue
            inlined_body.extend(item.body)
            changed = True
        new_body = inlined_body

        # Remove redundant population lookups
        temp = logical.mut(body=tuple(new_body)) if changed else logical
        lookups_to_remove = RedundantPopulationFilter().find_removable_populations(temp)
        if lookups_to_remove:
            new_body = [x for x in new_body if x not in lookups_to_remove]
            changed = True

        return logical.mut(body=tuple(new_body)) if changed else None


#------------------------------------------------------
# RecursionChecker
#------------------------------------------------------

class RecursionChecker(Walker):
    """
    Builds a dependency map of Update relations to Lookup relations that appear
    in the same scope=True Logical block.

    For each scope=True Logical, collects the relations written via Update and
    the relations read via Lookup at that scope level (not crossing into nested
    scope=True Logicals). The result maps each updated relation to the set of
    lookup relations used alongside it.

    A relation is recursive if it appears in both the key and its own value set.
    """

    def check(self, node: Task) -> bool:
        """
        Walk the normalized task, print any detected recursion cycles, and
        return whether any recursion was found.

        Returns:
            True if at least one relation is recursive (appears in its own
            transitive dependency set), False otherwise.
        """
        self.deps: dict[Relation, set[Relation]] = defaultdict(set)
        self.scope_stack: list[tuple[list[Update], list[Lookup]]] = []
        self(node)

        self.direct_deps = {}
        for rel, rs in self.deps.items():
            self.direct_deps[rel] = rs.copy()

        # Transitive closure of self.deps
        changed = True
        while changed:
            changed = False
            for _, deps in self.deps.items():
                new_deps = set()
                for dep in deps:
                    if dep in self.deps:
                        new_deps.update(self.deps[dep])
                before = len(deps)
                deps.update(new_deps)
                if len(deps) > before:
                    changed = True

        has_recursion = False
        for rel, rs in self.deps.items():
            if rel in rs:
                has_recursion = True
                if cycle := self._print_rec(rel, rel, {rel}):
                    print(f"{' <- '.join(str(x) for x in ([rel] + cycle))}")
        return has_recursion

    def _print_rec(self, rel, r, visited):
        """DFS over direct_deps to find a path from r back to rel.
        Returns the list of nodes from r's dependency onward that closes the cycle,
        or None if no cycle is reachable."""
        rs = self.direct_deps.get(r, None)
        if not rs:
            return None
        for a in rs:
            if a == rel:
                return [a]
            if a not in visited:
                visited.add(a)
                if result := self._print_rec(rel, a, visited):
                    return [a] + result
        return None

    def enter_logical(self, logical: Logical):
        if logical.scope:
            self.scope_stack.append(([], []))

    def logical(self, logical: Logical):
        if logical.scope:
            updates, lookups = self.scope_stack.pop()
            lookup_relations = {lookup.relation for lookup in lookups}
            for update in updates:
                self.deps[update.relation].update(lookup_relations)

    def update(self, node: Update):
        if self.scope_stack:
            self.scope_stack[-1][0].append(node)

    def lookup(self, node: Lookup):
        if node.relation in b.core:
            return
        if self.scope_stack:
            self.scope_stack[-1][1].append(node)


#------------------------------------------------------
# Normalize
#------------------------------------------------------

class Normalize(Walker):
    var_finder = VarFinder()
    rule_builder = RuleBuilder(unnest=False)

    def normalize(self, node: Task) -> Task:
        self.ctx = Context()
        self(node)
        return self.ctx.nodes[0]

    #------------------------------------------------------
    # Logical
    #------------------------------------------------------

    def enter_logical(self, logical: Logical):
        if logical.optional:
            self.ctx.add_post(logical)
            return NO_WALK

        self.ctx.push()

    def logical(self, logical: Logical):
        cur = self.ctx.pop()
        parent = self.ctx.frame

        # We need to build our nodes only in the context of our parent frame
        cleaned = self.rule_builder.build(self.ctx, Logical(tuple(cur.nodes)), parent)
        cleaned_nodes = []
        if cleaned:
            assert isinstance(cleaned[0], Logical)
            cleaned_nodes.extend(cleaned[0].body)

        # we want to build the post nodes in the context of the current frame, so
        # we temporarily push it back on the stack
        self.ctx.push()
        for cleaned_node in cleaned_nodes:
            self.ctx.add(cleaned_node)
        if logical.scope:
            for node in cur.post_nodes:
                cleaned_nodes.extend(self.rule_builder.build(self.ctx, node, cur))

        # make sure the frame is gone
        self.ctx.pop()

        # get vars used outside a container, and vars used in containers
        noncontainer_vs, container_vs = ContainerVarFinder().find_vars(logical.body)
        filtered_nodes = []
        for cleaned_node in cleaned_nodes:
            if isinstance(cleaned_node, Not):
                vs = self.var_finder.find_vars(cleaned_node)
                # the vars used in a negation, are also only used in a container.
                # the negation will move to the container, so no need to keep it here.
                if vs.isdisjoint(noncontainer_vs) and not vs.isdisjoint(container_vs):
                    continue
            filtered_nodes.append(cleaned_node)
        cleaned_nodes = filtered_nodes

        self.ctx.add(Logical(tuple(cleaned_nodes), scope=logical.scope, annotations=logical.annotations, source=logical.source))

    #------------------------------------------------------
    # Require
    #------------------------------------------------------

    def enter_require(self, node: Require):
        self.ctx.push()

    def require(self, node: Require):
        cur = self.ctx.pop()
        self.ctx.extend(self.rule_builder.build(self.ctx, node, cur))

    #------------------------------------------------------
    # Relation ops
    #------------------------------------------------------

    def lookup(self, node:Lookup):
        self.ctx.add(node)

    def construct(self, node:Construct):
        self.ctx.add(node)

    def enter_aggregate(self, node:Aggregate):
        self.ctx.add(node)
        return NO_WALK

    def update(self, update: Update):
        self.ctx.add(update)

    #------------------------------------------------------
    # Indirect filters
    #------------------------------------------------------

    def enter_not(self, node:Not):
        self.ctx.add(node)
        vs = self.var_finder.find_vars(node.task)
        for v in vs:
            self.ctx.frame.indirect_filters[v].add(node)
        return NO_WALK

    def enter_match(self, match:Match):
        self.ctx.add(match)
        return NO_WALK

    def enter_union(self, union:Union):
        self.ctx.add(union)
        return NO_WALK

    # TODO: preserve reordering and unwrapping invariants, but still normalize

    def enter_loop(self, node:Loop):
        normalized_body = Normalize().normalize(node.body)
        self.ctx.add(node.mut(body=normalized_body))
        return NO_WALK

    def enter_break(self, node:Break):
        normalized_check = Normalize().normalize(node.check)
        self.ctx.add(node.mut(check=normalized_check))
        return NO_WALK

    def enter_sequence(self, node:Sequence):
        normalized_tasks = []
        for task in node.tasks:
            normalized_tasks.append(Normalize().normalize(task))
        self.ctx.add(node.mut(tasks=tuple(normalized_tasks)))
        return NO_WALK

#------------------------------------------------------
# SplittingNormalize
#------------------------------------------------------

class SplittingNormalize(Walker):
    var_finder = VarFinder()
    rule_builder = RuleBuilder()

    def flatten(self, node: Node) -> list[Task]:
        self.ctx = Context()
        self.rules: list[Task] = []
        self(node)
        return self.rules

    def enter_logical(self, logical: Logical):
        if logical.optional:
            self.ctx.add_post(logical)
            return NO_WALK
        self.ctx.push()

    def logical(self, logical: Logical):
        if logical.scope:
            cur = self.ctx.pop()
            if len(cur.post_nodes) > 1:
                self._split_rules(cur, logical)
            else:
                post = cur.post_nodes
                first = next(iter(post), None)
                if isinstance(first, Logical):
                    post = first.body
                cur_nodes = (*cur.nodes, *post)
                self.rules.extend(self.rule_builder.build(self.ctx, Logical(cur_nodes), cur))

    def _split_rules(self, cur:Frame, logical:Logical):
        # determine vars that cross the shared root boundary
        vs = self.var_finder.find_vars(list(cur.nodes), positive_only=True)
        needed = OrderedSet[Var]()
        for item in cur.post_nodes:
            item_vars = self.var_finder.find_vars(item)
            for v in item_vars:
                if v in vs:
                    needed.add(v)

        # generate an intermediate to carry the needed vars
        intermediate = Relation("temp", tuple(Field(f.name, f.type) for f in needed), source=logical.source)

        # add the trunk rule
        cur_nodes = (*cur.nodes, Update(intermediate, tuple(needed)))
        base = self.rule_builder.build(self.ctx, Logical(cur_nodes, source=logical.source), cur)
        self.rules.extend(base)
        print("-------------------------------------------------")

        # add the child rules
        self.ctx.push()
        self.ctx.add(Lookup(intermediate, tuple(needed)))

        for node in cur.post_nodes:
            self.rules.extend(self.rule_builder.build(self.ctx, node, cur))

        self.ctx.pop()


    #------------------------------------------------------
    # Relation ops
    #------------------------------------------------------

    def lookup(self, node:Lookup):
        self.ctx.add(node)

    def construct(self, node:Construct):
        self.ctx.add(node)

    def enter_aggregate(self, node:Aggregate):
        self.ctx.add(node)
        return NO_WALK

    def update(self, update: Update):
        self.ctx.add_post(update)

    #------------------------------------------------------
    # Indirect filters
    #------------------------------------------------------

    def enter_not(self, node:Not):
        self.ctx.add(node)
        vs = self.var_finder.find_vars(node.task)
        for v in vs:
            self.ctx.frame.indirect_filters[v].add(node)
        return NO_WALK

    def enter_match(self, match:Match):
        self.ctx.add(match)
        vs = self.var_finder.find_vars(match.tasks)
        for v in vs:
            self.ctx.frame.indirect_filters[v].add(match)
        return NO_WALK

    # TODO: preserve reordering and unwrapping invariants, but still normalize

    def enter_loop(self, node:Loop):
        normalized_body = Normalize().normalize(node.body)
        self.ctx.add(node.mut(body=normalized_body))
        return NO_WALK

    def enter_break(self, node:Break):
        normalized_check = Normalize().normalize(node.check)
        self.ctx.add(node.mut(check=normalized_check))
        return NO_WALK

    def enter_sequence(self, node:Sequence):
        normalized_tasks = []
        for task in node.tasks:
            normalized_tasks.append(Normalize().normalize(task))
        self.ctx.add(node.mut(tasks=tuple(normalized_tasks)))
        return NO_WALK

    def exists(self, node):
        print("EXISTS")

#------------------------------------------------------
# Groundness Analysis - equality-output extension (temporary workaround)
#------------------------------------------------------

class _EqualityCollector(Walker):
    """Collect (left, right) Var pairs from equality lookups, skipping isolating scopes."""

    def collect(self, tasks) -> list[tuple[Var, Var]]:
        self._pairs: list[tuple[Var, Var]] = []
        self(tasks)
        return self._pairs

    def lookup(self, node: Lookup):
        if node.relation == b.core.eq and len(node.args) == 2:
            left, right = node.args
            if isinstance(left, Var) and isinstance(right, Var):
                self._pairs.append((left, right))

    # Do not cross Not or nested Match/Union boundaries (equalities inside those
    # scopes don't flow to the outer Match/Union output)
    def enter_not(self, _: Not):     return NO_WALK
    def enter_match(self, _: Match): return NO_WALK
    def enter_union(self, _: Union): return NO_WALK
    # Logical boundaries (scope=True, optional) are intentionally NOT blocked:
    # the branch itself may be a scope=True Logical and we still need to find
    # equalities inside it.


def _extend_outputs_via_equality(task: Match|Union) -> set[Var]:
    """
    TEMPORARY WORKAROUND: extend Match/Union outputs to include variables unified
    with an output variable via equality (=(v, x)) inside a branch.

    This compensates for a bug in the metamodel generator that omits such variables
    from the .outputs set.  Remove this function (and its two call sites in
    Groundness.enter_match / enter_union) once the generator bug is fixed.
    """
    pairs = _EqualityCollector().collect(list(task.tasks))
    # Keyed by var.id — never use == on nodes directly.
    extended: dict[int, Var] = {v.id: v for v in task.outputs}

    # Fixed-point: chains like x=v, y=x where v is output → y is also output
    changed = True
    while changed:
        changed = False
        for left, right in pairs:
            if left.id in extended and right.id not in extended:
                extended[right.id] = right
                changed = True
            elif right.id in extended and left.id not in extended:
                extended[left.id] = left
                changed = True

    return set(extended.values())

#--------------------------------------------------
# Groundness Testing
#--------------------------------------------------
UNGROUND_VARS = None
@contextmanager
def groundness_error(var_name: str):
    """
    Check that this type error is raised within the context.

    :param var_name: the name of the variable expected to cause a groundness error.
    """
    global UNGROUND_VARS
    UNGROUND_VARS = []
    try:
        yield
    except Exception:
        # we expect the groundness analyzer to raise an error
        pass
    finally:
        if var_name not in UNGROUND_VARS:
            exc("Expected Error", f"Expected groundness error for variable '{var_name}' but none was raised.")
        UNGROUND_VARS = None

#------------------------------------------------------
# Groundness Analysis
#------------------------------------------------------

class Groundness(Walker):
    """
    Analyzes variable groundness in metamodel tasks.

    A variable is "grounded" (bound) when it's assigned a value. Variables must be
    grounded before being used in certain contexts (Update, Construct, etc.).

    Scoping rules:
    - Variables bound in Not blocks don't escape to parent scope
    - Variables bound in scope=True containers don't escape to parent scope
    - Variables in regular Logical blocks unify across all tasks
    """
    var_finder = VarFinder()

    def enforce(self, task: Task):
        """
        Analyze the task for groundness violations and raise errors if any are found.

        This method calls analyze() and then provides nice error messages with source locations
        for any ungrounded variables.

        Args:
            task: The task to analyze and enforce groundness on

        Raises:
            Exception via exc() if any ungrounded variables are found
        """
        # Analyze the task to find ungrounded variables
        vars = self.analyze(task)

        if UNGROUND_VARS is not None and vars:
            # If we're in a testing context, record the ungrounded variable names instead of raising errors
            for key in vars.keys():
                UNGROUND_VARS.append(key.name)
            # raise to avoid further processing
            raise Exception("Groundness errors detected.")

        if vars:
            # For each ungrounded variable, build an error with source locations
            for key, val in vars.items():
                parts: list[Part] = []

                if key.source and Path(key.source.file).exists():
                    # Show where the variable was defined
                    parts.append("Variable defined at:")
                    parts.append(Source(key.source))

                # Collect unique source locations where variable is used ungrounded
                # Use a dict keyed by "line-file" to deduplicate sources
                uses = dict()
                for v in val:
                    if v.source:
                        uses[f"{v.source.line:04d}-{v.source.file}"] = Source(v.source)
                    else:
                        uses[v.id] = None

                # Show each unique ungrounded use location
                for u in uses:
                    if uses[u] is not None:
                        parts.append("Unground use at:")
                        parts.append(uses[u])

                    # Report the error for this variable
                    err("Unground Variable", f"Variable '{key.name}' is unground.", parts)

            # Raise an exception indicating groundness violations were found
            exc("Unground Variables", "Ungrounded variables found (see above for details).")

    def analyze(self, task: Task) -> dict[Var, list[Task]]:
        """
        Analyze a task to find ungrounded variables.

        Args:
            task: The task to analyze

        Returns:
            Dict mapping ungrounded variables to the tasks where they're used ungrounded
        """
        self.grounded: set[Var] = set()
        self.ungrounded_uses: dict[Var, list[Task]] = defaultdict(list)
        self.in_isolating_scope = 0  # Track depth of Not blocks or scope=True containers
        self.saved_grounded_stack: list[set[Var]] = []  # Stack for nested scopes

        self(task)

        return dict(self.ungrounded_uses)

    def _check_grounded(self, vars_needed: Iterable[Var], task: Task):
        """Check if variables are grounded, record violations."""
        for v in vars_needed:
            if v not in self.grounded:
                self.ungrounded_uses[v].append(task)

    def _ground_vars(self, vars: set[Var]):
        """Mark variables as grounded."""
        self.grounded.update(vars)

    #------------------------------------------------------
    # Logical - variables unify across all tasks
    #------------------------------------------------------

    def enter_logical(self, logical: Logical):
        # Optional logicals and scope=True act as upward scope boundaries
        # Variables bound inside don't escape upward
        if logical.scope or logical.optional:
            self.saved_grounded_stack.append(self.grounded.copy())
            self.in_isolating_scope += 1

        # First pass: collect all variables that will be grounded in this logical block
        # This is necessary because Logical blocks have order independence - all tasks
        # execute together, so we need to know all grounded variables upfront
        #
        # We use a fixed-point iteration because equality can ground variables
        # bidirectionally, creating chains: if A grounds B via =, and B grounds C via =,
        # then we need multiple passes to ground all variables
        grounded_here = set()

        # Keep iterating until no new variables are grounded
        changed = True
        while changed:
            changed = False
            prev_size = len(grounded_here)

            for task in logical.body:
                if isinstance(task, Lookup):
                    # Special case: equality (=) can ground variables bidirectionally
                    if task.relation.id == b.core.eq.id and len(task.args) == 2:
                        left, right = task.args
                        left_is_var = isinstance(left, Var)
                        right_is_var = isinstance(right, Var)

                        # Check which args are already grounded (either previously or in outer scope)
                        left_grounded = not left_is_var or left in self.grounded or left in grounded_here
                        right_grounded = not right_is_var or right in self.grounded or right in grounded_here

                        # Ground ungrounded vars if the other side is grounded
                        if left_grounded and right_is_var and right not in grounded_here:
                            grounded_here.add(right)
                        if right_grounded and left_is_var and left not in grounded_here:
                            grounded_here.add(left)
                    else:
                        # Normal lookup: output fields ground their variables
                        for i, arg in enumerate(task.args):
                            if isinstance(arg, Var) and not task.relation.fields[i].input:
                                grounded_here.add(arg)
                elif isinstance(task, Aggregate):
                    # Aggregate body grounds variables, projection and group specify what escapes
                    for var in task.projection:
                        if isinstance(var, Var):
                            grounded_here.add(var)
                    for var in task.group:
                        if isinstance(var, Var):
                            grounded_here.add(var)
                elif isinstance(task, Construct):
                    # id_var is grounded by Construct
                    if isinstance(task.id_var, Var):
                        grounded_here.add(task.id_var)
                elif isinstance(task, (Match, Union)):
                    # TEMPORARY: workaround for metamodel generator bug — see
                    # _extend_outputs_via_equality for details.  Remove this
                    # branch when the generator bug is fixed.
                    for var in _extend_outputs_via_equality(task):
                        grounded_here.add(var)

            # Check if we grounded any new variables in this iteration
            if len(grounded_here) > prev_size:
                changed = True

        # Mark these variables as grounded (will be ignored if in_isolating_scope > 0)
        self._ground_vars(grounded_here)

    def logical(self, logical: Logical):
        # Exit isolating scope for optional or scope=True logicals
        if logical.scope or logical.optional:
            self.in_isolating_scope -= 1
            # Variables bound in scope don't escape - restore saved state
            self.grounded = self.saved_grounded_stack.pop()

    #------------------------------------------------------
    # Not - variables don't escape
    #------------------------------------------------------

    def enter_not(self, not_task: Not):
        self.in_isolating_scope += 1
        # Save current grounded set on stack
        self.saved_grounded_stack.append(self.grounded.copy())

    def not_(self, not_task: Not):
        self.in_isolating_scope -= 1
        # Variables bound in Not don't escape - restore saved state
        self.grounded = self.saved_grounded_stack.pop()

    #------------------------------------------------------
    # Lookup - check input fields, ground output fields
    #------------------------------------------------------

    def lookup(self, lookup: Lookup):
        # Special case: equality (=) has unification semantics
        # Even though both fields are marked as input, either arg can be bound if the other is grounded
        if lookup.relation.id == b.core.eq.id and len(lookup.args) == 2:
            left, right = lookup.args
            left_is_var = isinstance(left, Var)
            right_is_var = isinstance(right, Var)

            # If both are vars, at least one must be grounded
            if left_is_var and right_is_var:
                left_grounded = left in self.grounded
                right_grounded = right in self.grounded

                if left_grounded and not right_grounded:
                    # Left grounded, bind right to left
                    self._ground_vars({right})
                elif right_grounded and not left_grounded:
                    # Right grounded, bind left to right
                    self._ground_vars({left})
                elif not left_grounded and not right_grounded:
                    # Neither grounded - error
                    self._check_grounded({left, right}, lookup)
                # else: both grounded - equality check, no action needed
            elif left_is_var and not right_is_var:
                # Right is not a var (literal/constant), grounds left
                self._ground_vars({left})
            elif right_is_var and not left_is_var:
                # Left is not a var (literal/constant), grounds right
                self._ground_vars({right})
            # else: neither is a var - just a constant equality check
            return

        # Normal lookup: input fields must be grounded, output fields get grounded
        input_vars = set()
        output_vars = set()

        for i, arg in enumerate(lookup.args):
            if isinstance(arg, Var):
                if lookup.relation.fields[i].input:
                    input_vars.add(arg)
                else:
                    output_vars.add(arg)

        self._check_grounded(input_vars, lookup)
        self._ground_vars(output_vars)

    #------------------------------------------------------
    # Update - all variables must be grounded
    #------------------------------------------------------

    def update(self, update: Update):
        # Skip groundness check for @empty annotated updates - variables are just placeholders
        for annotation in update.annotations:
            if annotation.relation is empty_relation:
                return

        vars_used = self.var_finder.find_vars(update)
        self._check_grounded(vars_used, update)

    #------------------------------------------------------
    # Construct - grounds id_var, requires values to be grounded
    #------------------------------------------------------

    def construct(self, construct: Construct):
        # id_var is the output variable (gets bound by construct)
        if isinstance(construct.id_var, Var):
            self._ground_vars({construct.id_var})

        # Check that all values (to be hashed) are grounded
        if construct.values:
            value_vars = self.var_finder.find_vars(construct.values)
            self._check_grounded(value_vars, construct)

    #------------------------------------------------------
    # Aggregate - body variables must be grounded
    #------------------------------------------------------

    def enter_aggregate(self, aggregate: Aggregate):
        # Save state - variables grounded inside (including aggregation outputs)
        # should not escape unless they're in projection
        self.saved_grounded_stack.append(self.grounded.copy())

    def aggregate(self, aggregate: Aggregate):
        # After walking body, check that:
        # 1. Input args to aggregation function are grounded (by body)
        # 2. Aggregation output args get grounded (within Aggregate)
        # 3. Projection variables are grounded

        # Check aggregation input field args are grounded, ground output field args
        for i, arg in enumerate(aggregate.args):
            if isinstance(arg, Var) and i < len(aggregate.aggregation.fields):
                if aggregate.aggregation.fields[i].input:
                    if arg not in self.grounded:
                        self.ungrounded_uses[arg].append(aggregate)
                else:
                    # Output fields from aggregation ground their args (within Aggregate scope)
                    self._ground_vars({arg})

        # Check projection and group variables are grounded (by body or aggregation outputs)
        for var in aggregate.projection:
            if isinstance(var, Var) and var not in self.grounded:
                self.ungrounded_uses[var].append(aggregate)
        for var in aggregate.group:
            if isinstance(var, Var) and var not in self.grounded:
                self.ungrounded_uses[var].append(aggregate)

        # Restore state and ground projection variables + group variables + aggregation output args
        self.grounded = self.saved_grounded_stack.pop()

        # Ground projection and group variables
        self._ground_vars(set(aggregate.projection))
        self._ground_vars(set(aggregate.group))

        # Also ground aggregation output field args (they escape too)
        output_args = set()
        for i, arg in enumerate(aggregate.args):
            if isinstance(arg, Var) and i < len(aggregate.aggregation.fields):
                if not aggregate.aggregation.fields[i].input:
                    output_args.add(arg)
        self._ground_vars(output_args)

    #------------------------------------------------------
    # Containers
    #------------------------------------------------------

    def enter_match(self, match: Match):
        # Each branch must be analyzed independently with only the pre-match grounded
        # state. Without isolation, variables grounded in branch N leak into branch N+1,
        # masking real groundness errors in later branches.
        saved = self.grounded.copy()
        for branch in match.tasks:
            self.grounded = saved.copy()
            self(branch)
        # Restore pre-match state; only .outputs (present in all branches) escape upward
        self.grounded = saved
        self._ground_vars(_extend_outputs_via_equality(match))
        return NO_WALK

    def enter_union(self, union: Union):
        # Same isolation requirement as Match: each branch analyzed independently
        saved = self.grounded.copy()
        for branch in union.tasks:
            self.grounded = saved.copy()
            self(branch)
        # Restore pre-union state; only .outputs escape upward
        self.grounded = saved
        self._ground_vars(_extend_outputs_via_equality(union))
        return NO_WALK

    def enter_sequence(self, sequence: Sequence):
        # Tasks execute sequentially, so variables ground progressively
        pass

    def enter_loop(self, loop: Loop):
        # Loop body is analyzed separately
        pass
