#!/usr/bin/env python3
"""recount_on_the_kit.py: run this kit's recount over the kit itself, and check it.

    python3 recount_on_the_kit.py <the Reading the answer kit's data/ directory>

`recount_v2.py` is the recount as it ran over the benches' private record. It reads a git
checkout, and a few of the files it reads are held back from this kit. This script runs that
same code, unchanged, over the files that ARE public, and says which figures it could
re-derive and which it could not:

  1. It builds a scratch tree shaped like the record (`bench/<bench>/...`) out of links to this
     kit's folders, with the Reading the answer kit's files placed under
     `bench/jev-2026-09-21/`, where the recount looks for them.
  2. It loads `recount_v2.py` and swaps in three small things: a source object that answers
     "kit" where the original asked git for a commit; the six-way part of the cost series' kit
     lock (`KIT.lock.six-way.json`) where the original read the whole lock; and an empty
     series where the original read the cost series' own series file, which is held back.
  3. It runs `main()`'s own statements one at a time. A statement that needs a held-back file
     stops, and says so; the statements after it still run. The recount's five post-hoc sets
     run one by one for the same reason, and the memory set's own statements one at a time, so
     the one run whose receipts are held back (imajev-4B on card B, I-2, which read only
     held-back sets) does not take the other runs' memory rows with it. The page's own checks
     (`checked`) are one statement each in `main()`, so they need no help: the three that read
     held-back sets stop, and are reported as held back.
  4. It compares every figure it produced with `recount-v2.json`, the output over the full
     record, and prints what matched, what could only be partly re-derived, what differed and
     what could not be re-derived at all.

Standard library only; no model and no network. Exit 0 when no re-derived figure differs
from the recorded one.
"""
from __future__ import annotations

import sys

sys.dont_write_bytecode = True  # leave the kit as it came: no __pycache__ beside recount_v2.py

import ast  # noqa: E402
import importlib.util  # noqa: E402
import json  # noqa: E402
import os  # noqa: E402
import tempfile  # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))
LOCK_SIX_WAY = "bench/cost-of-intelligence-2026-09-23/KIT.lock.six-way.json"
SERIES = "bench/cost-of-intelligence-2026-09-23/series/series.json"
#: keys whose values name a file and the commit it was read at; a run over the kit reads the same
#: file at no commit, so a value is compared only after its "@<commit>" is dropped
CITE_KEYS = {"src", "kit", "kit_lock", "kit_order_src", "growth_src", "card_total_src"}
#: the one field of a runnable set that reads the held-back series file
HELD_FIELDS = {"post_hoc.author_family.cost_series_out_of_family": "the cost series' series.json"}
#: memory fields counted over EVERY file of a run, the held-back sets' telemetry and rows included,
#: so a run over the public files counts fewer: reported as partial, never compared
PARTIAL_FIELDS = {"samples_files", "samples", "longest_input_overall"}
#: rows of a set whose every input is held back here: reported as held back, never compared
HELD_ROWS = {"post_hoc.memory.rows.imajev-4b":
             "needs imajev-4B's card-B receipts (I-2), which are held back with the held-out, "
             "field-exam and judge sets, the only sets that run read",
             "checked.bound_on_0_of_12":
             "needs APUS-OpenJev-v1 9B's count of the doorman's harmless lines, whose rows are held back",
             "checked.heldout_interval_widths":
             "needs the held-out lines' counts, whose rows are held back",
             "checked.door_9b_against_the_live_door":
             "needs the doorman's rows and the 2026-09-13 live door's verdicts, which are held back"}


def link_tree(dst: str, src: str) -> None:
    """Link every file under src into dst, making directories as needed; never overwrite."""
    for root, _dirs, names in os.walk(src):
        rel = os.path.relpath(root, src)
        os.makedirs(os.path.join(dst, rel), exist_ok=True)
        for n in names:
            target = os.path.join(dst, rel, n)
            if not os.path.lexists(target):
                os.symlink(os.path.join(root, n), target)


