#!/usr/bin/env python3
"""Differential harness: run each case against tai.py and detz, compare.

Usage:
  tools/parity_diff.py [--detz PATH] [--py PATH] [--filter TEXT] [--no-llm]
                       [--cases DIR] [--jobs N] [--list]

Every case runs twice, once per side, in its own throwaway TAI_HOME and cwd
(never ~/.tai), with the real binaries and a real vault. Compared channels:
stdout, stderr, exit code, vault state (sqlite dump of every *.db in the home),
and declared files. Volatile data is normalized ONLY by rules declared in the
case itself (each with a `why`); there is no blanket scrubbing. A rule that
matched nothing on either side is reported as unused.

LLM cases (`"llm": true`) make a real model call on both sides (tests/README.md
rule, no mocks). Their outputs are not equal-compared; both sides must each
satisfy the case's `checks`. They are probabilistic: `attempts` / `min_pass`
report a ratio, and the result carries it.

Writes build/parity/diff-results.json and diff-report.md. Exit 1 if any case
that is not marked expected_fail failed, or an expected_fail case passed.
See docs/parity.md for the case format.
"""
import argparse
import concurrent.futures
import difflib
import glob as globmod
import hashlib
import json
import os
import re
import shutil
import signal
import sqlite3
import subprocess
import sys
import tempfile
import time

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
REPO = os.path.dirname(ROOT)
OUT = os.path.join(ROOT, "build", "parity")
CASES_DIR = os.path.join(ROOT, "tests", "diff", "cases")

# Hard caps (nothing unbounded).
MAX_JOBS = 6
MAX_STEPS = 8
MAX_STEP_TIMEOUT = 300
MAX_DIFF_LINES = 40
MAX_CAPTURE = 2_000_000
MAX_ATTEMPTS = 5
BASE_ENV = {"COLUMNS": "80", "LINES": "24", "TAI_PASSPHRASE": "", "LC_ALL": "C.UTF-8", "TERM": "dumb"}
VAULT_SKIP_PREFIXES = ("sqlite_", "fts_docs_")  # sqlite internals and fts shadow tables
CHANNELS = ("stdout", "stderr", "exit", "vault", "files")


class CaseError(Exception):
    pass


def sha256_file(path):
    digest = hashlib.sha256()
    with open(path, "rb") as fh:
        for block in iter(lambda: fh.read(1 << 20), b""):
            digest.update(block)
    return digest.hexdigest()


class Rule:
    """One declared normalization: regex -> replacement on named channels."""

    def __init__(self, spec, case_id):
        for key in ("on", "pattern", "replace", "why"):
            if key not in spec:
                raise CaseError("%s: normalize rule missing %r" % (case_id, key))
        if not str(spec["why"]).strip():
            raise CaseError("%s: normalize rule needs a non-empty why" % case_id)
        self.on = list(spec["on"])
        for ch in self.on:
            if ch not in CHANNELS:
                raise CaseError("%s: normalize rule on unknown channel %r" % (case_id, ch))
        self.regex = re.compile(spec["pattern"], re.M | (re.S if spec.get("dotall") else 0))
        self.replace = spec["replace"]
        self.why = spec["why"]
        self.hits = {"ref": 0, "cpp": 0}

    def apply(self, channel, text, side):
        if channel not in self.on:
            return text
        text, count = self.regex.subn(self.replace, text)
        self.hits[side] += count
        return text


class Step:
    def __init__(self, spec, case_id):
        self.argv = [str(a) for a in spec.get("argv", [])]
        self.stdin = spec.get("stdin", "")
        self.env = {k: str(v) for k, v in spec.get("env", {}).items()}
        self.files = dict(spec.get("files", {}))
        self.home_files = dict(spec.get("home_files", {}))
        self.mode = spec.get("mode", "pipe")
        self.timeout = min(float(spec.get("timeout", 60)), MAX_STEP_TIMEOUT)
        if self.mode not in ("pipe", "pty"):
            raise CaseError("%s: unknown step mode %r" % (case_id, self.mode))


