#!/usr/bin/env python3
"""test_datalog.py -- verification suite for datalog.py.

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

Exit code is the verdict: 0 = every check passed, 1 = at least one failed.

What is verified, and how
-------------------------
1. Parser KATs          -- facts, rules, comments, '!', zero-arity, int/string
                           constants, and six malformed inputs that must raise.
2. Static checks        -- arity conflicts, rule safety, stratification.
3. Evaluation KATs      -- transitive closure, the non-linear same-generation
                           program, stratified negation, anonymous variables.
4. Order independence   -- shuffled fact/rule order gives identical relations.
5. Differential testing -- 600 randomly generated programs are evaluated by the
                           engine AND by an *independent oracle* written from
                           scratch in this file (reachability-closure SCCs +
                           Kahn waves for stratification; cartesian-product
                           naive fixpoint for evaluation). Both the stratified/
                           unstratified verdict and every relation must agree.
6. Work reduction       -- semi-naive evaluation must enumerate strictly fewer
                           body bindings than the naive oracle on a chain
                           transitive closure (i.e. the delta rule is doing real
                           work, not just returning the right answer by luck).
7. CLI end-to-end       -- subprocess run of datalog.py, JSON contract, exit
                           codes for success and for an unstratified program.

The oracle deliberately shares NOTHING with the engine except the parser and the
term representation; it re-derives stratification and evaluation itself.
"""

from __future__ import annotations

import itertools
import json
import os
import random
import subprocess
import sys
import tempfile
import unittest

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)

import datalog  # noqa: E402
from datalog import (  # noqa: E402
    ArityError,
    Literal,
    ParseError,
    Program,
    Rule,
    SafetyError,
    StratificationError,
    check_arity,
    check_safety,
    evaluate,
    parse_program,
    stratify,
)

DIFFERENTIAL_RUNS = {"programs": 0, "unstratified": 0, "predicates": 0}


# ==========================================================================
# Independent oracle
# ==========================================================================


def oracle_dependency_graph(program):
    preds = sorted(program.predicates())
    adj = {p: set() for p in preds}
    for rule in program.rules:
        for lit in rule.body:
            adj[rule.head_pred].add(lit.pred)
    return preds, adj


def oracle_reachability(program):
    preds, adj = oracle_dependency_graph(program)
    reach = {p: set(adj[p]) | {p} for p in preds}
    changed = True
    while changed:
        changed = False
        for p in preds:
            grown = set(reach[p])
            for q in list(reach[p]):
                grown |= reach[q]
            if grown != reach[p]:
                reach[p] = grown
                changed = True
    return preds, adj, reach


def oracle_stratify(program):
    """Stratify via reachability closure. Returns waves (list[list[str]]) or
    None when the program is not stratified."""
    preds, adj, reach = oracle_reachability(program)
    neg = {(r.head_pred, lit.pred) for r in program.rules for lit in r.negated}

    sccs, seen = [], set()
    for p in preds:
        if p in seen:
            continue
        comp = sorted(q for q in preds if q in reach[p] and p in reach[q])
        seen.update(comp)
        sccs.append(comp)
    comp_of = {p: i for i, comp in enumerate(sccs) for p in comp}

    for head, body in neg:
        if comp_of[head] == comp_of[body]:
            return None

    before = {i: set() for i in range(len(sccs))}
    for p in preds:
        for q in adj[p]:
            if comp_of[p] != comp_of[q]:
                before[comp_of[p]].add(comp_of[q])

    waves, assigned = [], set()
    while len(assigned) < len(sccs):
        ready = sorted(
            (i for i in range(len(sccs)) if i not in assigned and before[i] <= assigned),
            key=lambda i: sccs[i][0],
        )
        if not ready:
            return None
        assigned.update(ready)
        waves.append(sorted(p for i in ready for p in sccs[i]))
    return waves


def oracle_unify(pattern, tup, env):
    cur = env
    copied = False
    for term, value in zip(pattern, tup):
        if datalog.is_var(term):
            if datalog.is_anon(term):
                continue
            name = term[1:]
            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 oracle_neg_fails(lit, env, rels):
    for tup in rels[lit.pred]:
        if oracle_unify(lit.args, tup, env) is not None:
            return True
    return False


