#!/usr/bin/env python3
"""Generate JDL plan-buffer words for simple dynamic lookup exchanges.

This intentionally covers the narrow case used by example_codegen.cpp:
one sender tile and one receiver tile on a single IPU, with the sender selected
at runtime by the receiver tile's tileSelector scalar.
"""

from __future__ import annotations

import argparse
import json


def sender_row(receiver_tile: int, count: int) -> list[int]:
    # The middle word is a delay until the receiver-side mux point.
    # Known-good for the tested C600/ipu21 routes below.
    delay_by_receiver = {
        16: 0x9C,
        64: 0xA2,
        128: 0xAF,
        256: 0x98,
        512: 0xA9,
        768: 0x93,
        1024: 0x88,
        1025: 0x89,
        1280: 0x8C,
        1286: 0x86,
        1471: 0x8B,
    }
    if receiver_tile not in delay_by_receiver:
        raise ValueError(
            f"receiver tile {receiver_tile} is not in the small recovered table"
        )
    return [
        0x41800003,
        0x40A00000 | delay_by_receiver[receiver_tile],
        0x78000001 | ((count - 1) << 21),
        0x43A00000,
        0,
        0,
        0,
        0,
        0,
    ]


def receiver_row(receiver_tile: int, count: int) -> list[int]:
    # The setup code records index 1 as the tileSelector-patched instruction.
    route_base_by_receiver = {
        16: 0x62180000,
        64: 0x62680000,
        128: 0x62B80000,
        256: 0x62280000,
        512: 0x62C80000,
        768: 0x62980000,
        1024: 0x61F80000,
        1025: 0x61F80000,
        1280: 0x61C80000,
        1286: 0x61980000,
        1471: 0x61980000,
    }
    if receiver_tile not in route_base_by_receiver:
        raise ValueError(
            f"receiver tile {receiver_tile} is not in the small recovered table"
        )
    return [
        1,
        0x41800003,
        0x641C0000,
        0x64000640 | ((count - 1) << 14),
        route_base_by_receiver[receiver_tile] - (count << 19),
        0x40A00000 | (count + 4),
        0x43A00000,
        0,
        0,
    ]


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--receiver", type=int, default=1286)
    parser.add_argument("--count", type=int, default=3)
    args = parser.parse_args()
    if args.count < 1 or args.count > 32:
        raise SystemExit("count must be in [1, 32] for this simple generator")

    rows = [sender_row(args.receiver, args.count), receiver_row(args.receiver, args.count)]
    hex_rows = [[f"0x{word:08x}" for word in row] for row in rows]
    print(json.dumps({"receiver": args.receiver, "count": args.count, "rows": hex_rows}, indent=2))


if __name__ == "__main__":
    main()