#!/usr/bin/env python3
"""Generate conformance vectors for the batch extension to wire v1.

A THIRD file, for the reason the tracked file is a second one: `wire_v1_vectors.json`
is frozen and is read by five polyglot demo clients whose decoders know message
types 0x01..0x05, and a decoder that rejects an 0x08 frame is CORRECT. A client
that does not implement batching stays conformant by ignoring this file.

The vectors are generated by an implementation that shares no code with weir:
`zlib.crc32` is the IEEE / ISO-3309 CRC-32 weir uses, and the bit arithmetic
below is written out longhand rather than imported. A vector can therefore only
pass in `weir-core` if both sides independently agree on every byte.

The bitmap convention is why this file exists at all. An AckBatch written with
the opposite bit order is a WELL-FORMED frame — header CRC valid, payload CRC
valid, payload_len correct — that reports failures as successes. No checksum,
no length check and no cap detects it; only a vector with an asymmetric pattern
and N > 8 does, which is what `ack_bitmap_asymmetric_n9` is.
"""
import json
import zlib

WIRE_VERSION = 1
BATCH_VERSION = 1
ACK_BATCH_VERSION = 1
BATCH_HEADER_LEN = 1 + 2
ACK_BATCH_HEADER_LEN = 1 + 2
MAX_BATCH_RECORDS_HARD_CAP = 2048
MAX_ACK_BATCH_PAYLOAD_LEN = ACK_BATCH_HEADER_LEN + (MAX_BATCH_RECORDS_HARD_CAP + 7) // 8
MAX_TRACKED_ACK_PAYLOAD_LEN = 298
MAX_PAYLOAD_HARD_CAP = 16 * 1024 * 1024
MAGIC = b"WEIR"

MT = {"PushBatch": 0x08, "AckBatch": 0x09}
DUR = {"Durable": 0x01, "Buffered": 0x03}


def crc32(b: bytes) -> int:
    return zlib.crc32(b) & 0xFFFFFFFF


def frame(mt: int, dur: int, payload: bytes) -> bytes:
    h = bytearray(16)
    h[0:4] = MAGIC
    h[4] = WIRE_VERSION
    h[5] = mt
    h[6] = dur
    h[7] = 0
    h[8:12] = len(payload).to_bytes(4, "little")
    h[12:16] = crc32(bytes(h[0:12])).to_bytes(4, "little")
    return bytes(h) + payload + crc32(payload).to_bytes(4, "little")


def body(records, declared=None) -> bytes:
    """A PushBatch body. `declared` overrides the count, to build disagreements."""
    n = len(records) if declared is None else declared
    out = bytearray([BATCH_VERSION])
    out += n.to_bytes(2, "little")
    for r in records:
        out += len(r).to_bytes(4, "little")
        out += r
    return bytes(out)