class Case:
    def __init__(self, spec, source):
        self.id = spec["id"]
        self.source = source
        self.doc = spec.get("doc", "")
        self.symbols = list(spec.get("symbols", []))
        self.tests = list(spec.get("tests", []))
        self.llm = bool(spec.get("llm", False))
        self.expected_fail = spec.get("expected_fail")
        self.steps = [Step(s, self.id) for s in spec["steps"]]
        if not self.steps or len(self.steps) > MAX_STEPS:
            raise CaseError("%s: 1..%d steps required" % (self.id, MAX_STEPS))
        compare = spec.get("compare", {})
        self.compare = {
            "stdout": compare.get("stdout", True),
            "stderr": compare.get("stderr", True),
            "exit": compare.get("exit", True),
            "vault": compare.get("vault", False),
            "files": compare.get("files", []),
        }
        self.vault_ignore = list(compare.get("vault_ignore_tables", []))
        if self.vault_ignore and not str(compare.get("vault_ignore_why", "")).strip():
            raise CaseError("%s: vault_ignore_tables needs vault_ignore_why" % self.id)
        self.rule_specs = list(spec.get("normalize", []))
        for rs in self.rule_specs:
            Rule(rs, self.id)  # validate at load
        self.checks = list(spec.get("checks", []))
        # Symbols named by the reference's own output (e.g. every tool /tools lists), claimed only
        # when the case passes, i.e. when detz printed the same thing.
        self.symbols_from_stdout = list(spec.get("symbols_from_stdout", []))
        self.attempts = min(int(spec.get("attempts", 1)), MAX_ATTEMPTS)
        self.min_pass = int(spec.get("min_pass", 1))
        if self.llm and not self.checks:
            raise CaseError("%s: llm case needs checks (outputs are not equal-compared)" % self.id)

    def rules(self):
        return [Rule(rs, self.id) for rs in self.rule_specs]


def load_cases(directory):
    cases, seen = [], {}
    for path in sorted(globmod.glob(os.path.join(directory, "*.json"))):
        with open(path, encoding="utf-8") as fh:
            doc = json.load(fh)
        for spec in doc["cases"]:
            case = Case(spec, os.path.basename(path))
            if case.id in seen:
                raise CaseError("duplicate case id %s (%s, %s)" % (case.id, seen[case.id], path))
            seen[case.id] = path
            cases.append(case)
    return cases


class Side:
    """One implementation under test."""

    def __init__(self, name, command):
        self.name = name
        self.command = list(command)


class Capture:
    """Everything observed from running one case once on one side."""

    def __init__(self):
        self.steps = []   # list of dict(argv, rc, stdout, stderr, timed_out)
        self.vault = ""
        self.files = ""
        self.sql = None   # open read-only vault path for checks
        self.home = None
        self.cwd = None