def build_root(rta: str) -> str:
    root = tempfile.mkdtemp(prefix="recount-on-the-kit-")
    for name in sorted(os.listdir(HERE)):
        if os.path.isdir(os.path.join(HERE, name)) and name[-11:-6] == "-2026":
            link_tree(os.path.join(root, "bench", name), os.path.join(HERE, name))
    link_tree(os.path.join(root, "bench", "jev-2026-09-21"), rta)
    return root


def load_recount():
    spec = importlib.util.spec_from_file_location("recount_v2", os.path.join(HERE, "recount_v2.py"))
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)

    class KitSrc(mod.Src):
        def __init__(self, root):
            self.root = root
            self.head = "kit"

        def sha(self, rel):
            return "kit"

    read_json = mod.read_json

    def read_json_kit(src, rel):
        return {"rows": []} if rel == SERIES else read_json(src, rel)

    mod.Src = KitSrc
    mod.KIT_LOCK = LOCK_SIX_WAY
    mod.read_json = read_json_kit
    return mod


def why(exc: Exception, root: str) -> str:
    if isinstance(exc, FileNotFoundError) and exc.filename:
        name = os.path.relpath(exc.filename, os.path.join(root, "bench"))
        return "needs %s, which is held back" % name
    if isinstance(exc, KeyError):
        return "reads a count from a set whose rows are held back"
    if isinstance(exc, NameError):
        return "uses a value from a statement above that could not run"
    return "%s: %s" % (type(exc).__name__, exc)


def run_body(mod, source: str, name: str, args: dict, stopped: list, root: str):
    """Execute one of the recount's functions statement by statement, in a namespace of its own
    over the module's globals, and return what its `return` names. A statement that needs a
    held-back file stops, and says so; the statements after it still run."""
    tree = ast.parse(source)
    fn = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == name)
    ns = dict(mod.__dict__)
    ns.update(args)
    for stmt in fn.body:
        if isinstance(stmt, ast.Return):
            return eval(compile(ast.Expression(stmt.value), "recount_v2.py", "eval"), ns)
        code = compile(ast.Module(body=[stmt], type_ignores=[]), "recount_v2.py", "exec")
        try:
            exec(code, ns)
        except Exception as exc:  # noqa: BLE001 (reported, never hidden)
            head = ast.get_source_segment(source, stmt).splitlines()[0].strip()
            stopped.append(("%s line %d: %s" % (name, stmt.lineno, head[:36]), why(exc, root)))
    raise SystemExit("recount_on_the_kit: %s returned nothing" % name)


def run(mod, root: str, out: str, source: str) -> list[tuple[str, str]]:
    """Execute main()'s body statement by statement in the module's own namespace."""
    tree = ast.parse(source)
    main = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "main")
    ns = mod.__dict__
    sys.argv = ["recount_v2.py", root, out]
    stopped = []
    for stmt in main.body:
        # post_hoc(src, R, S0, D) builds five sets; run them one at a time so one held-back
        # file does not take the other four with it
        if (isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Call)
                and getattr(stmt.value.func, "id", None) == "post_hoc"):
            R, src = ns["R"], ns["src"]
            R["post_hoc"] = {"label": mod.POST_HOC}
            for key, call in (("door_by_kind", lambda: mod.post_hoc_door(src, R, ns["D"])),
                              ("six_order", lambda: mod.post_hoc_six_order(src, R)),
                              ("author_family", lambda: mod.post_hoc_author_family(src, ns["S0"])),
                              ("imajev_gpu_latency", lambda: mod.post_hoc_imajev_gpu(src)),
                              ("memory", lambda: run_body(mod, source, "post_hoc_memory",
                                                          {"src": src}, stopped, root))):
                try:
                    R["post_hoc"][key] = call()
                except Exception as exc:  # noqa: BLE001 (reported, never hidden)
                    stopped.append(("post_hoc: " + key, why(exc, root)))
            continue
        code = compile(ast.Module(body=[stmt], type_ignores=[]), "recount_v2.py", "exec")
        try:
            exec(code, ns)
        except Exception as exc:  # noqa: BLE001 (reported, never hidden)
            head = ast.get_source_segment(source, stmt).splitlines()[0].strip()
            stopped.append(("line %d: %s" % (stmt.lineno, head[:48]), why(exc, root)))
    return stopped


