#!/usr/bin/env python3
"""datalog.py -- a self-contained Datalog engine.

Wizard: null    Lane: Codeforge (hard coding and agentic reasoning)

What this is
------------
A bottom-up Datalog engine with *stratified negation*, written against the
standard semantics and nothing else:

  * a hand-written tokenizer + recursive-descent parser for facts and rules,
  * static checks: arity consistency, rule safety (range restriction),
    and stratification (no negation may sit inside a recursive cycle),
  * stratification by Tarjan SCC + Kahn topological order on the condensation,
  * evaluation by *semi-naive* fixpoint iteration inside each stratum, with
    negation resolved against strictly lower (already final) strata,
  * a JSON / text CLI.

Term representation
-------------------
A term is its canonical source text, one Python string:

    '?X'   variable X            (leading '?' marks a variable)
    '?_'   anonymous variable    (each occurrence is independent)
    'foo'  atom constant
    '"a"'  string constant       (kept verbatim, quotes included)
    '-7'   integer constant      (kept verbatim)

Variables in heads and in negated literals must also occur in a positive body
literal (range restriction); a rule with no positive body literal is rejected.
Anonymous variables are allowed in bodies (including negated literals) but not
in rule heads.

Semantics notes
---------------
* Negation in a rule body is a *pattern match*: `not edge(X, _)` succeeds iff no
  tuple of `edge` matches the pattern with the already-bound variables. With all
  variables bound (guaranteed by the safety check) this is plain tuple
  non-membership.
* Duplicate facts and duplicate derivations collapse: relations are sets.
* The result is independent of the order of facts and rules.

Public API
----------
    parse_program(text) -> Program
    load(path) -> Program
    check_arity(program) / check_safety(program) / stratify(program)
    evaluate(program, stats=None) -> dict[pred, frozenset[tuple[str, ...]]]
    format_relations(rels, preds=None) -> str

Errors: ParseError, ArityError, SafetyError, StratificationError, all
subclasses of DatalogError.

Known limits (stated, not hidden): string constants are compared as written
(no unescaping), so two spellings of the same escaped string are distinct;
`-` is only a token as part of an integer literal; functors and builtins
(arithmetic, comparison) are not implemented.
"""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
from dataclasses import dataclass, field

__all__ = [
    "DatalogError",
    "ParseError",
    "ArityError",
    "SafetyError",
    "StratificationError",
    "Literal",
    "Rule",
    "Program",
    "parse_program",
    "load",
    "check_arity",
    "check_safety",
    "stratify",
    "evaluate",
    "format_relations",
    "main",
]


class DatalogError(Exception):
    """Base class for every error this engine raises."""


class ParseError(DatalogError):
    pass


class ArityError(DatalogError):
    pass


class SafetyError(DatalogError):
    pass


class StratificationError(DatalogError):
    pass


# --------------------------------------------------------------------------
# Terms
# --------------------------------------------------------------------------

_ANON = "?_"


def is_var(term: str) -> bool:
    return term.startswith("?")


def is_anon(term: str) -> bool:
    return term == _ANON


def var_name(term: str) -> str:
    return term[1:]


# --------------------------------------------------------------------------
# Syntax tree
# --------------------------------------------------------------------------


@dataclass(frozen=True)
class Literal:
    pred: str
    args: tuple
    negated: bool = False

    def __str__(self) -> str:
        head = "not " if self.negated else ""
        return f"{head}{self.pred}({', '.join(self.args)})"


@dataclass(frozen=True)
class Rule:
    head_pred: str
    head_args: tuple
    body: tuple

    @property
    def positive(self) -> tuple:
        return tuple(lit for lit in self.body if not lit.negated)

    @property
    def negated(self) -> tuple:
        return tuple(lit for lit in self.body if lit.negated)

    def __str__(self) -> str:
        head = f"{self.head_pred}({', '.join(self.head_args)})"
        body = ", ".join(str(lit) for lit in self.body)
        return f"{head} :- {body}."