def oracle_evaluate(program, stats=None):
    """Naive (product-based) fixpoint per wave. Returns None if unstratified."""
    waves = oracle_stratify(program)
    if waves is None:
        return None
    if stats is None:
        stats = {}
    stats.setdefault("bindings", 0)

    rels = {p: set(program.facts.get(p, ())) for p in program.predicates()}
    for wave in waves:
        wave = set(wave)
        changed = True
        while changed:
            changed = False
            for rule in program.rules:
                if rule.head_pred not in wave:
                    continue
                pools = [sorted(rels[lit.pred]) for lit in rule.positive]
                for combo in itertools.product(*pools):
                    stats["bindings"] += 1
                    env = {}
                    ok = True
                    for lit, tup in zip(rule.positive, combo):
                        env = oracle_unify(lit.args, tup, env)
                        if env is None:
                            ok = False
                            break
                    if not ok:
                        continue
                    for lit in rule.negated:
                        if oracle_neg_fails(lit, env, rels):
                            ok = False
                            break
                    if not ok:
                        continue
                    head = tuple(
                        env[t[1:]] if datalog.is_var(t) else t for t in rule.head_args
                    )
                    if head not in rels[rule.head_pred]:
                        rels[rule.head_pred].add(head)
                        changed = True
    return {p: frozenset(v) for p, v in rels.items()}


# ==========================================================================
# Random program generators
# ==========================================================================

CONSTS = ["a", "b", "c"]
VARS = ["?X", "?Y", "?Z"]
EDB2 = ["e0", "e1"]
IDB2 = ["p0", "p1"]
IDB1 = ["q0", "q1"]


def _random_facts(rng):
    facts = {}
    for pred in EDB2:
        fs = set()
        for a in CONSTS:
            for b in CONSTS:
                if rng.random() < 0.30:
                    fs.add((a, b))
        facts[pred] = fs
    return facts


def gen_program(rng, constrained=True):
    """Random program. constrained=True keeps negation strictly lower-stratum,
    so the program is stratified by construction; constrained=False lets
    negation point anywhere, which yields unstratified programs too."""
    arity = {p: 2 for p in EDB2 + IDB2}
    arity.update({p: 1 for p in IDB1})
    levels = {p: 0 for p in EDB2}
    for p in IDB2 + IDB1:
        levels[p] = rng.choice([1, 1, 2])

    program = Program(facts=_random_facts(rng), rules=[], arities=dict(arity))
    for _ in range(rng.randint(1, 5)):
        head_pred = rng.choice(IDB2 + IDB1)
        hl = levels[head_pred]
        if constrained:
            cands = [p for p in arity if levels[p] <= hl]
        else:
            cands = list(arity)
        positives, used_vars = [], []
        for _ in range(rng.randint(1, 2)):
            bp = rng.choice(cands)
            args = []
            for _ in range(arity[bp]):
                if rng.random() < 0.15:
                    args.append(rng.choice(CONSTS))
                else:
                    args.append(rng.choice(VARS))
            positives.append(Literal(bp, tuple(args), False))
            used_vars.extend(a for a in args if datalog.is_var(a))

        def pick():
            if used_vars and rng.random() < 0.85:
                return rng.choice(used_vars)
            return rng.choice(CONSTS)

        head_args = tuple(pick() for _ in range(arity[head_pred]))
        body = list(positives)
        if rng.random() < 0.35:
            if constrained:
                lower = [p for p in arity if levels[p] < hl]
            else:
                lower = list(arity)
            if lower:
                bp = rng.choice(lower)
                body.append(
                    Literal(bp, tuple(pick() for _ in range(arity[bp])), True)
                )
        program.add_rule(Rule(head_pred, head_args, tuple(body)))
    return program


# ==========================================================================
# 1. parser
# ==========================================================================


