"""
Walker v5: flexible NN range and multiple block-type lengths.

Hypothesis:
- [10 NN]: 4-bit-delta data block, length = NN/2 + 2
- [20 NN]: 8-bit-literal data block, length = NN + 2
- [00 NN]: 2-byte marker (no payload)
- [30 NN]: trailer/summary block, length = NN*4
- [40 NN]: footer-marker block, fixed 22 bytes
"""
import sys
sys.path.insert(0, ".")
from analysis.load_bundle import load_bundle
from collections import Counter


def walk(body, start, max_blocks=10000):
    i = start
    blocks = []
    while i + 1 < len(body) and len(blocks) < max_blocks:
        t0 = body[i]
        t1 = body[i + 1]
        if t0 == 0x10 and t1 % 4 == 0 and 0 < t1 <= 0xFC:
            length = t1 // 2 + 2
            if i + length > len(body):
                break
            data = bytes(body[i + 2 : i + length])
            blocks.append((i, "10", t1, data, length))
            i += length
        elif t0 == 0x20 and t1 % 4 == 0 and 0 < t1 <= 0xFC:
            length = t1 + 2
            if i + length > len(body):
                break
            data = bytes(body[i + 2 : i + length])
            blocks.append((i, "20", t1, data, length))
            i += length
        elif t0 == 0x00 and t1 % 4 == 0:
            # 2-byte marker
            blocks.append((i, "00", t1, b"", 2))
            i += 2
        elif t0 == 0x30 and t1 % 4 == 0:
            length = t1 * 4
            if i + length > len(body):
                break
            data = bytes(body[i + 2 : i + length])
            blocks.append((i, "30", t1, data, length))
            i += length
        elif t0 == 0x40 and t1 == 0x02:
            length = 22
            if i + length > len(body):
                break
            data = bytes(body[i + 2 : i + length])
            blocks.append((i, "40", t1, data, length))
            i += length
        else:
            blocks.append((i, "??", t0, bytes(body[i:i+8]), 0))
            break
    return blocks, i


def main():
    for name in ("event-c", "event-d", "event-a", "event-b"):
        b = load_bundle(name)
        body = b.body
        for s in range(15):
            if body[s] == 0x10 and body[s+1] % 4 == 0 and 0 < body[s+1] <= 0xFC:
                start = s; break
        else:
            start = 7
        blocks, end = walk(body, start)
        types = Counter(bb[1] for bb in blocks)
        print(f"\n=== {name} === body={len(body)} N={len(b.samples['Tran'])}  start={start}")
        print(f"  total blocks: {len(blocks)}, walk ended at {end}/{len(body)}")
        print(f"  type counts: {dict(types)}")
        if blocks and blocks[-1][1] == "??":
            print(f"  stopped at byte: 0x{blocks[-1][2]:02x}, prev 5 blocks: {[(bb[0], bb[1], bb[2]) for bb in blocks[-6:-1]]}")
        # Sum payload sizes by type
        payload_sizes = {t: sum(len(bb[3]) for bb in blocks if bb[1] == t) for t in types}
        print(f"  payload bytes by type: {payload_sizes}")


if __name__ == "__main__":
    main()