@dataclass
class Program:
    facts: dict = field(default_factory=dict)   # pred -> set of tuples
    rules: list = field(default_factory=list)   # list[Rule]
    arities: dict = field(default_factory=dict)  # pred -> int

    def predicates(self) -> set:
        preds = set(self.facts)
        preds.update(self.arities)
        for rule in self.rules:
            preds.add(rule.head_pred)
            for lit in rule.body:
                preds.add(lit.pred)
        return preds

    def add_fact(self, pred: str, args: tuple) -> None:
        self.facts.setdefault(pred, set()).add(tuple(args))

    def add_rule(self, rule: Rule) -> None:
        self.rules.append(rule)

    def text(self) -> str:
        lines = []
        for pred in sorted(self.facts):
            for tup in sorted(self.facts[pred]):
                lines.append(f"{pred}({', '.join(tup)}).")
        lines.extend(str(r) for r in self.rules)
        return "\n".join(lines) + "\n"


# --------------------------------------------------------------------------
# Tokenizer + parser
# --------------------------------------------------------------------------

_TOKEN_RE = re.compile(
    r"(?P<ws>[ \t\r\n]+)"
    r"|(?P<comment>%[^\n]*)"
    r"|(?P<string>\"(?:[^\"\\]|\\.)*\")"
    r"|(?P<int>-?\d+)"
    r"|(?P<name>[A-Za-z_][A-Za-z0-9_]*)"
    r"|(?P<punct>:-|[!(),.])"
)


def tokenize(text: str) -> list:
    tokens = []
    pos = 0
    n = len(text)
    while pos < n:
        m = _TOKEN_RE.match(text, pos)
        if m is None:
            line = text.count("\n", 0, pos) + 1
            raise ParseError(
                f"unexpected character {text[pos]!r} at line {line}, column "
                f"{pos - text.rfind(chr(10), 0, pos)}"
            )
        kind = m.lastgroup
        if kind not in ("ws", "comment"):
            tokens.append((kind, m.group()))
        pos = m.end()
    return tokens


class _Parser:
    def __init__(self, tokens):
        self.toks = tokens
        self.i = 0

    def peek(self, k: int = 0):
        j = self.i + k
        return self.toks[j] if j < len(self.toks) else (None, None)

    def next(self):
        tok = self.peek()
        self.i += 1
        return tok

    def expect(self, kind, value=None):
        k, v = self.next()
        if k != kind or (value is not None and v != value):
            want = value if value is not None else kind
            raise ParseError(f"expected {want!r} but found {v!r}")
        return v

    # ---- grammar ------------------------------------------------------

    def parse_term(self):
        kind, value = self.next()
        if kind in ("string", "int"):
            return value
        if kind == "name":
            if value in ("not", "!"):
                raise ParseError(f"{value!r} cannot be used as a term")
            if value[0].isupper() or value[0] == "_":
                return "?" + value
            return value
        raise ParseError(f"expected a term but found {value!r}")

    def parse_literal(self):
        negated = False
        kind, value = self.peek()
        if (kind == "name" and value == "not") or (kind == "punct" and value == "!"):
            self.next()
            negated = True
        pred = self.expect("name")
        if pred in ("not", "!"):
            raise ParseError(f"{pred!r} cannot be used as a predicate name")
        self.expect("punct", "(")
        args = []
        if self.peek() == ("punct", ")"):
            self.next()
        else:
            while True:
                args.append(self.parse_term())
                kind, value = self.peek()
                if (kind, value) == ("punct", ","):
                    self.next()
                    continue
                break
            self.expect("punct", ")")
        return Literal(pred, tuple(args), negated)

    def parse_program(self) -> Program:
        program = Program()
        while self.peek()[0] is not None:
            head = self.parse_literal()
            if head.negated:
                raise ParseError("a clause head cannot be negated")
            if self.peek() == ("punct", ":-"):
                self.next()
                body = [self.parse_literal()]
                while self.peek() == ("punct", ","):
                    self.next()
                    body.append(self.parse_literal())
                self.expect("punct", ".")
                program.add_rule(Rule(head.pred, head.args, tuple(body)))
            else:
                self.expect("punct", ".")
                for term in head.args:
                    if is_var(term):
                        raise ParseError(
                            f"fact {head.pred}/{len(head.args)} contains the "
                            f"variable {term}; facts must be ground"
                        )
                program.add_fact(head.pred, head.args)
        return program


