#!/usr/bin/env python3
"""Turn a lane-CB `shim.aiorders.txt` into a `TurnCommands` capture (JSON) for the replayer.

WHY THIS IS A SEPARATE, OFFLINE STEP.

The shim's block dump reads element BYTES, not element TYPES -- lane L4's design point, kept
deliberately: a wrong element record then shows up as a wrong value in a report instead of being
baked into an instrument that needs a VM run to correct. This script is where the typing happens,
so a corrected record costs a re-parse rather than a rebuild, a redeploy and a turn.

The consequence for anyone reading the output: `raw_words` is the GROUND TRUTH. `decoded` is a
best-effort typing, present only for the six lists the campaign has actually observed, and
regenerable from `raw_words` by re-running this script. If the two ever disagree, `raw_words` wins.

THE ONE THING THAT IS EASY TO GET WRONG. List 3's in-memory element is in the OPPOSITE order from
its wire record: its writer (0x00822870) emits +0x14, +0x10, +0x0c, +0x08, descending, while lists
5, 8, 10, 14 and 23 all write ascending. That is per-list, not a rule, and it is encoded as such
below (`REVERSED`). A decoder that applies one order to all six produces a build order whose
design id and ordinal are swapped -- which is a plausible-looking save with every id wrong.

Usage:
    uv run python3 tools/turncommands_capture.py \
        --log   verify/results/shim/aiorders/cb-turn2to3.txt \
        --input verify/results/saves/turn2-state.sav \
        --autosave "(Autosave EndTurn).sav=<sha256>:<bytes>" \
        --autosave "(Autosave).sav=<sha256>:<bytes>" \
        --out   verify/results/turncommands/cb-turn2to3.json
"""
import argparse, datetime, hashlib, json, os, re, struct, sys

CAPTURE_VERSION = 1

# ---------------------------------------------------------------------------------------------
# element records
# ---------------------------------------------------------------------------------------------
# Each entry: (record name, field names in WIRE order, whether the memory words run backwards,
# and which word index -- if any -- carries a counted heap vector whose contents belong on the
# wire in place of a scalar).
#
# Lists with no entry here are carried as raw words only. That is deliberate (method rule 6): no
# workload in this campaign has ever put an element in them, so any record would be a hypothesis,
# and a hypothesis written into a replayer's input file is indistinguishable from a fact.
REVERSED = {3}

# Lists whose element VALUES are certain but whose memory-to-wire mapping is NOT. Lane RB mapped
# list 5's sliders straight through and regressed two leaves; until the writer's own order is read,
# a capture must carry the values without asserting where they land. In the `.tcb` these words are
# emitted as `?` -- byte-for-byte what lane RB's own converter emits -- with the memory-order values
# on a `meta uncertain` line, so nothing is lost and nothing is claimed.
UNCERTAIN_WIRE_ORDER = {5}
# What IS known about list 5, recorded next to the values rather than baked into a field name.
RATES_NOTE = ("memory member 1 is wire member SRsc (lane RB, from a replay that regressed two "
              "leaves when the memory order was used as the wire order); the other six memory "
              "members are UNREAD. Values are certain, positions are not.")

