"""Corrupt one zstd block of an out-of-line TOAST value, on disk.

    corrupt_block.py CHUNKS RELFILE [BLOCK]

CHUNKS is "chunk_seq|hex of chunk_data" per line, for one value, as dumped
while the server ran; RELFILE the TOAST relation file, with the server
stopped.  Walks the zstd frame across the chunks and sets the type of the
header of block BLOCK (default 2, the third) to the reserved value 3, so any
decoder fails exactly there: deterministic, and in the middle of the frame,
after BLOCK good blocks.
"""
import sys

chunks_path, relfile = sys.argv[1], sys.argv[2]
target = int(sys.argv[3]) if len(sys.argv) > 3 else 2

chunks = []
for line in open(chunks_path):
    seq, hx = line.strip().split("|")
    chunks.append((int(seq), bytes.fromhex(hx[2:] if hx.startswith("\\x") else hx)))
chunks.sort()
data = b"".join(c for _, c in chunks)

m = data.find(b"\x28\xb5\x2f\xfd")
assert m >= 0, "no zstd magic"
fhd = data[m + 4]
single = (fhd >> 5) & 1
pos = m + 5 + (0 if single else 1) + [0, 1, 2, 4][fhd & 3] + [single, 2, 4, 8][fhd >> 6]
blocks = []
while True:
    bh = data[pos] | data[pos + 1] << 8 | data[pos + 2] << 16
    last, typ, size = bh & 1, (bh >> 1) & 3, bh >> 3
    blocks.append((pos, typ, size))
    pos += 3 + (1 if typ == 1 else size)
    if last:
        break
print(f"frame at {m}, {len(blocks)} blocks, types {[b[1] for b in blocks]}")
assert len(blocks) > target, "value too small for the target block"

# Find the chunk holding that header byte, then the chunk in the file.
off = blocks[target][0]
for seq, c in chunks:
    if off < len(c):
        break
    off -= len(c)
raw = bytearray(open(relfile, "rb").read())
at = raw.find(c)
assert at >= 0 and raw.find(c, at + 1) < 0, "chunk not found once in the file"
before = raw[at + off]
raw[at + off] |= 0x06
open(relfile, "wb").write(raw)
print(f"block {target} header: chunk_seq {seq} offset {off}, file offset {at + off}, "
      f"byte 0x{before:02x} -> 0x{raw[at + off]:02x}")