class ParserTests(unittest.TestCase):
    def test_facts_and_rule_shape(self):
        p = parse_program("edge(a, b).\npath(X, Y) :- edge(X, Y).\n")
        self.assertEqual(p.facts, {"edge": {("a", "b")}})
        self.assertEqual(len(p.rules), 1)
        rule = p.rules[0]
        self.assertEqual(rule.head_pred, "path")
        self.assertEqual(rule.head_args, ("?X", "?Y"))
        self.assertEqual(str(rule), "path(?X, ?Y) :- edge(?X, ?Y).")

    def test_comments_and_whitespace(self):
        text = "% a comment\n\nedge(a,b).  % trailing\npath(X,Y):-edge(X,Y).\n"
        p = parse_program(text)
        self.assertEqual(p.facts["edge"], {("a", "b")})
        self.assertEqual(len(p.rules), 1)

    def test_bang_is_not(self):
        p = parse_program("dom(a). p(X) :- dom(X), ! q(X). q(b).\n")
        self.assertTrue(p.rules[0].negated)
        self.assertEqual(p.rules[0].negated[0].pred, "q")

    def test_zero_arity_and_constants(self):
        p = parse_program('flag(). n(1). s("x").\n')
        self.assertEqual(p.facts["flag"], {()})
        self.assertEqual(p.facts["n"], {("1",)})
        self.assertEqual(p.facts["s"], {('"x"',)})

    def test_variable_in_fact_rejected(self):
        with self.assertRaises(ParseError):
            parse_program("p(X).\n")

    def test_bad_character_rejected(self):
        with self.assertRaises(ParseError):
            parse_program("p(a) # oops.\n")

    def test_missing_dot_rejected(self):
        with self.assertRaises(ParseError):
            parse_program("p(a)\n")

    def test_negated_head_rejected(self):
        with self.assertRaises(ParseError):
            parse_program("dom(a). not p(X) :- dom(X).\n")

    def test_unbalanced_parens_rejected(self):
        with self.assertRaises(ParseError):
            parse_program("p(a, b.\n")

    def test_empty_body_rejected(self):
        with self.assertRaises(ParseError):
            parse_program("p(a) :- .\n")


# ==========================================================================
# 2. static checks
# ==========================================================================


class StaticCheckTests(unittest.TestCase):
    def test_arity_conflict_between_fact_and_body(self):
        p = parse_program("q(a). p(X) :- q(X, Y).\n")
        with self.assertRaises(ArityError):
            evaluate(p)

    def test_arity_conflict_between_facts(self):
        p = parse_program("p(a, b). p(a).\n")
        with self.assertRaises(ArityError):
            check_arity(p)

    def test_unbound_head_variable(self):
        p = parse_program("q(a). p(X) :- q(Y).\n")
        with self.assertRaises(SafetyError):
            evaluate(p)

    def test_rule_without_positive_body(self):
        p = parse_program("dom(a). p(X) :- not dom(X).\n")
        with self.assertRaises(SafetyError):
            check_safety(p)

    def test_unbound_variable_inside_negation(self):
        p = parse_program("dom(a). q(a). p(X) :- q(Y), not dom(X).\n")
        with self.assertRaises(SafetyError):
            check_safety(p)

    def test_anonymous_variable_in_head(self):
        p = parse_program("q(a). p(_) :- q(a).\n")
        with self.assertRaises(SafetyError):
            check_safety(p)

    def test_negation_inside_recursive_cycle_rejected(self):
        p = parse_program("dom(a). p(X) :- dom(X), not q(X). q(X) :- p(X).\n")
        with self.assertRaises(StratificationError):
            stratify(p)
        self.assertIsNone(oracle_stratify(p))

    def test_mutual_negation_rejected(self):
        p = parse_program(
            "dom(a). p(X) :- dom(X), not q(X). q(X) :- dom(X), not p(X).\n"
        )
        with self.assertRaises(StratificationError):
            stratify(p)
        self.assertIsNone(oracle_stratify(p))

    def test_positive_recursion_is_not_a_stratification_error(self):
        p = parse_program(
            "edge(a,b). edge(b,c).\n"
            "path(X,Y) :- edge(X,Y).\n"
            "path(X,Y) :- edge(X,Z), path(Z,Y).\n"
        )
        waves = stratify(p)
        self.assertEqual(len(waves), 2)
        self.assertEqual(waves[0], frozenset({"edge"}))
        self.assertEqual(waves[1], frozenset({"path"}))
        self.assertIsNotNone(oracle_stratify(p))

    def test_negation_to_lower_stratum_is_fine(self):
        p = parse_program(
            "dom(a). dom(b). q(a).\n"
            "r(X) :- dom(X), not q(X).\n"
        )
        waves = stratify(p)
        names = [sorted(w) for w in waves]
        self.assertLess(names.index(["q"]), names.index(["r"]))
        self.assertEqual(evaluate(p)["r"], frozenset({("b",)}))