RECORDS = {
    1:  ("NewDesign",   None),                       # polymorphic; carried raw + its strings
    3:  ("BuildOrder",  ["ordinal", "designId", "systemId", "trailing"]),
    # NO RATES RECORD, ON PURPOSE -- the system id only.
    #
    # This element is dumped in MEMORY order and its memory field order is NOT its wire order.
    # Lane Q named the wire frame {SRs, SRt, SRsc, SRtf, SRi, SRoh, SRnr} from a HUMAN orders save;
    # lane RB mapped the AI's memory words straight onto it, predicted zero regressions on its
    # first replay and got TWO -- the AI's one non-zero slider landed in `SRt` where the oracle
    # holds `SRsc = 1.0`. The one mapping that is known is **memory member 1 -> wire member SRsc**;
    # the other six are unread.
    #
    # A capture is supposed to outlive whatever adapter happened to read it, so this ships the
    # values and NOT a naming. `rates_memory_order` carries the seven words as dumped, `wire` is
    # None, and a replayer that needs them must refuse rather than guess. A missing field is
    # honest; a mislabelled one propagates.
    5:  ("SystemRates", ["systemId"]),
    7:  ("Colonize",    ["shipId", "trailing"]),
    8:  ("FleetMove",   ["fleetId", "@route"]),      # `@` = the counted vector at this position
    10: ("List10",      ["systemId", "fleetId", "@counted"]),
    14: ("FleetTask",   ["fleetId", "mode", "flag"]),
    # CORRECTED from lane L4 §0, which read a trailing int here ("Population{vptr, vector(24 B),
    # -1}"). There is no trailing int. Two runs of THIS turn on two guests -- lane L4's on VM145
    # and lane CB's on VM146, both of which wrote byte-identical autosaves -- disagree on that
    # word (-1 vs 0x0035765f, an allocator cookie that also shows up mid-element in list 8).
    # A word that differs between two runs whose outputs are byte-identical cannot be a word the
    # applier reads. The element is {systemId, vptr, vector}.
    23: ("Population",  ["systemId", "@vptr", "@body"]),
}
# Which fields of which record are floats rather than ints, by wire-field name.
FLOAT_FIELDS = set()   # the only float frame in the corpus is list 5's, and it is not named
# The bool is ONE BYTE on the wire; the other three bytes of its word are padding and carry
# heap garbage that differs run to run (0x00B5B601 vs 0x01138601 on two runs that produced
# byte-identical autosaves). Masking to the low byte is not tidiness, it is the record.
BOOL_FIELDS = {("FleetTask", "flag")}

GATES = [
    ("researchRate",    "rate"),
    ("researchTarget",  "target"),
    ("researchBoost",   "boost"),
    ("group4",          "g4"),
    ("group5",          "f3"),
    ("civilianRatios",  "civ"),
]


def u32_to_f32(v):
    return struct.unpack("<f", struct.pack("<I", v & 0xFFFFFFFF))[0]


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


# ---------------------------------------------------------------------------------------------
# the log
# ---------------------------------------------------------------------------------------------
RE_CFG   = re.compile(r"^aicfg (.*)$")
RE_SEED  = re.compile(r"^aiseed call=(\d+) netId=(-?\d+) observed=0x([0-9a-f]+) "
                      r"used=0x([0-9a-f]+) pinned=(\d)")
RE_BATCH = re.compile(r"^---- aibatch seq=(\d+) blocks=0x([0-9a-f]+) n=(-?\d+) stride=0x([0-9a-f]+)")
RE_BLK   = re.compile(r"^aiblk seq=(\d+) blk=(\d+)/(\d+) at=0x([0-9a-f]+) pid=(-?\d+) (.*)$")
RE_SIZES = re.compile(r"^ailists seq=(\d+) blk=(\d+) pid=(-?\d+) nonEmpty=(\d+) "
                      r"sizes\(1\.\.27\)=\[ (.*?)\]")
RE_LIST  = re.compile(r"^ailist blk=(\d+) pid=(-?\d+) list=(\d+) off=0x([0-9a-f]+) size=(\d+)")
RE_MISM  = re.compile(r"^ailist MISMATCH blk=(\d+) pid=(-?\d+) list=(\d+) _Mysize=(\d+) walked=(\d+)")
# `words=` is lane CB's; a lane-L4 log has no such field and is still parsed, at its fixed 12.
RE_ELEM  = re.compile(r"^aielem blk=(\d+) pid=(-?\d+) list=(\d+) idx=(\d+) node=0x([0-9a-f]+) "
                      r"(?:words=(\d+) )?f0=\S+ f1=\S+ ints=\[ .*?\] hex=\[ (.*?)\]")
RE_VEC   = re.compile(r"^aivec blk=(\d+) pid=(-?\d+) list=(\d+) idx=(\d+) at=w(\d+) "
                      r"first=0x([0-9a-f]+) cap=(\d+) count=(\d+) ints=\[ .*?\] hex=\[ (.*?)\]"
                      r"( TRUNCATED)?")
