#!/usr/bin/env python3
"""Live parity: behavior delegating facade (Rust core) vs the Python-pin
oracle path, over the SAME case matrix the truth generator froze.
Run: python3 scripts/check_behavior_parity.py  -> x/y lines, exit 1 on any
divergence.

FROZEN layer (gap-close): the new pure-fn surface (connect schema, core
probe, python-pin predicates, reinforce plans, list_items,
behavior_stats) replayed through the cdylib's ctypes FFI against the
frozen oracle truth — live parity for the ported fns.
"""
from __future__ import annotations

import ctypes
import json
import os
import sys
import tempfile
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

_tmp = tempfile.mkdtemp(prefix="bh-parity-")
os.environ["RUNE_HOME"] = _tmp
os.environ["RUNE_MEMORY_DIR"] = os.path.join(_tmp, "memory")

from rune import behavior as bh  # noqa: E402
from rune import bh_loader  # noqa: E402

if not bh_loader.available():
    print("0/0 — rust behavior core not built")
    sys.exit(0)

PASS = 0
FAIL = 0
DIVERGENCES = []

_CLOCK_KEYS = {"created_at", "last_reinforced", "last_activated"}


def strip_clocks(obj):
    """Timestamps are pinned by the shim (`now`) — wall-clock values
    differ per run by design; compare everything else."""
    if isinstance(obj, dict):
        return {k: strip_clocks(v) for k, v in obj.items() if k not in _CLOCK_KEYS}
    if isinstance(obj, list):
        return [strip_clocks(x) for x in obj]
    return obj


def check(name, rust, py):
    global PASS, FAIL
    if strip_clocks(rust) == strip_clocks(py):
        PASS += 1
    else:
        FAIL += 1
        DIVERGENCES.append((name, rust, py))


def row(table, rid):
    db = bh.connect()
    try:
        r = db.execute(f"SELECT * FROM {table} WHERE id=?", (rid,)).fetchone()
        return dict(r) if r else None
    finally:
        db.close()


def both(name, fn):
    """Drive fn() twice: pinned-python first (fresh db), rust second
    (fresh db). Databases are reset between runs so ids align."""
    os.environ["RUNE_BEHAVIOR_CORE"] = "python"
    bh.DB_PATH = os.path.join(os.environ["RUNE_MEMORY_DIR"], "behavior.db")
    if os.path.exists(bh.DB_PATH):
        os.remove(bh.DB_PATH)
    py = fn()
    os.environ["RUNE_BEHAVIOR_CORE"] = "auto"
    if os.path.exists(bh.DB_PATH):
        os.remove(bh.DB_PATH)
    rust = fn()
    os.environ["RUNE_BEHAVIOR_CORE"] = "python"
    check(name, rust, py)


# ── FROZEN layer: the gap-close fns through the cdylib FFI ────────────
FROZEN_PASS = 0
FROZEN_FAIL = 0


def frozen_check(name, got, want):
    global FROZEN_PASS, FROZEN_FAIL
    if got == want:
        FROZEN_PASS += 1
    else:
        FROZEN_FAIL += 1
        DIVERGENCES.append((f"FROZEN {name}", got, want))


_new_lib = ctypes.CDLL(str(bh_loader._LIB_PATH))
for _sym in ("bh_connect", "bh_core_probe", "bh_python_predicates",
             "bh_reinforce_plan", "bh_list_items", "bh_behavior_stats"):
    _fn = getattr(_new_lib, _sym)
    _fn.argtypes = [ctypes.c_void_p]
    _fn.restype = ctypes.c_void_p


def ffi(sym, payload):
    data = json.dumps(payload).encode()
    buf = ctypes.create_string_buffer(data)
    ptr = getattr(_new_lib, sym)(ctypes.cast(buf, ctypes.c_void_p))
    if not ptr:
        raise RuntimeError(f"{sym} returned null")
    out = ctypes.string_at(ptr).decode()
    bh_loader._free(ptr)
    doc = json.loads(out)
    if not doc.get("ok"):
        raise RuntimeError(doc.get("err", f"{sym} failed"))
    return doc["value"]


os.environ["RUNE_BEHAVIOR_CORE"] = "python"
_truth = json.loads((Path(__file__).resolve().parents[1]
                     / "src" / "rust" / "behavior" / "src" / "bh_truth.json").read_text())