def parse_program(text: str) -> Program:
    """Parse Datalog source text into a Program."""
    return _Parser(tokenize(text)).parse_program()


def load(path: str) -> Program:
    with open(path, "r", encoding="utf-8") as fh:
        return parse_program(fh.read())


# --------------------------------------------------------------------------
# Static checks
# --------------------------------------------------------------------------


def check_arity(program: Program) -> None:
    """Every predicate must be used with one arity throughout."""
    seen = dict(program.arities)

    def note(pred, n, where):
        if pred in seen and seen[pred] != n:
            raise ArityError(
                f"predicate {pred!r} used with arity {n} in {where} "
                f"but arity {seen[pred]} elsewhere"
            )
        seen[pred] = n

    for pred, tuples in program.facts.items():
        for tup in tuples:
            note(pred, len(tup), f"a fact of {pred!r}")
    for rule in program.rules:
        note(rule.head_pred, len(rule.head_args), f"the head of {rule}")
        for lit in rule.body:
            note(lit.pred, len(lit.args), f"the body of {rule}")


def check_safety(program: Program) -> None:
    """Range restriction: head and negated-literal variables must be bound by a
    positive body literal, and every rule needs at least one positive literal."""
    for rule in program.rules:
        if not rule.positive:
            raise SafetyError(f"rule {rule} has no positive body literal (unsafe)")
        bound = set()
        for lit in rule.positive:
            for term in lit.args:
                if is_var(term) and not is_anon(term):
                    bound.add(term)
        for term in rule.head_args:
            if is_anon(term):
                raise SafetyError(f"anonymous variable in the head of {rule}")
            if is_var(term) and term not in bound:
                raise SafetyError(
                    f"head variable {term} of {rule} is not bound by a positive "
                    f"body literal (unsafe)"
                )
        for lit in rule.negated:
            for term in lit.args:
                if is_var(term) and not is_anon(term) and term not in bound:
                    raise SafetyError(
                        f"variable {term} in the negated literal {lit} of {rule} "
                        f"is not bound by a positive body literal (unsafe)"
                    )


def _tarjan_scc(nodes, adj) -> list:
    """Iterative Tarjan SCC. Returns components as sorted lists, deterministic."""
    index: dict = {}
    low: dict = {}
    on_stack: dict = {}
    stack: list = []
    out: list = []
    counter = [0]

    for root in nodes:
        if root in index:
            continue
        index[root] = low[root] = counter[0]
        counter[0] += 1
        stack.append(root)
        on_stack[root] = True
        work = [(root, iter(sorted(adj.get(root, ()))))]
        while work:
            node, it = work[-1]
            advanced = False
            for nxt in it:
                if nxt not in index:
                    index[nxt] = low[nxt] = counter[0]
                    counter[0] += 1
                    stack.append(nxt)
                    on_stack[nxt] = True
                    work.append((nxt, iter(sorted(adj.get(nxt, ())))))
                    advanced = True
                    break
                if on_stack.get(nxt):
                    low[node] = min(low[node], index[nxt])
            if advanced:
                continue
            work.pop()
            if work:
                parent = work[-1][0]
                low[parent] = min(low[parent], low[node])
            if low[node] == index[node]:
                comp = []
                while True:
                    top = stack.pop()
                    on_stack[top] = False
                    comp.append(top)
                    if top == node:
                        break
                out.append(sorted(comp))
    return out