RE_STR   = re.compile(r'^aistr blk=(\d+) pid=(-?\d+) list=(\d+) idx=(\d+) at=w(\d+) sso=(\d) '
                      r'len=(\d+) text="(.*)"$')


def parse_gates(rest):
    """`rate=1:0.8 target=0:880612328 boost=0:0,0 g4=0:0,0 f3=0:1,1,1 civ=0` -> a dict."""
    out = {}
    fields = dict(tok.split("=", 1) for tok in rest.split() if "=" in tok)
    for name, key in GATES:
        raw = fields.get(key)
        if raw is None:
            continue
        if ":" in raw:
            setbit, payload = raw.split(":", 1)
            out[name] = {"set": setbit == "1", "payload": payload}
        else:
            out[name] = {"set": raw == "1"}
    return out


def hexwords(s):
    return [int(w, 16) for w in s.split()]


def decode(list_no, words, vectors, strings):
    """Type one element. Returns None where the campaign has no record for the list."""
    entry = RECORDS.get(list_no)
    if entry is None:
        return None
    name, fields = entry
    if fields is None:
        return {"record": name, "wire": None,
                "note": "polymorphic object; carried as raw words and decoded strings only"}
    src = list(reversed(words[:len(fields)])) if list_no in REVERSED else words[:len(fields)]
    byword = {v["at_word"]: v for v in vectors}
    out, wire, wi = {"record": name}, [], 0
    for i, f in enumerate(fields):
        if f.startswith("@"):
            key = f[1:]
            if key == "vptr":
                out[key] = f"0x{words[i]:08x}"
                wi += 1
                continue
            v = byword.get(wi)
            if v is None:
                out[key] = None
                out.setdefault("incomplete", []).append(key)
            else:
                out[key] = v["words"]
                wire.append(len(v["words"]))
                wire.extend(v["words"])
            wi += 3  # a vector occupies three words of the element
            continue
        val = src[i] if list_no in REVERSED else words[wi]
        if (name, f) in FLOAT_FIELDS:
            val = u32_to_f32(val)
        elif (name, f) in BOOL_FIELDS:
            val = bool(val & 0xFF)
        elif val >= 0x80000000:
            val -= 1 << 32
        out[f] = val
        wire.append(val)
        wi += 1
    out["wire"] = wire
    # How many words of `raw_words` are the element. The window is deliberately wider than any
    # record (`aiorders.words=`), so everything past this is the NEXT heap node and is noise --
    # which is why the widened window also produces spurious `vectors` entries at word offsets
    # no record names. Measured independently: see the cross-run stability table in the report.
    out["record_words"] = wi
    if list_no == 5:
        # The seven words after the system id, exactly as dumped, as raw u32 and as f32.
        out["rates_memory_order_u32"] = words[1:8]
        out["rates_memory_order_f32"] = [u32_to_f32(w) for w in words[1:8]]
        out["record_words"] = 8
        out["wire"] = None
        out["wire_order_unread"] = True
        out["note"] = RATES_NOTE
        return out
    if list_no in UNCERTAIN_WIRE_ORDER:
        out["wire_order_uncertain"] = True
        out["wire_order_note"] = (
            "the VALUES are measured; their positions in the wire frame are NOT. `wire` is memory "
            "order. A replayer must not map these onto named members without reading the writer.")
    if strings:
        out["strings"] = [s["text"] for s in strings]
    return out