class Runner:
    def __init__(self, work_root):
        self.work_root = work_root

    def run_step(self, side, step, home, cwd):
        for rel, content in step.files.items():
            self._write(os.path.join(cwd, rel), content)
        for rel, content in step.home_files.items():
            self._write(os.path.join(home, rel), content)
        # Hermetic environment: inherited TAI_*, CAP_*, TMUX* leak the caller's session into the
        # run (parent shell session, pane resolution), so they are dropped. Declared per step env wins.
        env = {k: v for k, v in os.environ.items() if not k.startswith(("TAI_", "CAP_", "TMUX"))}
        env.update(BASE_ENV)
        user_home = os.path.join(os.path.dirname(home), "user")
        os.makedirs(user_home, exist_ok=True)
        env.update({"TAI_HOME": home, "HOME": user_home})
        env.update(step.env)
        cmd = side.command + step.argv
        start = time.monotonic()
        if step.mode == "pty":
            return self._run_pty(cmd, step, env, cwd, start)
        proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                cwd=cwd, env=env, start_new_session=True)
        timed_out = False
        try:
            out, err = proc.communicate(step.stdin.encode("utf-8"), timeout=step.timeout)
        except subprocess.TimeoutExpired:
            timed_out = True
            self._kill(proc)
            out, err = proc.communicate()
        return {"argv": step.argv, "rc": proc.returncode, "timed_out": timed_out,
                "stdout": out[:MAX_CAPTURE].decode("utf-8", "replace"),
                "stderr": err[:MAX_CAPTURE].decode("utf-8", "replace"),
                "secs": round(time.monotonic() - start, 2)}

    def _run_pty(self, cmd, step, env, cwd, start):
        import asyncio
        sys.path.insert(0, REPO)
        import tai_test  # the one base harness (tests/README.md); pty merges stderr into stdout

        async def go():
            proc = tai_test.ManagedProcess(cmd, mode="pty", cwd=cwd, env=env)
            await proc.start()
            timed_out = False
            text = ""
            try:
                for line in step.stdin.splitlines():
                    await proc.send_line(line)
                text = await proc.drain_until_eof(timeout=step.timeout)
            except tai_test.TestTimeout:
                timed_out = True
                text = "".join(l.text for l in proc.transcript if l.stream == "stdout")
            rc = await proc.terminate() if timed_out else await proc.wait(5)
            await proc.terminate()
            return text, rc, timed_out
        text, rc, timed_out = asyncio.run(go())
        return {"argv": step.argv, "rc": rc, "timed_out": timed_out, "stdout": tai_test.strip_ansi(text)[:MAX_CAPTURE],
                "stderr": "", "secs": round(time.monotonic() - start, 2)}

    @staticmethod
    def _write(path, content):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        with open(path, "w", encoding="utf-8") as fh:
            fh.write(content)

    @staticmethod
    def _kill(proc):
        for sig in (signal.SIGTERM, signal.SIGKILL):
            try:
                os.killpg(proc.pid, sig)
            except (ProcessLookupError, PermissionError):
                return
            try:
                proc.wait(3)
                return
            except subprocess.TimeoutExpired:
                continue

    def capture(self, case, side, tag):
        base = os.path.join(self.work_root, case.id.replace("/", "_"), "%s-%s" % (side.name, tag))
        shutil.rmtree(base, ignore_errors=True)
        home, cwd = os.path.join(base, "home"), os.path.join(base, "cwd")
        os.makedirs(home)
        os.makedirs(cwd)
        cap = Capture()
        cap.home, cap.cwd = home, cwd
        for step in case.steps:
            cap.steps.append(self.run_step(side, step, home, cwd))
        if case.compare["vault"]:
            cap.vault = VaultDump(home, case.vault_ignore).text()
        if case.compare["files"]:
            cap.files = FilesDump(home, cwd, case.compare["files"]).text()
        return cap


class VaultDump:
    """Deterministic text dump of every sqlite db under the home (read-only copy)."""

    def __init__(self, home, ignore):
        self.home = home
        self.ignore = set(ignore)

    def text(self):
        lines = []
        for base, _dirs, names in sorted(os.walk(self.home)):
            for name in sorted(names):
                if name.endswith(".db"):
                    path = os.path.join(base, name)
                    lines.append("# db %s" % os.path.relpath(path, self.home))
                    lines.extend(self._dump(path))
        return "\n".join(lines) + ("\n" if lines else "")

    def _dump(self, path):
        out = []
        tmp = tempfile.mkdtemp()
        try:
            copy = os.path.join(tmp, "v.db")
            shutil.copy(path, copy)
            for suffix in ("-wal", "-shm"):
                if os.path.exists(path + suffix):
                    shutil.copy(path + suffix, copy + suffix)
            db = sqlite3.connect(copy)
            db.text_factory = lambda b: b.decode("utf-8", "replace")
            tables = [r for r in db.execute("SELECT name, sql FROM sqlite_master WHERE type='table' ORDER BY name")]
            for name, sql in tables:
                if name.startswith(VAULT_SKIP_PREFIXES) or name in self.ignore:
                    continue
                out.append("## table %s :: %s" % (name, " ".join((sql or "").split())))
                rows = [json.dumps(list(r), ensure_ascii=False, default=repr) for r in db.execute('SELECT * FROM "%s"' % name)]
                out.extend("  " + r for r in sorted(rows))
            idx = [r[0] for r in db.execute("SELECT name FROM sqlite_master WHERE type='index' AND name NOT LIKE 'sqlite_%' ORDER BY name")]
            out.append("## indexes %s" % ",".join(idx))
            db.close()
        except sqlite3.Error as exc:
            out.append("## unreadable: %s" % exc)
        finally:
            shutil.rmtree(tmp, ignore_errors=True)
        return out