def stratify(program: Program) -> list:
    """Split predicates into strata (a list of frozensets, dependencies first).

    Raises StratificationError when a negated literal sits inside a recursive
    cycle, i.e. when head and negated body predicate share an SCC.
    """
    preds = sorted(program.predicates())
    adj = {p: set() for p in preds}
    neg_edges = []
    for rule in program.rules:
        for lit in rule.body:
            adj[rule.head_pred].add(lit.pred)
            if lit.negated:
                neg_edges.append((rule.head_pred, lit.pred))

    comps = _tarjan_scc(preds, adj)
    comp_of = {}
    for idx, comp in enumerate(comps):
        for pred in comp:
            comp_of[pred] = idx

    for head, body in neg_edges:
        if comp_of[head] == comp_of[body]:
            members = ", ".join(comps[comp_of[head]])
            raise StratificationError(
                f"negation is not stratified: {head} negates {body} but they are "
                f"mutually recursive (SCC: {members})"
            )

    # condensation: edge comp(p) -> comp(q) means "p depends on q", so q must be
    # computed first. Kahn over the reversed relation.
    must_precede = {i: set() for i in range(len(comps))}
    for pred in preds:
        for dep in adj[pred]:
            a, b = comp_of[pred], comp_of[dep]
            if a != b:
                must_precede[a].add(b)

    indeg = {i: len(must_precede[i]) for i in must_precede}
    dependents = {i: set() for i in must_precede}
    for i, deps in must_precede.items():
        for dep in deps:
            dependents[dep].add(i)

    ready = sorted((i for i, d in indeg.items() if d == 0), key=lambda i: comps[i][0])
    order = []
    while ready:
        cur = ready.pop(0)
        order.append(cur)
        for nxt in sorted(dependents[cur], key=lambda i: comps[i][0]):
            indeg[nxt] -= 1
            if indeg[nxt] == 0:
                ready.append(nxt)
        ready.sort(key=lambda i: comps[i][0])
    if len(order) != len(comps):  # pragma: no cover - condensation is a DAG
        raise StratificationError("internal error: condensation is not a DAG")
    return [frozenset(comps[i]) for i in order]


# --------------------------------------------------------------------------
# Evaluation
# --------------------------------------------------------------------------


def _match(pattern, tup, env):
    """Unify a literal's argument pattern against a tuple. Returns a new env or
    None. Anonymous variables never bind. Bindings made while walking this
    literal are visible to its later arguments, so `p(X, X)` really does mean
    "both arguments equal"."""
    cur = env
    copied = False
    for term, value in zip(pattern, tup):
        if is_var(term):
            if is_anon(term):
                continue
            name = var_name(term)
            if name in cur:
                if cur[name] != value:
                    return None
            else:
                if not copied:
                    cur = dict(env)
                    copied = True
                cur[name] = value
        elif term != value:
            return None
    return cur


def _subst(term, env):
    if is_var(term):
        if is_anon(term):
            raise SafetyError("anonymous variable cannot be substituted")
        return env[var_name(term)]
    return term


def _negation_fails(lit, env, rels):
    """True when the negated literal matches some tuple (so the body fails)."""
    rel = rels.get(lit.pred, ())
    if not rel:
        return False
    for tup in rel:
        if _match(lit.args, tup, env) is not None:
            return True
    return False


def _order_pos(rule, R, rels):
    """Positive body literals, smallest relation first (deterministic)."""
    def size(lit):
        if lit.pred in R:
            return len(R[lit.pred])
        return len(rels.get(lit.pred, ()))

    return sorted(rule.positive, key=lambda lit: (size(lit), lit.pred))


def _join(rule, pos, R, rels, override, stats):
    """Enumerate satisfying environments for the positive body literals."""
    def relation_for(idx, lit):
        if override and idx in override:
            return override[idx]
        if lit.pred in R:
            return R[lit.pred]
        return rels.get(lit.pred, ())

    def rec(idx, env):
        if idx == len(pos):
            for lit in rule.negated:
                if _negation_fails(lit, env, rels):
                    return
            yield env
            return
        lit = pos[idx]
        for tup in sorted(relation_for(idx, lit)):
            stats["bindings"] += 1
            env2 = _match(lit.args, tup, env)
            if env2 is not None:
                yield from rec(idx + 1, env2)

    yield from rec(0, {})