def parse_log(path):
    cfg, seeds, batches = {}, [], []
    batch = block = None
    elems = {}          # (blk, list, idx) -> element dict, so aivec/aistr can attach
    mismatches = []
    with open(path, "r", errors="replace") as fh:
        for line in fh:
            line = line.rstrip("\n")

            m = RE_CFG.match(line)
            if m:
                cfg.update(dict(t.split("=", 1) for t in m.group(1).split() if "=" in t))
                continue

            m = RE_SEED.match(line)
            if m:
                seeds.append({"call": int(m.group(1)), "netId": int(m.group(2)),
                              "observed": f"0x{m.group(3)}", "used": f"0x{m.group(4)}",
                              "pinned": m.group(5) == "1"})
                continue

            m = RE_BATCH.match(line)
            if m:
                batch = {"seq": int(m.group(1)), "blocks_va": f"0x{m.group(2)}",
                         "n": int(m.group(3)), "stride": int(m.group(4), 16), "blocks": []}
                batches.append(batch)
                elems = {}
                continue

            m = RE_BLK.match(line)
            if m and batch is not None:
                block = {"index": int(m.group(2)), "n": int(m.group(3)),
                         "block_va": f"0x{m.group(4)}", "playerId": int(m.group(5)),
                         "gates": parse_gates(m.group(6)), "lists": []}
                batch["blocks"].append(block)
                continue

            m = RE_SIZES.match(line)
            if m and block is not None:
                block["list_sizes"] = [int(x) for x in m.group(5).split()]
                continue

            m = RE_MISM.match(line)
            if m:
                mismatches.append({"blk": int(m.group(1)), "list": int(m.group(3)),
                                   "mysize": int(m.group(4)), "walked": int(m.group(5))})
                continue

            m = RE_LIST.match(line)
            if m and batch is not None:
                blk = int(m.group(1))
                tgt = next(b for b in batch["blocks"] if b["index"] == blk)
                tgt["lists"].append({"list": int(m.group(3)),
                                     "member_offset": f"0x{m.group(4)}",
                                     "size": int(m.group(5)), "elements": []})
                continue

            m = RE_ELEM.match(line)
            if m and batch is not None:
                blk, lst, idx = int(m.group(1)), int(m.group(3)), int(m.group(4))
                e = {"index": idx, "node_va": f"0x{m.group(5)}",
                     "raw_words": hexwords(m.group(7)), "vectors": [], "strings": []}
                tgt = next(b for b in batch["blocks"] if b["index"] == blk)
                lt = next(l for l in tgt["lists"] if l["list"] == lst)
                lt["elements"].append(e)
                elems[(blk, lst, idx)] = e
                continue

            m = RE_VEC.match(line)
            if m:
                e = elems.get((int(m.group(1)), int(m.group(3)), int(m.group(4))))
                if e is not None:
                    e["vectors"].append({"at_word": int(m.group(5)),
                                         "first_va": f"0x{m.group(6)}",
                                         "capacity": int(m.group(7)),
                                         "count": int(m.group(8)),
                                         "words": hexwords(m.group(9)),
                                         "truncated": bool(m.group(10))})
                continue

            m = RE_STR.match(line)
            if m:
                e = elems.get((int(m.group(1)), int(m.group(3)), int(m.group(4))))
                if e is not None:
                    e["strings"].append({"at_word": int(m.group(5)),
                                         "sso": m.group(6) == "1",
                                         "len": int(m.group(7)), "text": m.group(8)})
                continue

    for b in batches:
        for blk in b["blocks"]:
            for lt in blk["lists"]:
                for e in lt["elements"]:
                    d = decode(lt["list"], e["raw_words"], e["vectors"], e["strings"])
                    if d is not None:
                        e["decoded"] = d
    return cfg, seeds, batches, mismatches