# ==========================================================================
# 3. evaluation KATs
# ==========================================================================


class EvaluationTests(unittest.TestCase):
    def test_transitive_closure(self):
        p = parse_program(
            "edge(a,b). edge(b,c). edge(c,d).\n"
            "path(X,Y) :- edge(X,Y).\n"
            "path(X,Y) :- edge(X,Z), path(Z,Y).\n"
        )
        rels = evaluate(p)
        self.assertEqual(
            rels["path"],
            frozenset({("a", "b"), ("a", "c"), ("a", "d"),
                       ("b", "c"), ("b", "d"), ("c", "d")}),
        )

    def test_same_generation_nonlinear_recursion(self):
        p = parse_program(
            "person(a1). person(a2). person(b1). person(b2).\n"
            "parent(b1, a1). parent(b1, a2). parent(b2, a1). parent(b2, a2).\n"
            "sg(X, X) :- person(X).\n"
            "sg(X, Y) :- parent(X, A), parent(Y, B), sg(A, B).\n"
        )
        rels = evaluate(p)
        self.assertEqual(
            rels["sg"],
            frozenset({("a1", "a1"), ("a2", "a2"), ("b1", "b1"), ("b2", "b2"),
                       ("b1", "b2"), ("b2", "b1")}),
        )

    def test_stratified_negation(self):
        p = parse_program(
            "node(a). node(b). node(c).\n"
            "source(a).\n"
            "edge(a,b).\n"
            "reach(X) :- source(X).\n"
            "reach(X) :- edge(Y,X), reach(Y).\n"
            "unreachable(X) :- node(X), not reach(X).\n"
        )
        rels = evaluate(p)
        self.assertEqual(rels["reach"], frozenset({("a",), ("b",)}))
        self.assertEqual(rels["unreachable"], frozenset({("c",)}))

    def test_anonymous_variable_in_negated_literal(self):
        p = parse_program(
            "node(a). node(b). node(c). edge(a,b). edge(c,a).\n"
            "sink(X) :- node(X), not edge(X, _).\n"
        )
        rels = evaluate(p)
        self.assertEqual(rels["sink"], frozenset({("b",)}))

    def test_repeated_variable_inside_one_literal(self):
        # Regression: the join reorders positive literals smallest-first, so
        # `p1(X, X)` can be the FIRST literal joined. A repeated variable must
        # still constrain both argument positions of that single tuple.
        p = parse_program(
            "e0(a, a). e0(a, b). e0(b, a). e1(b, a). e1(b, b). e1(c, a).\n"
            "q1(X) :- e0(Z, X).\n"
            "p1(Y, c) :- e0(Y, Z), not q0(Z).\n"
            "p0(X, c) :- e1(X, Y), p1(X, X), not q0(X).\n"
        )
        rels = evaluate(p)
        self.assertEqual(rels["p0"], frozenset())
        self.assertEqual(rels["p1"], frozenset({("a", "c"), ("b", "c")}))
        self.assertEqual(rels["p0"], oracle_evaluate(p)["p0"])

    def test_negation_over_empty_relation_always_succeeds(self):
        p = parse_program("dom(a). p(X) :- dom(X), not missing(X).\n")
        rels = evaluate(p)
        self.assertEqual(rels["p"], frozenset({("a",)}))
        self.assertEqual(rels["missing"], frozenset())

    def test_zero_arity_derivation(self):
        p = parse_program("flag(). ready() :- flag().\n")
        rels = evaluate(p)
        self.assertEqual(rels["ready"], frozenset({()}))

    def test_constant_arguments_and_repeated_variable(self):
        p = parse_program("e(a,a). e(a,b). loop(X) :- e(X, X).\n")
        rels = evaluate(p)
        self.assertEqual(rels["loop"], frozenset({("a",)}))

    def test_duplicate_facts_collapse(self):
        p = parse_program("edge(a,b). edge(a,b). edge(a,b).\n")
        rels = evaluate(p)
        self.assertEqual(rels["edge"], frozenset({("a", "b")}))

    def test_ground_rule_without_variables(self):
        p = parse_program("dom(a). mark(a) :- dom(a).\n")
        rels = evaluate(p)
        self.assertEqual(rels["mark"], frozenset({("a",)}))

    def test_order_independence(self):
        text = (
            "edge(a,b). edge(b,c). edge(c,d).\n"
            "path(X,Y) :- edge(X,Y).\n"
            "path(X,Y) :- edge(X,Z), path(Z,Y).\n"
        )
        base = evaluate(parse_program(text))
        rng = random.Random(7)
        for _ in range(5):
            clauses = [c.strip() for c in text.strip().split("\n") if c.strip()]
            rng.shuffle(clauses)
            shuffled = evaluate(parse_program("\n".join(clauses) + "\n"))
            self.assertEqual(shuffled, base)

    def test_repeated_evaluation_is_deterministic(self):
        text = (
            "n(a). n(b). n(c). n(d).\n"
            "e(a,b). e(b,c). e(c,d). e(d,a).\n"
            "r(X,Y) :- e(X,Y).\n"
            "r(X,Y) :- r(X,Z), e(Z,Y).\n"
            "lonely(X) :- n(X), not r(X, _).\n"
        )
        first = evaluate(parse_program(text))
        for _ in range(3):
            self.assertEqual(evaluate(parse_program(text)), first)
        self.assertEqual(len(first["r"]), 16)


