#!/usr/bin/env python3
"""Convert a lane-CB JSON turn-command capture into the `.tcb` file `sots_turn` reads.

TWO FORMATS ON PURPOSE, AND THIS IS THE JOIN.

`tools/turncommands_capture.py` (lane CB) produces the **capture of record**: raw element words as
ground truth, heap vectors and strings the deep dump followed, the input save's hash bound to the
output autosaves' hashes, the per-client AI seeds, and the container self-check. That is what an
experiment should leave behind and none of it belongs in an engine's input file.

`.tcb` is the **engine's input**: line-oriented, no parser, nothing but the commands and their
provenance. `sots_turn --turn-commands` reads it and nothing else, so the engine never grows a
JSON reader and never has an opinion about how a capture was taken.

This script is the only thing that has to know both, and it is deliberately the narrow part:
it takes `decoded.wire` where lane CB's decoder produced one, falls back to raw words where it did
not, and writes `?` for every field neither could reach.

    tools/tcb_from_json.py CAPTURE.json [-o OUT.tcb] [--batch SEQ]

ONE FIELD MAPPING IS KNOWN WRONG AND IS OVERRIDDEN HERE, WITH ITS REASON.

List 5 (system rates) is decoded by lane CB as `{systemId, ship, terraform, sciences, ...}` --
the element's memory words mapped straight onto the frame's WIRE order. Lane RB measured that and
it is wrong: replaying it wrote the AI's one non-zero slider into `SRt` and regressed two leaves
on the canonical pair, where the oracle holds `SRsc = 1.0`. The only non-zero word of every list-5
element ever dumped is at memory index 2, and the same command on the wire -- where the frame is
NAMED -- puts its only non-zero in `SRsc`, the third member. So memory member 1 is wire member
`SRsc`, and the frame's memory order is not its wire order: ONE correspondence pinned, six unread.

Rather than carry a mapping that is known to be off by at least one, this converter emits the
system id and seven `?`. The replayer then counts the command -- the count is right either way --
and refuses to apply it, which is the correct behaviour for a payload nobody has read.

Settling it is one run: push two DIFFERENT sliders to two DIFFERENT values and read the
permutation off the element. Cheaper, a save taken after issuing rates carries the same command on
the wire with every field named, and needs no memory mapping at all.
"""
import argparse
import json
import struct
import sys

# Lists whose decoded `wire` this converter trusts.  List 5 is deliberately absent (see above);
# lists with no entry are carried as a single `?`, which counts the command and applies nothing.
TRUSTED = {
    3:  "iiii",      # ordinal, designId, systemId, trailing
    7:  "ii",        # shipId, trailing
    14: "iib",       # fleetId, mode, flag
}
# Lists whose element leads with scalars and then a counted vector the deep dump may have read.
VECTOR_TAIL = {
    8:  ("i", 1),    # fleetId, then the route
    10: ("ii", 2),   # systemId, fleetId, then a counted vector
}
UNMAPPED_HEAD = {
    5:  (1, 7),      # one trusted leading int (the system id) and seven unread fields
    23: (1, 1),      # the system id and a Population body behind a vftable
}


def as_int(v):
    if isinstance(v, bool):
        return 1 if v else 0
    if isinstance(v, float):
        return int(v)
    return int(v)


def tok_i(v):
    return "i%d" % as_int(v)