# ---------------------------------------------------------------------------------------------
# lane RB's `.tcb`, with the payloads filled in
# ---------------------------------------------------------------------------------------------
# Lane RB landed `tools/aiorders_to_tcb.py` and a line-oriented `.tcb` grammar before this lane
# landed its JSON, and `sots_turn --turn-commands` reads `.tcb`. Rather than ask RB to read a
# second format, this emits theirs -- with ONE additive extension, flagged here because a silent
# superset is how two tools drift apart:
#
#     `vN`      a vector of N elements whose CONTENTS ARE UNKNOWN   (RB's; the dump could not follow
#               the pointer)
#     `vN:a,b`  a vector of N elements whose contents ARE known     (lane CB's; the dump followed it)
#
# A reader that only knows `vN` must treat `vN:...` as unreadable rather than as `vN`, because the
# whole point of RB's `vN` is that the replayer counts such a command and REFUSES to apply it. The
# two mean opposite things about whether the command can be reproduced, and getting that backwards
# turns "we cannot replay this" into "we replayed it wrong".
#
# `?` is preserved exactly: a word the record says is part of an object this window cannot read.
def emit_tcb(cap, source):
    out = ["tcb 1",
           "meta source %s" % source,
           "meta input %s" % cap["binding"]["input"]["file"],
           "meta input_sha256 %s" % cap["binding"]["input"]["sha256"]]
    for o in cap["binding"]["outputs"]:
        out.append("meta output %s %d %s%s" % (
            o["file"], o["bytes"], o["sha256"],
            "" if "matches_oracle" not in o else (" oracle=%s" % o["matches_oracle"])))
    if cap["binding"].get("control_run"):
        out.append("meta control %s" % cap["binding"]["control_run"])
    out.append("meta lane CB capture_version %d build %s guest %s"
               % (cap["capture_version"], cap.get("build"), cap.get("guest")))
    out.append("meta extension vN:a,b -- a vector whose contents ARE known; bare vN still means "
               "unknown and must still be refused")

    # The End-Turn batch is the LAST one: a process applies a batch at LOAD as well, and replaying
    # the load-time one against a save written after it would double-count. RB's default, kept.
    batch = cap["batches"][-1]
    out.append("meta batch seq=%d n=%d" % (batch["seq"], batch["n"]))
    for s in cap.get("ai_seeds", []):
        out.append("seed %d %s%s" % (s["netId"], s["used"], " pinned" if s["pinned"] else ""))

    for blk in batch["blocks"]:
        out.append("block %d %d" % (blk["index"], blk["playerId"]))
        g = blk["gates"]
        if g.get("researchRate", {}).get("set"):
            out.append("gate %d rate %s" % (blk["index"], g["researchRate"]["payload"]))
        if g.get("researchTarget", {}).get("set"):
            out.append("gate %d target %s" % (blk["index"], g["researchTarget"]["payload"]))
        if g.get("researchBoost", {}).get("set"):
            out.append("gate %d boost %s" % (blk["index"], g["researchBoost"]["payload"]))
        for lt in blk["lists"]:
            out.append("list %d %d %d" % (blk["index"], lt["list"], lt["size"]))
            for e in lt["elements"]:
                out.append("elem %d %d %d %s"
                           % (blk["index"], lt["list"], e["index"], tcb_tokens(lt["list"], e)))
                d = e.get("decoded") or {}
                for s in e.get("strings", []):
                    # List 1's element is a polymorphic ShipDesignDef the window cannot type, but
                    # its NAME is inline (short-string optimisation) and is the only human-readable
                    # thing in the whole block. It is recorded, not decoded.
                    out.append('meta string %d %d %d w%d "%s"'
                               % (blk["index"], lt["list"], e["index"], s["at_word"], s["text"]))
                if d.get("wire_order_unread"):
                    out.append("meta rates %d %d %d memory-order-f32 %s"
                               % (blk["index"], lt["list"], e["index"],
                                  ",".join(repr(v) for v in d["rates_memory_order_f32"])))
                    out.append("meta rates-note %s" % RATES_NOTE)
    return "\n".join(out) + "\n"