def _derive_one(rule, pos, R, rels, override, out, stats):
    head_pred = rule.head_pred
    for env in _join(rule, pos, R, rels, override, stats):
        stats["rule_runs"] += 1
        tup = tuple(_subst(term, env) for term in rule.head_args)
        if tup in R[head_pred] or tup in out[head_pred]:
            continue
        out[head_pred].add(tup)
        stats["derived"] += 1


def _derive(rules_here, R, rels, wave, override, stats):
    out = {pred: set() for pred in wave}
    for rule in rules_here:
        pos = _order_pos(rule, R, rels)
        _derive_one(rule, pos, R, rels, override, out, stats)
    return out


def _eval_stratum(program, wave, rels, stats):
    rules_here = [r for r in program.rules if r.head_pred in wave]
    if not rules_here:
        return
    R = {pred: rels[pred] for pred in wave}  # live sets, mutated in place

    # Round 0 is naive: it also fires rules whose positive body lives entirely in
    # lower strata (such a rule can never see a same-stratum delta).
    new = _derive(rules_here, R, rels, wave, None, stats)
    for pred, added in new.items():
        R[pred] |= added
    delta = {pred: set(added) for pred, added in new.items()}

    while any(delta.values()):
        stats["rounds"] += 1
        new = {pred: set() for pred in wave}
        for rule in rules_here:
            pos = _order_pos(rule, R, rels)
            for idx, lit in enumerate(pos):
                if lit.pred in wave and delta.get(lit.pred):
                    _derive_one(
                        rule, pos, R, rels, {idx: delta[lit.pred]}, new, stats
                    )
        for pred, added in new.items():
            R[pred] |= added
        delta = new


def evaluate(program: Program, stats: dict = None, check: bool = True) -> dict:
    """Semi-naive bottom-up evaluation. Returns {pred: frozenset(tuples)}."""
    if check:
        check_arity(program)
        check_safety(program)
    strata = stratify(program)
    if stats is None:
        stats = {}
    stats.update(rounds=0, bindings=0, derived=0, rule_runs=0)

    rels = {pred: set(program.facts.get(pred, ())) for pred in program.predicates()}
    for wave in strata:
        _eval_stratum(program, wave, rels, stats)
    return {pred: frozenset(tuples) for pred, tuples in rels.items()}


# --------------------------------------------------------------------------
# Output
# --------------------------------------------------------------------------


def format_relations(rels: dict, preds=None) -> str:
    lines = []
    for pred in sorted(rels) if preds is None else list(preds):
        tuples = sorted(rels.get(pred, ()))
        lines.append(f"{pred}: {len(tuples)} tuple(s)")
        for tup in tuples:
            lines.append(f"    {pred}({', '.join(tup)}).")
    return "\n".join(lines)


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(
        prog="datalog.py",
        description="Evaluate a Datalog program (stratified negation, semi-naive).",
    )
    ap.add_argument("program", help="path to a Datalog source file")
    ap.add_argument("-q", "--query", action="append", default=None,
                    help="only report this predicate (repeatable)")
    ap.add_argument("--text", action="store_true", help="human-readable output")
    ap.add_argument("--stats", action="store_true", help="include run statistics")
    args = ap.parse_args(argv)

    try:
        program = load(args.program)
        stats = {}
        rels = evaluate(program, stats)
        strata = stratify(program)
    except DatalogError as exc:
        print(f"error: {type(exc).__name__}: {exc}", file=sys.stderr)
        return 2
    except OSError as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 2

    if args.text:
        print(format_relations(rels, args.query))
        if args.stats:
            print(json.dumps(stats, sort_keys=True))
        return 0

    payload = {
        "predicates": {
            pred: [list(t) for t in sorted(rels[pred])]
            for pred in (sorted(rels) if args.query is None else args.query)
        },
        "strata": [sorted(wave) for wave in strata],
    }
    if args.stats:
        payload["stats"] = stats
    print(json.dumps(payload, sort_keys=True, indent=2))
    return 0


if __name__ == "__main__":  # pragma: no cover
    sys.exit(main())