#!/usr/bin/env python3

import zmq

def main() -> None:
    ctx = zmq.Context()

    # Subscribe to the three point-cloud publishers
    data_ports = [5556, 5557, 5558]
    subs = []
    for port in data_ports:
        sock = ctx.socket(zmq.SUB)
        sock.connect(f"tcp://localhost:{port}")
        sock.setsockopt(zmq.SUBSCRIBE, b"")
        subs.append(sock)

    # Publisher for the merged cloud
    pub = ctx.socket(zmq.PUB)
    pub.bind("tcp://*:5563")

    # Poller to wait for incoming clouds
    poller = zmq.Poller()
    for sock in subs:
        poller.register(sock, zmq.POLLIN)

    pc_bufs = [None] * 3  # buffers for the three incoming clouds

    while True:
        # Collect a full set of three point clouds
        ready = dict(poller.poll())
        for idx, sock in enumerate(subs):
            if sock in ready:
                pc_bufs[idx] = sock.recv()

        # Once all three parts arrive, concatenate and publish
        if all(pc_bufs):
            combined = b"".join(pc_bufs)
            pub.send(combined)
            print(f"Published combined point cloud ({len(combined)} bytes)")
            pc_bufs = [None] * 3  # reset buffers


if __name__ == "__main__":
    main()