def tcb_tokens(list_no, e):
    entry = RECORDS.get(list_no)
    d = e.get("decoded")
    if entry is None or entry[1] is None or d is None:
        return "?"
    _, fields = entry
    name = entry[0]
    if list_no == 5:
        # Byte-for-byte what lane RB's own converter emits: the id, then seven unreadable words.
        return "i%d ? ? ? ? ? ? ?" % d["systemId"]
    byword = {v["at_word"]: v for v in e["vectors"]}
    toks, wi = [], 0
    for i, f in enumerate(fields):
        if f.startswith("@"):
            key = f[1:]
            if key == "vptr":
                toks.append("?")          # a vftable pointer is not a wire field
                wi += 1
                continue
            v = byword.get(wi)
            if v is None:
                toks.append("?")
            else:
                toks.append("v%d:%s" % (v["count"], ",".join(str(w) for w in v["words"])))
            wi += 3
            continue
        val = d.get(f)

        if (name, f) in BOOL_FIELDS:
            toks.append("b%d" % (1 if val else 0))
        elif (name, f) in FLOAT_FIELDS:
            toks.append("f%s" % repr(float(val)))
        elif val is None:
            toks.append("?")
        else:
            toks.append("i%d" % val)
        wi += 1
    return " ".join(toks)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--log", required=True)
    ap.add_argument("--input", required=True, help="the save the turn was run from")
    ap.add_argument("--autosave", action="append", default=[],
                    help="NAME=PATH -- an output autosave of THIS run")
    ap.add_argument("--oracle", action="append", default=[],
                    help="NAME=SHA256 -- the published oracle hash, if this pair has one")
    ap.add_argument("--control", default=None,
                    help="name of the hooks=off run whose autosaves this run's must match")
    ap.add_argument("--guest", default="VM146")
    ap.add_argument("--build", default=None)
    ap.add_argument("--workload", default=None)
    ap.add_argument("--note", default=None)
    ap.add_argument("--out", required=True)
    ap.add_argument("--tcb", default=None,
                    help="also emit lane RB's .tcb grammar, with the heap payloads "
                         "filled in (see emit_tcb for the one extension)")
    a = ap.parse_args()

    cfg, seeds, batches, mismatches = parse_log(a.log)
    oracle = dict(x.split("=", 1) for x in a.oracle)

    outputs = []
    for spec in a.autosave:
        name, path = spec.split("=", 1)
        h = sha256_of(path)
        row = {"file": name, "bytes": os.path.getsize(path), "sha256": h,
               "path": os.path.relpath(path, start=os.path.dirname(os.path.abspath(a.out)))}
        if name in oracle:
            row["oracle_sha256"] = oracle[name]
            row["matches_oracle"] = (h == oracle[name])
        outputs.append(row)

    cap = {
        "capture_version": CAPTURE_VERSION,
        "lane": "CB",
        "guest": a.guest,
        "build": a.build,
        "captured_utc": datetime.datetime.now(datetime.timezone.utc)
                                 .replace(microsecond=0).isoformat(),
        "instrument": cfg,
        "workload": a.workload,
        "note": a.note,
        "binding": {
            "input": {"file": os.path.basename(a.input),
                      "bytes": os.path.getsize(a.input),
                      "sha256": sha256_of(a.input)},
            "outputs": outputs,
            "control_run": a.control,
        },
        # Three words per process (lane L1). WITHOUT THESE THE CAPTURE IS A LOG FILE: the AI is
        # MT19937 from one word per client, so the block is the AI's answer and the seed is its
        # input. `pinned` says whether the run FORCED the value rather than observing it.
        "ai_seeds": seeds,
        "list_size_self_check": {
            "mismatches": mismatches,
            "note": ("every list is measured twice, by walking its node chain and by reading "
                     "_Mysize; an empty list here means the container layout agreed on every "
                     "list of every block, which is the evidence that a zero size is a real "
                     "zero rather than a wrong offset"),
        },
        "batches": batches,
    }

    if not outputs:
        cap["binding"]["WARNING"] = ("no output autosave recorded -- this capture CANNOT be "
                                     "checked by a byte-match and must not be used as one")
    if cfg.get("deep") != "1":
        cap["WARNING_INCOMPLETE"] = ("captured with aiorders.deep=0: every heap payload (list 8's "
                                     "route, list 10's counted vector, list 23's body) is ABSENT, "
                                     "not empty")

    os.makedirs(os.path.dirname(os.path.abspath(a.out)), exist_ok=True)
    with open(a.out, "w") as fh:
        json.dump(cap, fh, indent=1, sort_keys=False)
        fh.write("\n")
    if a.tcb:
        os.makedirs(os.path.dirname(os.path.abspath(a.tcb)), exist_ok=True)
        with open(a.tcb, "w") as fh:
            fh.write(emit_tcb(cap, os.path.basename(a.log)))
        print("wrote", a.tcb)

    nblocks = sum(len(b["blocks"]) for b in batches)
    nelem = sum(len(l["elements"]) for b in batches for k in b["blocks"] for l in k["lists"])
    print(f"wrote {a.out}: {len(batches)} batches, {nblocks} blocks, {nelem} elements, "
          f"{len(seeds)} seeds, {len(mismatches)} layout mismatches")


if __name__ == "__main__":
    main()