def fields_for(list_no, elem):
    """The `.tcb` field tokens for one element."""
    decoded = elem.get("decoded") or {}
    wire = decoded.get("wire")
    raw = elem.get("raw_words") or []
    vectors = {v["at_word"]: v for v in (elem.get("vectors") or [])}

    if list_no in UNMAPPED_HEAD:
        lead, unread = UNMAPPED_HEAD[list_no]
        head = [tok_i(raw[i]) for i in range(min(lead, len(raw)))]
        return head + ["?"] * unread if head else ["?"]

    if list_no in TRUSTED:
        spec = TRUSTED[list_no]
        if not wire or len(wire) < len(spec):
            return ["?"]
        out = []
        for i, kind in enumerate(spec):
            v = wire[i]
            out.append("b%d" % (1 if v else 0) if kind == "b" else tok_i(v))
        return out

    if list_no in VECTOR_TAIL:
        spec, vec_word = VECTOR_TAIL[list_no]
        if len(raw) < len(spec):
            return ["?"]
        out = [tok_i(struct.unpack("<i", struct.pack("<I", raw[i] & 0xffffffff))[0])
               for i in range(len(spec))]
        v = vectors.get(vec_word)
        if v is None:
            # The dump did not follow the pointer.  Its LENGTH is still derivable from the
            # begin/end pair, and a length with no values is exactly the honest field: the
            # command is counted and not applied.
            if len(raw) > vec_word + 1:
                out.append("v%d" % max(0, (raw[vec_word + 1] - raw[vec_word]) // 4))
            else:
                out.append("?")
        elif v.get("truncated"):
            out.append("v%d" % v["count"])
        else:
            vals = [struct.unpack("<i", struct.pack("<I", w & 0xffffffff))[0] for w in v["words"]]
            out.append("v%d:%s" % (len(vals), ",".join(str(x) for x in vals)) if vals
                       else "v0:")
        return out

    return ["?"]


def emit(cap, batch):
    src = cap.get("workload") or cap.get("note") or "lane-CB capture"
    lines = ["tcb 1",
             "meta source %s" % str(src).replace("\n", " "),
             "meta lane %s" % cap.get("lane", "?"),
             "meta guest %s" % cap.get("guest", "?"),
             "meta build %s" % cap.get("build", "?"),
             "meta captured %s" % cap.get("captured_utc", "?"),
             "meta batch seq=%s n=%s" % (batch.get("seq"), batch.get("n"))]
    binding = cap.get("binding") or {}
    if binding.get("input"):
        lines.append("meta input %s %s" % (binding["input"].get("file"),
                                           binding["input"].get("sha256")))
    for o in binding.get("outputs") or []:
        lines.append("meta output %s %s%s" % (o.get("file"), o.get("sha256"),
                                              "" if o.get("matches_oracle") is None
                                              else (" matches_oracle=%s" % o["matches_oracle"])))
    if cap.get("WARNING_INCOMPLETE"):
        lines.append("meta warning heap-payloads-absent")
    # The seeds. Without these the capture is a log file: the AI is one MT19937 per client seeded
    # with one word, so the block is the AI's answer and the seed is its input.
    for s in cap.get("ai_seeds") or []:
        val = s.get("used") or s.get("observed")
        if val is None:
            continue
        lines.append("seed %s %s%s" % (s.get("netId"), val,
                                       ""))
    for blk in batch.get("blocks") or []:
        lines.append("block %d %d" % (blk["index"], blk.get("playerId", blk.get("pid", 0))))
        idx = blk["index"]
        gates = blk.get("gates") or {}

        def payload(name):
            g = gates.get(name) or {}
            return g.get("payload") if g.get("set") else None

        p = payload("researchRate")
        if p is not None:
            lines.append("gate %d rate %s" % (idx, p))
        p = payload("researchTarget")
        if p is not None:
            name = (blk.get("researchTargetName") or "").strip()
            lines.append("gate %d target %s%s" % (idx, p, (" name %s" % name) if name else ""))
        p = payload("researchBoost")
        if p is not None:
            spend, _, frac = p.partition(",")
            lines.append("gate %d boost %s %s" % (idx, spend, frac or "0"))
        p = payload("group4")
        if p is not None:
            flag, _, val = p.partition(",")
            lines.append("gate %d group4 %s %s" % (idx, flag, val or "0"))
        p = payload("group5")
        if p is not None:
            a, b, c = (p.split(",") + ["0", "0", "0"])[:3]
            lines.append("gate %d group5 %s %s %s" % (idx, a, b, c))
        if (gates.get("civilianRatios") or {}).get("set"):
            lines.append("gate %d civilian" % idx)

        for lt in sorted(blk.get("lists") or [], key=lambda l: l["list"]):
            n = lt.get("size", 0)
            if not n:
                continue
            lines.append("list %d %d %d" % (idx, lt["list"], n))
            by_index = {e["index"]: e for e in lt.get("elements") or []}
            for e in range(n):
                elem = by_index.get(e)
                f = fields_for(lt["list"], elem) if elem else ["?"]
                lines.append("elem %d %d %d %s" % (idx, lt["list"], e, " ".join(f)))
    return "\n".join(lines) + "\n"


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("capture")
    ap.add_argument("-o", "--out")
    ap.add_argument("--batch", type=int,
                    help="which batch seq to convert (default: the last, which is the End-Turn "
                         "submission; a process also applies a batch at LOAD and replaying that "
                         "one against a save written after it would double-count)")
    a = ap.parse_args(argv)

    with open(a.capture) as f:
        cap = json.load(f)
    batches = cap.get("batches") or []
    if not batches:
        print("no batches in %s" % a.capture, file=sys.stderr)
        return 2
    if a.batch is not None:
        picked = [b for b in batches if b.get("seq") == a.batch]
        if not picked:
            print("no batch seq=%d; the capture holds %s"
                  % (a.batch, [b.get("seq") for b in batches]), file=sys.stderr)
            return 2
        batch = picked[0]
    else:
        batch = batches[-1]

    text = emit(cap, batch)
    if a.out:
        with open(a.out, "w") as f:
            f.write(text)
        print("wrote %s (batch seq=%s, %d blocks)"
              % (a.out, batch.get("seq"), len(batch.get("blocks") or [])))
    else:
        sys.stdout.write(text)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())