class FilesDump:
    """Path + content hash listing for declared globs, rooted at home or cwd."""

    def __init__(self, home, cwd, specs):
        self.roots = {"home": home, "cwd": cwd}
        self.specs = specs

    def text(self):
        lines = []
        for spec in self.specs:
            root = self.roots[spec["root"]]
            for path in sorted(globmod.glob(os.path.join(root, spec["glob"]), recursive=True)):
                if os.path.isfile(path):
                    rel = "%s:%s" % (spec["root"], os.path.relpath(path, root))
                    with open(path, "rb") as fh:
                        data = fh.read(MAX_CAPTURE)
                    try:
                        body = data.decode("utf-8")
                        lines.append("%s\t%s" % (rel, json.dumps(body, ensure_ascii=False)))
                    except UnicodeDecodeError:
                        lines.append("%s\tsha256:%s" % (rel, hashlib.sha256(data).hexdigest()))
        return "\n".join(sorted(lines)) + ("\n" if lines else "")


class Normalizer:
    def __init__(self, case):
        self.rules = case.rules()

    def apply(self, channel, text, side):
        for rule in self.rules:
            text = rule.apply(channel, text, side)
        return text

    def unused(self):
        return [r.why for r in self.rules if r.hits["ref"] == 0 and r.hits["cpp"] == 0]


class Checker:
    """Evaluates `checks` of an llm case against one side's capture."""

    def __init__(self, case):
        self.case = case

    def evaluate(self, cap, side):
        fails = []
        for chk in self.case.checks:
            if chk.get("side", "both") not in ("both", side):
                continue
            kind = chk["kind"]
            if kind == "exit":
                got = cap.steps[chk["step"]]["rc"]
                if got != chk["value"]:
                    fails.append("exit of step %d is %r, want %r" % (chk["step"], got, chk["value"]))
            elif kind in ("stdout_regex", "stderr_regex"):
                text = cap.steps[chk["step"]]["stdout" if kind == "stdout_regex" else "stderr"]
                if not re.search(chk["pattern"], text, re.S | re.I):
                    fails.append("%s step %d: /%s/ not found" % (kind, chk["step"], chk["pattern"]))
            elif kind == "vault_sql":
                got = self._sql(cap, chk["sql"])
                if not self._cmp(got, chk["op"], chk["value"]):
                    fails.append("vault_sql %r gave %r, want %s %r" % (chk["sql"], got, chk["op"], chk["value"]))
            else:
                fails.append("unknown check kind %r" % kind)
        return fails

    @staticmethod
    def _cmp(got, op, want):
        try:
            return {"==": got == want, ">=": got >= want, "<=": got <= want, ">": got > want, "<": got < want}[op]
        except TypeError:
            return False

    @staticmethod
    def _sql(cap, sql):
        path = os.path.join(cap.home, "memory.db")
        if not os.path.exists(path):
            return None
        tmp = tempfile.mkdtemp()
        try:
            copy = os.path.join(tmp, "v.db")
            shutil.copy(path, copy)
            for suffix in ("-wal", "-shm"):
                if os.path.exists(path + suffix):
                    shutil.copy(path + suffix, copy + suffix)
            db = sqlite3.connect(copy)
            row = db.execute(sql).fetchone()
            db.close()
            return row[0] if row else None
        except sqlite3.Error as exc:
            return "sqlerror: %s" % exc
        finally:
            shutil.rmtree(tmp, ignore_errors=True)


class Verdict:
    def __init__(self, case):
        self.case = case
        self.status = "fail"      # pass | fail | xfail | xpass | skipped
        self.problems = []        # human lines, capped
        self.ratio = None         # llm: "ok/attempts" per side
        self.unused_rules = []
        self.extra_symbols = []
        self.secs = 0.0

    def as_dict(self):
        return {"id": self.case.id, "status": self.status, "llm": self.case.llm, "symbols": self.case.symbols + self.extra_symbols,
                "tests": self.case.tests, "expected_fail": self.case.expected_fail, "problems": self.problems,
                "ratio": self.ratio, "unused_rules": self.unused_rules, "secs": round(self.secs, 1),
                "source": self.case.source}