# connect: schema + pragma through the FFI (the frozen truth carries no
# path — no absolute paths in frozen truth; the path law is asserted
# separately: memory-dir relative, behavior.db suffix)
_connect_got = ffi("bh_connect", {})
frozen_check("connect-schema",
             {"journal_mode": _connect_got["journal_mode"], "schema": _connect_got["schema"]},
             {"journal_mode": _truth["cases"]["connect"]["journal_mode"],
              "schema": _truth["cases"]["connect"]["schema"]})
if not _connect_got["path"].endswith("/behavior.db"):
    FROZEN_FAIL += 1
    DIVERGENCES.append(("FROZEN connect-path", _connect_got["path"], "*/behavior.db"))
else:
    FROZEN_PASS += 1
# core probe under the pin (frozen core: null)
frozen_check("core-probe-pin", ffi("bh_core_probe", {}), {"available": False})

# python-pin predicate matrix
for c in _truth["cases"]["predicates"]["corrections"]:
    got = ffi("bh_python_predicates", {"kind": "correction", "row": c["row"],
                                       "context": c["context"], "context_low": c["context_low"]})
    frozen_check(f"pred-corr-{c['row']['id']}-{c['context'][:12]}", got, c["matched"])
for h in _truth["cases"]["predicates"]["habits"]:
    got = ffi("bh_python_predicates", {"kind": "habit", "row": h["row"], "context": h["context"]})
    frozen_check(f"pred-habit-{h['row']['id']}-{h['context'][:12]}", got, h["matched"])

# reinforce plans: full pre->final threading through BOTH paths
for key in ("reinforce_habit", "reinforce_correction"):
    for case in _truth["cases"][key]:
        name = case["name"]
        if case["pre_row"] is None:
            item_id = 99991 if name == "habit-not-found" else 99992
            payload = {"item_id": item_id, "outcome": case["outcome"],
                       "item_type": case["item_type"], "row": None, "ts": "ts"}
            frozen_check(f"reinf-{name}", ffi("bh_reinforce_plan", payload)["result"],
                         case["results"][0])
            continue
        ts = case["pre_row"]["created_at"]
        cur = dict(case["pre_row"])
        for i in range(case["n"]):
            payload = {"item_id": case["pre_row"]["id"], "outcome": case["outcome"],
                       "item_type": case["item_type"], "row": cur, "ts": ts, "path": "rust"}
            plan = ffi("bh_reinforce_plan", payload)
            twin = ffi("bh_reinforce_plan", {**payload, "path": "python"})
            frozen_check(f"reinf-{name}-step{i}-halves", plan, twin)
            frozen_check(f"reinf-{name}-step{i}", plan["result"], case["results"][i])
            cur.update(plan["writeback"])
        frozen_check(f"reinf-{name}-final", cur, case["final_row"])

# list_items + behavior_stats over the frozen row collections
_sf = _truth["cases"]["stats_full"]
for case in _truth["cases"]["list_full"]:
    got = ffi("bh_list_items", {"corrections": _sf["corrections"], "habits": _sf["habits"],
                                "item_type": case["item_type"], "limit": case["limit"]})
    frozen_check(f"list-{case['name']}", got, case["result"])
frozen_check("stats", ffi("bh_behavior_stats", {"corrections": _sf["corrections"],
                                                "habits": _sf["habits"],
                                                "activations": _sf["activations"]}),
             _sf["stats"])
os.environ["RUNE_BEHAVIOR_CORE"] = "auto"

# ── LIVE layer: the delegating facade vs the python pin ───────────────


# ── add shapes ────────────────────────────────────────────────────────
def add_shapes():
    out = [bh.add_correction("say X", "say Y"),
           bh.add_correction("orig", "corr", scope_tags={"k": "v"}, domain="shell", pattern="shell"),
           bh.add_habit("trig", "act"),
           bh.add_habit("trig2", "act2", name="n", context_tags=["a"], min_activations=5)]
    return json.loads(json.dumps(out))


both("add-shapes", add_shapes)

# ── seed the shared match fixture in BOTH modes ───────────────────────
def seed_match():
    bh.add_correction("o1", "scoped corr", scope_tags={}, domain="shell", pattern="deploy")
    bh.add_correction("o2", "domain corr", scope_tags={}, domain="editor")
    bh.add_correction("o3", "plain corr")
    bh.add_correction("o4", "bad regex corr", pattern="([unclosed")
    low = bh.add_correction("o5", "low conf corr")["added"]
    bh.add_habit("deploy", "run deploy script")
    bh.add_habit("test (unit|e2e)", "run tests")
    bh.add_habit("([broken", "bad regex habit")
    hslow = bh.add_habit("zzz-no-match-ever", "quiet")["added"]
    db = bh.connect()
    db.execute("UPDATE corrections SET confidence=0.2 WHERE id=?", (low,))
    db.execute("UPDATE habits SET confidence=0.2 WHERE id=?", (hslow,))
    db.commit()
    db.close()


