#!/usr/bin/env python3

"""Prove a harness change is measurement-identical (ADR-0005).

Compares two result JSONL files row-by-row on the fields that must not move

when only harness code changed: the commands that ran, the operations

recorded, tool-call accounting, and agent-emitted token estimates (command

text is deterministic; output tokens and timings legitimately vary run to

run). Coverage changes (operations present in only one file) are reported

separately from drift.

This codifies the verification used for the oakbench extraction: before

trusting a refactor, run the same lane before and after and require zero

invariant drift.

Usage:

    python3 scripts/row_parity.py old/latest.jsonl new/latest.jsonl

    python3 scripts/row_parity.py old.jsonl new.jsonl --strict-coverage

Exit codes: 0 parity, 1 invariant drift (or coverage change with

--strict-coverage), 2 usage/input error.

"""

from __future__ import annotations

import argparse

import json

import re

from pathlib import Path

from typing import Any

# Run-timestamp directory segments (workdirs are keyed by run id, e.g.

# runs/20260611T194547Z/...) legitimately differ between the two runs being

# compared; commands embedding such paths (clone destinations, bare remotes)

# are normalized before comparison so the run id itself never reads as drift.

RUN_TS_PATTERN = re.compile(r"\d{8}T\d{6}Z")

# Disposable remote branch names (oakbench.remotes.disposable_branch) end in

# per-run pid/token uniquifiers that must never collide on the server — and

# therefore never match across two runs either.

BENCH_BRANCH_PATTERN = re.compile(r"bench-[\w.-]+-r\d+-[0-9a-f]+-[0-9a-f]+")

def normalize_value(value: Any) -> Any:

    if isinstance(value, str):

        value = RUN_TS_PATTERN.sub("<RUN_TS>", value)

        return BENCH_BRANCH_PATTERN.sub("<BENCH_BRANCH>", value)

    if isinstance(value, list):

        return [normalize_value(item) for item in value]

    return value

# Fields that must be identical for the same row identity when only harness

# code changed. Timings, output bytes, and output-derived token estimates are

# deliberately NOT here: they vary run to run.

INVARIANT_FIELDS = (

    "command",

    "subject_kind",

    "benchmark_track",

    "command_semantics_version",

    "git_mode",

    "step_kind",

    "phase",

    "expected_returncodes",

    "tool_call_count",

    "vcs_tool_call_count",

    "terminal_tool_call_count",

    "test_tool_call_count",

    "estimated_tokens_input",

    "estimated_tokens_agent_emitted",

)

# Fields where drift is reported as a warning, not a failure: deterministic in

# most scenarios but legitimately variable in some (e.g. retry-dependent).

WARN_FIELDS = ("returncode", "process_returncode")

IDENTITY_FIELDS = ("scenario", "subject", "operation", "run", "phase", "workflow", "worker", "task_id")

def parse_args() -> argparse.Namespace:

    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)

    parser.add_argument("old", type=Path, help="Result JSONL from before the change.")

    parser.add_argument("new", type=Path, help="Result JSONL from after the change.")

    parser.add_argument(

        "--strict-coverage",

        action="store_true",

        help="Also fail when row identities appear in only one file (default: report only).",

    )

    parser.add_argument(

        "--ignore-fields",

        default="",

        help="Comma-separated invariant fields to skip, e.g. when a change intentionally moves them.",

    )

    parser.add_argument("--limit", type=int, default=30, help="Max differences to print per section.")

    return parser.parse_args()

def load_rows(path: Path) -> list[dict[str, Any]]:

    rows: list[dict[str, Any]] = []

    for line_number, line in enumerate(path.read_text().splitlines(), 1):

        if not line.strip():

            continue

        try:

            row = json.loads(line)

        except json.JSONDecodeError as exc:

            raise SystemExit(f"{path}:{line_number}: invalid JSON: {exc}")

        if isinstance(row, dict):

            rows.append(row)

    return rows

def identity(row: dict[str, Any]) -> tuple:

    return tuple(json.dumps(row.get(field), sort_keys=True) for field in IDENTITY_FIELDS)

def identity_display(row: dict[str, Any]) -> str:

    parts = [f"{field}={row.get(field)}" for field in IDENTITY_FIELDS if row.get(field) is not None]

    return " ".join(parts)

def index_rows(rows: list[dict[str, Any]], label: str) -> dict[tuple, dict[str, Any]]:

    indexed: dict[tuple, dict[str, Any]] = {}

    duplicates = 0

    for row in rows:

        key = identity(row)

        if key in indexed:

            duplicates += 1

        indexed[key] = row

    if duplicates:

        print(f"note: {label} contains {duplicates} duplicate row identities; last occurrence wins")

    return indexed

def main() -> int:

    args = parse_args()

    for path in (args.old, args.new):

        if not path.exists():

            print(f"input not found: {path}")

            return 2

    ignored = {item.strip() for item in args.ignore_fields.split(",") if item.strip()}

    fields = [field for field in INVARIANT_FIELDS if field not in ignored]

    old_rows = index_rows(load_rows(args.old), "old")

    new_rows = index_rows(load_rows(args.new), "new")

    shared = [key for key in old_rows if key in new_rows]

    only_old = [key for key in old_rows if key not in new_rows]

    only_new = [key for key in new_rows if key not in old_rows]

    drift: list[str] = []

    warnings: list[str] = []

    for key in shared:

        old_row, new_row = old_rows[key], new_rows[key]

        for field in fields:

            if field in old_row or field in new_row:

                if normalize_value(old_row.get(field)) != normalize_value(new_row.get(field)):

                    drift.append(

                        f"{identity_display(old_row)} :: {field}: {old_row.get(field)!r} -> {new_row.get(field)!r}"

                    )

        for field in WARN_FIELDS:

            if field in old_row and old_row.get(field) != new_row.get(field):

                warnings.append(

                    f"{identity_display(old_row)} :: {field}: {old_row.get(field)!r} -> {new_row.get(field)!r}"

                )

    print("# Row Parity Report")

    print()

    print(f"- Old: `{args.old}` ({len(old_rows)} unique row identities)")

    print(f"- New: `{args.new}` ({len(new_rows)} unique row identities)")

    print(f"- Shared identities compared: {len(shared)}")

    print(f"- Invariant fields checked: {', '.join(fields)}")

    print()

    def section(title: str, items: list[str]) -> None:

        print(f"## {title} ({len(items)})")

        print()

        for item in items[: args.limit]:

            print(f"- {item}")

        if len(items) > args.limit:

            print(f"- ... and {len(items) - args.limit} more")

        if not items:

            print("- none")

        print()

    section("Invariant drift (FAIL)", drift)

    section("Returncode drift (warning)", warnings)

    section("Coverage only in old", [identity_display(old_rows[k]) for k in only_old])

    section("Coverage only in new", [identity_display(new_rows[k]) for k in only_new])

    failed = bool(drift) or (args.strict_coverage and (only_old or only_new))

    print(f"VERDICT: {'DRIFT' if failed else 'PARITY'}")

    return 1 if failed else 0

if __name__ == "__main__":

    raise SystemExit(main())