class CaseExecutor:
    def __init__(self, runner, ref, cpp):
        self.runner, self.ref, self.cpp = runner, ref, cpp

    def execute(self, case):
        start = time.monotonic()
        verdict = Verdict(case)
        try:
            if case.llm:
                self._llm(case, verdict)
            else:
                self._equal(case, verdict)
        except Exception as exc:  # harness fault is a failure, loudly
            verdict.problems.append("harness error: %r" % exc)
        passed = not verdict.problems
        if case.expected_fail:
            verdict.status = "xpass" if passed else "xfail"
        else:
            verdict.status = "pass" if passed else "fail"
        verdict.secs = time.monotonic() - start
        return verdict

    def _equal(self, case, verdict):
        norm = Normalizer(case)
        caps = {"ref": self.runner.capture(case, self.ref, "a"), "cpp": self.runner.capture(case, self.cpp, "a")}
        for i, step in enumerate(case.steps):
            for side, cap in caps.items():
                if cap.steps[i]["timed_out"]:
                    verdict.problems.append("step %d timed out on %s (%.0fs)" % (i, side, step.timeout))
        views = {}
        for side, cap in caps.items():
            view = {"exit": "\n".join(str(s["rc"]) for s in cap.steps) + "\n"}
            view["stdout"] = "".join("--- step %d\n%s" % (i, s["stdout"]) for i, s in enumerate(cap.steps))
            view["stderr"] = "".join("--- step %d\n%s" % (i, s["stderr"]) for i, s in enumerate(cap.steps))
            view["vault"] = cap.vault
            view["files"] = cap.files
            views[side] = {ch: norm.apply(ch, txt, side) for ch, txt in view.items()}
        for ch in CHANNELS:
            if not case.compare[ch]:
                continue
            if views["ref"][ch] != views["cpp"][ch]:
                verdict.problems.extend(self._diff(ch, views["ref"][ch], views["cpp"][ch]))
        verdict.unused_rules = norm.unused()
        for spec in case.symbols_from_stdout:
            text = views["ref"]["stdout"].split("--- step %d\n" % spec["step"], 1)[-1]
            names = re.findall(spec["pattern"], text, re.M)
            if not names:
                verdict.problems.append("symbols_from_stdout matched nothing for %s" % spec["prefix"])
            verdict.extra_symbols += ["%s%s#%s" % (spec["prefix"], n, spec["facet"]) for n in names]

    @staticmethod
    def _diff(channel, ref, cpp):
        lines = list(difflib.unified_diff(ref.splitlines(), cpp.splitlines(), "tai.py:" + channel, "detz:" + channel, lineterm="", n=1))
        out = ["%s differs (%d diff lines)" % (channel, len(lines))]
        out.extend(lines[:MAX_DIFF_LINES])
        if len(lines) > MAX_DIFF_LINES:
            out.append("... %d more diff lines" % (len(lines) - MAX_DIFF_LINES))
        return out

    def _llm(self, case, verdict):
        checker = Checker(case)
        ratios = {}
        for name, side in (("ref", self.ref), ("cpp", self.cpp)):
            ok, last = 0, []
            for attempt in range(case.attempts):
                cap = self.runner.capture(case, side, "t%d" % attempt)
                fails = checker.evaluate(cap, name)
                for i, s in enumerate(cap.steps):
                    if s["timed_out"]:
                        fails.append("step %d timed out" % i)
                if fails:
                    last = fails
                else:
                    ok += 1
            ratios[name] = "%d/%d" % (ok, case.attempts)
            if ok < case.min_pass:
                verdict.problems.append("%s: %d/%d attempts passed, need %d" % (name, ok, case.attempts, case.min_pass))
                verdict.problems.extend("  " + f for f in last[:MAX_DIFF_LINES])
        verdict.ratio = ratios