def uncite(v):
    return v.split("@", 1)[0] if isinstance(v, str) else v


def compare(mine, recorded, path, tally, notes):
    if path in HELD_FIELDS:
        tally["held back"] += 1
        notes.append(("held", "%s needs %s, which is held back" % (path, HELD_FIELDS[path])))
        return
    if path.startswith("post_hoc.memory.") and path.rsplit(".", 1)[-1] in PARTIAL_FIELDS:
        if mine != recorded:
            tally["partial"] += 1
            notes.append(("partial", "%s: recorded %r, from the public files %r"
                          % (path, recorded, mine)))
            return
    if isinstance(recorded, dict):
        if isinstance(mine, dict) and mine.get("missing") is True and "missing" not in recorded:
            tally["held back"] += 1
            return
        for k, v in recorded.items():
            if k in CITE_KEYS:
                continue
            sub = path + "." + k
            if sub in HELD_ROWS and (not isinstance(mine, dict) or k not in mine):
                tally["held back"] += 1
                notes.append(("held", "%s %s" % (sub, HELD_ROWS[sub])))
                continue
            if not isinstance(mine, dict) or k not in mine:
                tally["not run"] += 1
                continue
            compare(mine[k], v, path + "." + k, tally, notes)
        return
    if isinstance(recorded, list) and isinstance(mine, list) and len(mine) == len(recorded):
        for i, (a, b) in enumerate(zip(mine, recorded)):
            compare(a, b, "%s[%d]" % (path, i), tally, notes)
        return
    if uncite(mine) == uncite(recorded):
        tally["equal"] += 1
    else:
        tally["different"] += 1
        notes.append(("different", "%s: recorded %r, re-derived %r" % (path, recorded, mine)))


def main() -> int:
    if len(sys.argv) != 2 or not os.path.isfile(os.path.join(sys.argv[1], "kit", "task_c.json")):
        raise SystemExit(__doc__.split("\n\n")[1])
    rta = os.path.abspath(sys.argv[1])
    root = build_root(rta)
    source = open(os.path.join(HERE, "recount_v2.py"), encoding="utf-8").read()
    mod = load_recount()
    stopped = run(mod, root, os.path.join(root, "recount-on-the-kit.json"), source)
    mine = mod.__dict__["R"]
    recorded = json.load(open(os.path.join(HERE, "recount-v2.json"), encoding="utf-8"))

    print("recount_on_the_kit: recount_v2.py over this kit and the Reading the answer kit")
    print("  statements that could not run here:")
    for where, reason in stopped:
        print("    %-58s %s" % (where, reason))
    print("  %-14s %7s %8s %10s %10s %8s" % ("section", "equal", "partial", "different",
                                             "held back", "not run"))
    totals = {"equal": 0, "partial": 0, "different": 0, "held back": 0, "not run": 0}
    notes = []
    for section in sorted(k for k in recorded if k != "ref"):
        tally = {"equal": 0, "partial": 0, "different": 0, "held back": 0, "not run": 0}
        compare(mine.get(section, {}), recorded[section], section, tally, notes)
        for k in totals:
            totals[k] += tally[k]
        print("  %-14s %7d %8d %10d %10d %8d" % (section, tally["equal"], tally["partial"],
                                                 tally["different"], tally["held back"],
                                                 tally["not run"]))
    for kind, label in (("held", "held back"), ("partial", "partial (the recorded figure also "
                        "counted held-back files)"), ("different", "DIFFERENT")):
        rows = [text for k, text in notes if k == kind]
        if rows:
            print("  %s:" % label)
            for text in rows:
                print("    " + text)
    print("recount_on_the_kit: %s (%d equal, %d partial, %d held back, %d not run)"
          % ("no re-derived figure differs from recount-v2.json" if not totals["different"]
             else "%d figure(s) DIFFER from recount-v2.json" % totals["different"],
             totals["equal"], totals["partial"], totals["held back"], totals["not run"]))
    return 1 if totals["different"] else 0


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