# ==========================================================================
# 4. differential testing against the oracle
# ==========================================================================


def differential_run(programs, counters=None):
    """Evaluate every program with the engine and with the oracle; return a list
    of (index, program_text, explanation) mismatches. `counters` accumulates the
    coverage numbers."""
    if counters is None:
        counters = {"programs": 0, "unstratified": 0, "predicates": 0}
    failures = []
    for idx, program in enumerate(programs):
        stats = {}
        try:
            got = evaluate(program, stats)
            got_error = None
        except (StratificationError, ArityError, SafetyError) as exc:
            got, got_error = None, type(exc).__name__
        want = oracle_evaluate(program)
        counters["programs"] += 1
        if want is None:
            counters["unstratified"] += 1
            if got_error != "StratificationError":
                failures.append(
                    (idx, program.text(),
                     f"oracle says unstratified, engine said {got_error!r}")
                )
            continue
        if got_error is not None:
            failures.append(
                (idx, program.text(), f"oracle stratified, engine raised {got_error}")
            )
            continue
        counters["predicates"] += len(got)
        if got != want:
            diff = {
                p: (sorted(got.get(p, ())), sorted(want.get(p, ())))
                for p in sorted(set(got) | set(want))
                if got.get(p) != want.get(p)
            }
            failures.append((idx, program.text(), f"relation mismatch {diff}"))
    return failures


class DifferentialTests(unittest.TestCase):
    def _run(self, programs):
        failures = differential_run(programs, DIFFERENTIAL_RUNS)
        return failures

    def test_differential_constrained_300(self):
        rng = random.Random(20261024)
        programs = [gen_program(rng, constrained=True) for _ in range(300)]
        failures = self._run(programs)
        self.assertEqual(failures[:3], [], f"{len(failures)}/300 mismatches")

    def test_differential_free_negation_300(self):
        rng = random.Random(20261025)
        programs = [gen_program(rng, constrained=False) for _ in range(300)]
        failures = self._run(programs)
        self.assertEqual(failures[:3], [], f"{len(failures)}/300 mismatches")
        self.assertGreater(
            DIFFERENTIAL_RUNS["unstratified"], 0,
            "the free generator produced no unstratified program, so the "
            "stratification check was not exercised",
        )

    def test_generator_actually_builds_nontrivial_programs(self):
        rng = random.Random(1)
        programs = [gen_program(rng, constrained=True) for _ in range(50)]
        self.assertTrue(all(p.rules for p in programs))
        self.assertGreater(sum(len(p.rules) for p in programs), 100)
        self.assertTrue(any(any(lit.negated for lit in r.body)
                            for p in programs for r in p.rules))