def match_shapes():
    seed_match()
    cases = []
    for ctx, mt in [("deploy to shell now", None), ("deploy the editor", None),
                    ("unrelated chatter here", None), ("deploy to shell now", "habit"),
                    ("deploy to shell now", "correction"), ("run the unit tests", None),
                    ("zzz-no-match-ever quiet", None), ("([unclosed anything", None)]:
        got = bh.match_context(ctx, match_type=mt)
        cases.append({
            "corrections": sorted(c["corrected_behavior"] for c in got["corrections"]),
            "habits": sorted(h["action_template"] for h in got["habits"]),
            "full": json.loads(json.dumps(got)),
        })
    return cases


both("match-matrix", match_shapes)

# ── reinforce matrix ──────────────────────────────────────────────────
def reinforce_shapes():
    seed_match()
    out = []
    for name, item_type, outcome, seed_fn, pre, n in [
        ("success-promotes-at-3", "habit", "success", lambda: bh.add_habit("t1", "a1"), {}, 3),
        ("failure-counts", "habit", "failure", lambda: bh.add_habit("t2", "a2"), {}, 2),
        ("neutral-bumps-activation-only", "habit", "neutral", lambda: bh.add_habit("t3", "a3"), {}, 1),
        ("already-active-stays", "habit", "success", lambda: bh.add_habit("t4", "a4", min_activations=1), {}, 2),
        ("min5-gate", "habit", "success", lambda: bh.add_habit("t5", "a5", min_activations=5), {}, 4),
        ("success-plus-point-one", "correction", "success", lambda: bh.add_correction("x1", "y1"), {}, 1),
        ("failure-archives", "correction", "failure", lambda: bh.add_correction("x2", "y2"), {}, 1),
        ("double-failure-floor", "correction", "failure", lambda: bh.add_correction("x3", "y3"), {}, 3),
        ("ceiling-at-one", "correction", "success", lambda: bh.add_correction("x4", "y4"), {"confidence": 0.95}, 1),
        ("neutral-noop", "correction", "neutral", lambda: bh.add_correction("x5", "y5"), {}, 1),
    ]:
        rid = seed_fn()["added"]
        if pre:
            db = bh.connect()
            table = "habits" if item_type == "habit" else "corrections"
            for k, v in pre.items():
                db.execute(f"UPDATE {table} SET {k}=? WHERE id=?", (v, rid))
            db.commit()
            db.close()
        results = [bh.reinforce(rid, outcome, item_type=item_type) for _ in range(n)]
        out.append({"name": name, "results": json.loads(json.dumps(results)),
                    "final": json.loads(json.dumps(row("habits" if item_type == "habit" else "corrections", rid)))})
    out.append({"name": "not-found-habit",
                "results": [bh.reinforce(99991, "success", item_type="habit")], "final": None})
    out.append({"name": "not-found-correction",
                "results": [bh.reinforce(99992, "success", item_type="correction")], "final": None})
    return out


both("reinforce-matrix", reinforce_shapes)

# ── list + stats ──────────────────────────────────────────────────────
def list_stats_shapes():
    seed_match()
    out = {}
    check("list-keys", sorted(bh.list_items().keys()), ["corrections", "habits"])
    check("limit-1", len(bh.list_items(limit=1).get("corrections", [])), 1)
    check("type-correction-keys", sorted(bh.list_items(item_type="correction").keys()), ["corrections"])
    check("type-habit-keys", sorted(bh.list_items(item_type="habit").keys()), ["habits"])
    check("type-all-keys", sorted(bh.list_items(item_type="all").keys()), ["corrections", "habits"])
    out["stats"] = bh.behavior_stats()
    return json.loads(json.dumps(out))


both("list-stats", list_stats_shapes)

print(f"FROZEN {FROZEN_PASS}/{FROZEN_PASS + FROZEN_FAIL} pure-fn replay through cdylib FFI")
print(f"{PASS}/{PASS + FAIL} live parity (rust facade vs python pin)")
for name, rust, py in DIVERGENCES:
    print(f"DIVERGENCE {name}:")
    print(f"  rust: {json.dumps(rust, sort_keys=True)[:400]}")
    print(f"  py:   {json.dumps(py, sort_keys=True)[:400]}")
sys.exit(1 if (FAIL or FROZEN_FAIL) else 0)