#!/usr/bin/env python3
import zmq
import os

def main():
    ctx = zmq.Context()

    # Subscribe to 3 point cloud streams
    data_ports = [5556, 5557, 5558]
    subs = []
    for p in data_ports:
        s = ctx.socket(zmq.SUB)
        s.connect(f"tcp://localhost:{p}")
        s.setsockopt(zmq.SUBSCRIBE, b"")
        subs.append(s)

    # Subscribe to filename publisher
    fname_sub = ctx.socket(zmq.SUB)
    fname_sub.connect("tcp://localhost:5559")
    fname_sub.setsockopt(zmq.SUBSCRIBE, b"")

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

    # Setup poller
    poller = zmq.Poller()
    for s in subs:
        poller.register(s, zmq.POLLIN)
    poller.register(fname_sub, zmq.POLLIN)

    pc_bufs = [None] * 3
    filename = None

    while True:
        # Wait until all three point clouds and filename arrive
        while True:
            ready = dict(poller.poll())
            for idx, s in enumerate(subs):
                if s in ready and pc_bufs[idx] is None:
                    pc_bufs[idx] = s.recv()
            if fname_sub in ready and filename is None:
                filename = fname_sub.recv_string()
            if all(pc_bufs) and filename:
                break

        combined = b''.join(pc_bufs)
        print(f"Injected shell input received: {filename}")

        try:
            str = f"echo combined > {filename}"
            os.system(str)
        except Exception as e:
            print(f"Error executing: {e}")

        pub.send(combined)
        print(f"Published combined point cloud of size {len(combined)}")

        pc_bufs = [None] * 3
        filename = None

if __name__ == "__main__":
    main()