class Report:
    def __init__(self, verdicts, meta):
        self.verdicts, self.meta = verdicts, meta

    def write(self):
        os.makedirs(OUT, exist_ok=True)
        suffix = "-partial" if self.meta["filtered"] else ""  # a filtered run never feeds the ledger
        with open(os.path.join(OUT, "diff-results%s.json" % suffix), "w") as fh:
            json.dump({"meta": self.meta, "results": [v.as_dict() for v in self.verdicts]}, fh, indent=1)
        counts = {}
        for v in self.verdicts:
            counts[v.status] = counts.get(v.status, 0) + 1
        lines = ["# Differential report (generated)", "",
                 "tai.py sha256 `%s`, detz sha256 `%s`" % (self.meta["tai_sha"][:16], self.meta["detz_sha"][:16]), "",
                 "Counts: " + ", ".join("%s=%d" % kv for kv in sorted(counts.items())), "",
                 "| case | status | llm | secs | note |", "|---|---|---|---|---|"]
        for v in sorted(self.verdicts, key=lambda x: x.case.id):
            note = (v.case.expected_fail or "") if v.status in ("xfail", "xpass") else ""
            if v.ratio:
                note += " ratio %s" % json.dumps(v.ratio)
            lines.append("| %s | %s | %s | %.1f | %s |" % (v.case.id, v.status, "yes" if v.case.llm else "", v.secs, note.strip()))
        lines += ["", "## Details of non-passing cases", ""]
        for v in sorted(self.verdicts, key=lambda x: x.case.id):
            if v.status in ("pass", "skipped"):
                continue
            lines += ["### %s (%s)" % (v.case.id, v.status), ""]
            if v.case.expected_fail:
                lines += ["expected_fail: " + v.case.expected_fail, ""]
            lines += ["```"] + v.problems[:MAX_DIFF_LINES * 3] + ["```", ""]
        unused = [(v.case.id, w) for v in self.verdicts for w in v.unused_rules]
        if unused:
            lines += ["## Unused normalization rules (suspect)", ""] + ["- %s: %s" % u for u in unused] + [""]
        with open(os.path.join(OUT, "diff-report%s.md" % suffix), "w") as fh:
            fh.write("\n".join(lines))
        return counts


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--detz", default=os.environ.get("DETZ_BIN", os.path.join(ROOT, "build", "detz")))
    ap.add_argument("--py", default=os.environ.get("TAI_PY", os.path.join(REPO, "tai.py")))
    ap.add_argument("--cases", default=CASES_DIR)
    ap.add_argument("--filter", default=None)
    ap.add_argument("--no-llm", action="store_true")
    ap.add_argument("--jobs", type=int, default=4)
    ap.add_argument("--list", action="store_true")
    args = ap.parse_args()
    cases = load_cases(args.cases)
    if args.filter:
        cases = [c for c in cases if args.filter in c.id]
    if args.list:
        for c in cases:
            print(c.id, "llm" if c.llm else "", c.expected_fail or "")
        return 0
    if not os.path.isfile(args.py):
        print("tai.py not found: %s" % args.py, file=sys.stderr)
        return 2
    detz_ok = os.path.isfile(args.detz) and os.access(args.detz, os.X_OK)
    if not detz_ok:
        print("detz binary missing: %s (every case will fail; build it with `make`)" % args.detz, file=sys.stderr)
    work = os.path.join(OUT, "work")
    shutil.rmtree(work, ignore_errors=True)
    runner = Runner(work)
    ref = Side("ref", [sys.executable, args.py])
    cpp = Side("cpp", [args.detz] if detz_ok else ["/nonexistent/detz"])
    executor = CaseExecutor(runner, ref, cpp)
    todo, verdicts = [], []
    for c in cases:
        if c.llm and args.no_llm:
            v = Verdict(c)
            v.status = "skipped"
            verdicts.append(v)
        else:
            todo.append(c)
    with concurrent.futures.ThreadPoolExecutor(max_workers=max(1, min(args.jobs, MAX_JOBS))) as pool:
        for v in pool.map(executor.execute, todo):
            print("%-7s %s" % (v.status.upper(), v.case.id), flush=True)
            verdicts.append(v)
    meta = {"tai_sha": sha256_file(args.py), "detz_sha": sha256_file(args.detz) if detz_ok else "missing",
            "detz_path": args.detz, "ts": time.strftime("%Y-%m-%dT%H:%M:%S"), "filtered": bool(args.filter),
            "no_llm": args.no_llm}
    counts = Report(verdicts, meta).write()
    print("counts:", counts, "(report dir: %s)" % OUT)
    bad = counts.get("fail", 0) + counts.get("xpass", 0)
    return 1 if bad else 0


if __name__ == "__main__":
    sys.exit(main())