# ==========================================================================
# 5. semi-naive does less work than naive
# ==========================================================================


class WorkReductionTests(unittest.TestCase):
    @staticmethod
    def chain_program(n):
        facts = "".join(f"edge(n{i}, n{i + 1}).\n" for i in range(n - 1))
        return parse_program(
            facts
            + "path(X,Y) :- edge(X,Y).\n"
            + "path(X,Y) :- edge(X,Z), path(Z,Y).\n"
        )

    def test_semi_naive_enumerates_fewer_bindings(self):
        program = self.chain_program(12)
        engine_stats, oracle_stats = {}, {}
        rels = evaluate(program, engine_stats)
        oracle_rels = oracle_evaluate(program, oracle_stats)
        self.assertEqual(rels, oracle_rels)
        self.assertEqual(len(rels["path"]), 66)  # C(12,2)
        self.assertLess(
            engine_stats["bindings"], oracle_stats["bindings"],
            f"semi-naive bindings {engine_stats['bindings']} vs naive "
            f"{oracle_stats['bindings']}",
        )
        # The delta rule must buy a real margin, not a rounding error.
        self.assertLess(engine_stats["bindings"], oracle_stats["bindings"] // 2)

    def test_stats_are_populated(self):
        stats = {}
        evaluate(self.chain_program(6), stats)
        for key in ("rounds", "bindings", "derived", "rule_runs"):
            self.assertIn(key, stats)
        self.assertGreater(stats["derived"], 0)
        self.assertGreater(stats["rounds"], 1)


# ==========================================================================
# 6. CLI
# ==========================================================================


class CliTests(unittest.TestCase):
    SCRIPT = os.path.join(HERE, "datalog.py")

    def _run(self, text, *args):
        with tempfile.TemporaryDirectory(prefix=".datalog-fixture-", dir=HERE) as tmp:
            path = os.path.join(tmp, "program.dl")
            with open(path, "w", encoding="utf-8") as fh:
                fh.write(text)
            proc = subprocess.run(
                [sys.executable, self.SCRIPT, path, *args],
                capture_output=True, text=True, timeout=120,
            )
        return proc

    def test_json_contract(self):
        proc = self._run(
            "node(a). node(b). node(c). source(a). edge(a,b).\n"
            "reach(X) :- source(X).\n"
            "reach(X) :- edge(Y,X), reach(Y).\n"
            "unreach(X) :- node(X), not reach(X).\n",
            "--stats",
        )
        self.assertEqual(proc.returncode, 0, proc.stderr)
        payload = json.loads(proc.stdout)
        self.assertEqual(payload["predicates"]["reach"], [["a"], ["b"]])
        self.assertEqual(payload["predicates"]["unreach"], [["c"]])
        self.assertEqual(payload["predicates"]["node"], [["a"], ["b"], ["c"]])
        self.assertIn("stats", payload)
        self.assertTrue(payload["strata"])

    def test_unstratified_program_exits_2(self):
        proc = self._run("dom(a). p(X) :- dom(X), not q(X). q(X) :- p(X).\n")
        self.assertEqual(proc.returncode, 2)
        self.assertIn("StratificationError", proc.stderr)

    def test_text_output(self):
        proc = self._run("e(a,b).\n", "--text")
        self.assertEqual(proc.returncode, 0, proc.stderr)
        self.assertIn("e(a, b).", proc.stdout)


def tearDownModule():
    print(
        "\n[differential] programs evaluated: {programs} "
        "(unstratified: {unstratified}, predicate-relations compared: {predicates})".format(
            **DIFFERENTIAL_RUNS
        ),
        file=sys.stderr,
    )


if __name__ == "__main__":
    unittest.main(verbosity=1)