def ack(n: int, set_indices) -> bytes:
    """An AckBatch payload with `set_indices` accepted, LSB-first within each byte.

    Written as an explicit shift rather than a helper so the convention is
    visible in the generator too: bit i lives in byte i//8 at mask 1 << (i%8).
    """
    bits = bytearray((n + 7) // 8)
    for i in set_indices:
        bits[i // 8] |= 1 << (i % 8)
    return bytes([ACK_BATCH_VERSION]) + n.to_bytes(2, "little") + bytes(bits)


frames, bodies, acks = [], [], []


def ok_frame(name, notes, raw, mt_name, dur_name, payload):
    frames.append({
        "name": name, "notes": notes, "hex": raw.hex(), "decode": "ok",
        "message_type": mt_name, "durability": dur_name, "flags": 0,
        "payload_hex": payload.hex(),
    })


def ok_body(name, notes, raw, records):
    bodies.append({
        "name": name, "notes": notes, "hex": raw.hex(), "decode": "ok",
        "records_hex": [r.hex() for r in records],
    })


def bad_body(name, notes, raw, tag):
    bodies.append({"name": name, "notes": notes, "hex": raw.hex(), "decode": tag})


def ok_ack(name, notes, raw, expected, accepted=None, accepted_all=None):
    """`accepted` lists the per-record verdicts; `accepted_all` is the shorthand.

    The at-cap vector would otherwise carry 2048 booleans, one per line, which
    is 20 KB of noise in every diff of this file and tells a reader nothing the
    hex does not. A runner must accept exactly one of the two forms.
    """
    v = {"name": name, "notes": notes, "hex": raw.hex(), "decode": "ok",
         "expected": expected}
    if accepted is not None:
        v["accepted"] = accepted
    else:
        v["accepted_all"] = accepted_all
    acks.append(v)


def bad_ack(name, notes, raw, expected, tag):
    acks.append({
        "name": name, "notes": notes, "hex": raw.hex(), "decode": tag,
        "expected": expected,
    })


# ── Frames ────────────────────────────────────────────────────────────────────

R3 = [b"alpha", b"beta", b"gamma"]

ok_frame(
    "push_batch_three",
    "PushBatch of three records at Durable durability. The header is a v1 header "
    "in every byte but the message type (0x08) and the CRC that covers it — "
    "which is the whole claim of additivity: a v1 reader's header parse is "
    "unchanged, and only the type dispatch is new.",
    frame(MT["PushBatch"], DUR["Durable"], body(R3)), "PushBatch", "Durable", body(R3),
)

ok_frame(
    "push_batch_buffered",
    "PushBatch at Buffered durability. The tier applies to every record in the "
    "batch; there is no per-record tier, because a batch shares one ack.",
    frame(MT["PushBatch"], DUR["Buffered"], body(R3)), "PushBatch", "Buffered", body(R3),
)

ACK3 = ack(3, [0, 1, 2])
ok_frame(
    "ack_batch_three_all_ok",
    "The daemon's reply to `push_batch_three`: all three accepted. The "
    "durability byte is fixed filler (0x01) exactly as on a plain Ack.",
    frame(MT["AckBatch"], DUR["Durable"], ACK3), "AckBatch", "Durable", ACK3,
)

ACK_MAX = ack(MAX_BATCH_RECORDS_HARD_CAP, range(MAX_BATCH_RECORDS_HARD_CAP))
ok_frame(
    "ack_batch_at_the_hard_cap",
    f"The largest AckBatch that exists: {MAX_BATCH_RECORDS_HARD_CAP} records, "
    f"{MAX_ACK_BATCH_PAYLOAD_LEN} payload bytes. This vector is what pins the "
    f"claim the hard cap was chosen for — {MAX_ACK_BATCH_PAYLOAD_LEN} is under "
    f"the {MAX_TRACKED_ACK_PAYLOAD_LEN}-byte AckTracked payload, so batching "
    "introduces no new largest response and no client gains a new allocation "
    "bound. If the cap is ever raised past that, this vector fails.",
    frame(MT["AckBatch"], DUR["Durable"], ACK_MAX), "AckBatch", "Durable", ACK_MAX,
)

# ── PushBatch bodies ──────────────────────────────────────────────────────────

ok_body(
    "body_three_records",
    "The body carried by `push_batch_three`: version, u16 count, then each "
    "record as a u32 little-endian length followed by its bytes.",
    body(R3), R3,
)

ok_body(
    "body_single_record",
    "Boundary: N=1. Structurally legal and worth no special case — a batch of "
    "one is a batch, and a client that degrades it to a plain Push is making an "
    "optimisation, not obeying the protocol.",
    body([b"solo"]), [b"solo"],
)

bad_body(
    "reject_body_truncated",
    "Shorter than the fixed prefix every body has. Rejected before any field "
    "offset is interpreted.",
    body(R3)[:2], "Truncated",
)

bad_body(
    "reject_body_unknown_version",
    "Leading version byte is 0x02. The version leads so the body can grow inside "
    "wire v1: a reader meeting a layout it does not know must reject rather than "
    "parse a prefix of a layout it has never seen.",
    bytes([0x02]) + body(R3)[1:], "UnsupportedVersion",
)

bad_body(
    "reject_body_empty_batch",
    "Declares zero records. Not a no-op to be tolerated — nothing legitimate "
    "produces one, so it is a framing disagreement and is refused.",
    body([], declared=0), "EmptyBatch",
)

bad_body(
    "reject_body_too_many_records",
    f"Declares {MAX_BATCH_RECORDS_HARD_CAP + 1} records in a body of "
    f"{BATCH_HEADER_LEN} bytes. The rejection must happen on the DECLARED count, "
    "before any allocation is sized by it — that is the whole point of checking "
    "the cap before the walk, and this vector cannot pass any implementation "
    "that reserves first and validates after.",
    body([], declared=MAX_BATCH_RECORDS_HARD_CAP + 1), "TooManyRecords",
)

bad_body(
    "reject_body_empty_record",
    "Record 1 declares length 0. The WAB uses a zero payload_len as its "
    "end-of-records sentinel and cannot represent an empty record, so this is "
    "refused at the edge rather than deeper where the sentinel would be ambiguous.",
    body([b"alpha", b"", b"gamma"]), "EmptyRecord",
)

bad_body(
    "reject_body_truncated_record",
    "Record 0 declares 64 bytes but the body holds 5. The length prefix is "
    "peer-chosen and must be checked against what is actually present.",
    bytes([BATCH_VERSION]) + (1).to_bytes(2, "little")
    + (64).to_bytes(4, "little") + b"alpha",
    "TruncatedRecord",
)

bad_body(
    "reject_body_record_over_the_payload_cap",
    f"Record 0 declares {MAX_PAYLOAD_HARD_CAP + 1} bytes in a body that carries "
    "none of them. The vector is deliberately tiny: it pins the ORDER of the "
    "checks, because the cap must be tested before the body is consulted. An "
    "implementation that checked `cursor + len > body.len()` first would reject "
    "this as TruncatedRecord, which is a different and less informative answer, "
    "and would have computed `cursor + len` on a peer-chosen length first.",
    bytes([BATCH_VERSION]) + (1).to_bytes(2, "little")
    + (MAX_PAYLOAD_HARD_CAP + 1).to_bytes(4, "little"),
    "RecordTooLarge",
)

bad_body(
    "reject_body_more_records_than_declared",
    "Declares 2 records and carries 3. A decoder that stopped at the declared "
    "count would silently drop the third; the extra bytes mean the two ends "
    "disagree about the body, so the frame is refused.",
    body(R3, declared=2), "LengthMismatch",
)

bad_body(
    "reject_body_fewer_records_than_declared",
    "Declares 4 records and carries 3. The walk ends at the body's end having "
    "found fewer than declared.",
    body(R3, declared=4), "LengthMismatch",
)

# ── AckBatch payloads ─────────────────────────────────────────────────────────

ok_ack(
    "ack_all_accepted_n3",
    "The payload carried by `ack_batch_three_all_ok`. Three set bits in one "
    "byte: 0b00000111.",
    ack(3, [0, 1, 2]), 3, [True, True, True],
)

ok_ack(
    "ack_all_rejected_n3",
    "Every bit clear. A clear bit is the WEAK statement: not durable as of this "
    "reply, retry, and expect it may nonetheless have been written. Only a set "
    "bit inherits the crown invariant.",
    ack(3, []), 3, [False, False, False],
)

ok_ack(
    "ack_bitmap_asymmetric_n9",
    "The vector this file exists for. N=9 with records 0, 2 and 8 accepted "
    "encodes as bytes 05 01 under weir's LSB-first convention and a0 80 under "
    "MSB-first. Every vector with N<=8, and every symmetric pattern at any N, "
    "encodes IDENTICALLY under both conventions and would pass against an "
    "inverted implementation — which produces a well-formed frame, valid CRCs "
    "and correct length, that reports failures as successes. N>8 forces the "
    "byte order to show, and the asymmetric pattern forces the bit order to "
    "show. Both are needed; neither alone is sufficient.",
    ack(9, [0, 2, 8]), 9,
    [True, False, True, False, False, False, False, False, True],
)

ok_ack(
    "ack_exactly_one_byte_n8",
    "Boundary: N=8 fills the byte exactly, so there are no padding bits and the "
    "padding rule has nothing to check. The adjacent N=9 case above is where it "
    "starts mattering.",
    ack(8, [1, 3, 5, 7]), 8,
    [False, True, False, True, False, True, False, True],
)

ok_ack(
    "ack_at_the_hard_cap",
    f"The payload of `ack_batch_at_the_hard_cap`: {MAX_BATCH_RECORDS_HARD_CAP} "
    f"records, all accepted, {MAX_ACK_BATCH_PAYLOAD_LEN} bytes.",
    ACK_MAX, MAX_BATCH_RECORDS_HARD_CAP, accepted_all=True,
)

bad_ack(
    "reject_ack_truncated",
    "Shorter than the fixed prefix.",
    ack(3, [0])[:2], 3, "Truncated",
)

bad_ack(
    "reject_ack_unknown_version",
    "Leading version byte is 0x02.",
    bytes([0x02]) + ack(3, [0])[1:], 3, "UnsupportedVersion",
)

bad_ack(
    "reject_ack_empty",
    "Declares zero records.",
    ack(0, []), 0, "EmptyBatch",
)

bad_ack(
    "reject_ack_count_disagrees_with_request",
    "The payload declares 3 but the client sent 4. A bitmap is the first weir "
    "response whose meaning depends on client-held state — an Ack says 'your "
    "last record' and an AckTracked carries its own coordinate, but a bitmap is "
    "meaningless without knowing which batch it answers. Since ceil(N/8) is not "
    "injective (N=1017..1024 all give 128 bitmap bytes), the echoed count is the only "
    "thing that can catch a desync.",
    ack(3, [0, 1, 2]), 4, "LengthMismatch",
)

bad_ack(
    "reject_ack_bitmap_too_short",
    "Declares 9 records but carries one bitmap byte where ceil(9/8)=2 are "
    "required.",
    ack(9, [0])[:-1], 9, "LengthMismatch",
)

bad_ack(
    "reject_ack_padding_bits_set",
    "N=9, so byte 1 holds one used bit and seven padding bits; bit 1 of that "
    "byte is set. Ignoring padding would make popcount(bitmap) == N — the "
    "obvious way to ask 'did the whole batch succeed' — silently wrong.",
    ack(9, [0, 2, 8])[:-1] + bytes([0x03]), 9, "PaddingNotZero",
)

doc = {
    "wire_version": WIRE_VERSION,
    "batch_version": BATCH_VERSION,
    "ack_batch_version": ACK_BATCH_VERSION,
    "batch_header_len": BATCH_HEADER_LEN,
    "ack_batch_header_len": ACK_BATCH_HEADER_LEN,
    "max_batch_records_hard_cap": MAX_BATCH_RECORDS_HARD_CAP,
    "max_ack_batch_payload_len": MAX_ACK_BATCH_PAYLOAD_LEN,
    "bitmap_bit_order": "LSB-first within each byte: bit i is byte i/8, mask 1 << (i%8)",
    "description": (
        "Conformance vectors for the batch extension to the weir v1 wire "
        "protocol: message types PushBatch (0x08) and AckBatch (0x09), the "
        "length-prefixed record body, and the positional bitmap reply. Kept "
        "apart from wire_v1_vectors.json, which is frozen and is read by "
        "decoders that know only message types 0x01..0x05 — such a decoder is "
        "CORRECT to reject an 0x08 frame and stays conformant by ignoring this "
        "file. 'frame_vectors' are complete frames; 'body_vectors' are bare "
        "PushBatch bodies; 'ack_vectors' are bare AckBatch payloads, each with "
        "the record count the client sent, which a decoder must be given "
        "because a bitmap alone cannot say which batch it answers."
    ),
    "frame_vectors": frames,
    "body_vectors": bodies,
    "ack_vectors": acks,
}

print(json.dumps(doc, indent=2))