#!/usr/bin/env pf

# this script serves as a demonstration of PF's TCP capabilities.  it
# sets up a maximally connected (all to all) network of nodes.  the
# connections themselves are unidirectional, giving a total of N(N-1)
# TCP pipes for N nodes.

# the script takes the following startup arguments
# - a number giving the total amount of interconnected nodes
# - a number giving this node's number
# - base port number

# it works like this:
# - script opens a listening socket from which it accepts other connections
# - it tries to connect to all other peers
# - if it receives data from a peer, it's put in the big bag
# - it randomly sends out a zero to a peer
# - it randomly sends out a number from the big bag, incremented by one



# things that are assumed to simplify the example
# - all clients are on localhost, and identified by a single number.
#   the latter fact is only used to simply map client -> port number
# - a connection is made only once. if it breaks, no retry is attempted.



# to try this script, make sure it's executable (chmod +x) and run in
# 3 terminals:

# ./maxcon.pf 3 0 50000
# ./maxcon.pf 3 1 50000
# ./maxcon.pf 3 2 50000


# UTILITY


# quickly interpret command line args.
: next-arg  args pop interpret-string ;

# don't remember if this is a core word somewhere.. if not, it should be.
# get a list from listvar and replace listvar with empty list.
: slurp   dup @> >r () swap ! r> ;



# MAIN SCRIPT

variable me
variable population
variable base-port


: port base-port @ + ;	# ( ID - TCP.port )

# will fail if there are not enough args, and ignores the extra args.
: parse-command-line-arguments
	args pop drop	# ignore script name
	next-arg population !
	next-arg me !
	next-arg base-port ! ;


# connection
() variable! peers-out   # connected output streams
() variable! peers-in	 # accepted input streams
() variable! ids         # wanted client IDs


# CLIENT (connect to other peer)

# connect to a peer
: open-peer  port 1 pack (localhost) swap + "w" open-connect ;


# create initial ID list = all peers except me
: generate-ids
	me @			# start with me, so i'll be excluded (pre-inc)
	population @ 1 - for	# number of connections
		1 + dup		# pre-inc
		population @ mod
		ids push
	next drop ;
	

# try to connect ( ID -- )
# fail -> push ID to ids
# ok   -> push stream to connections
: try-peer
	try
		dup	# need dup here so recover can recover ID 
		open-peer
		peers-out push
		"ID" p p "connected" p cr
		
	recover
		drop 
		"ID" p dup p "failed" p cr
		ids push
	endtry ;

# retry all in ID list
: try-connect  ids slurp ' try-peer for-each ;


# SERVER (accept from other peers)

# the timeout used here is 0. sometimes it can be simpler to
# incorporate the main loop timing in 'select-input', however, since
# i'm polling 2 things in the main loop, each poll (select) just
# checks the operating system queues.

variable server	       # server listening socket stream

: open-server
    me @ port "localhost"     # listen only on local host socket
    2 pack "r" open-listen    # create listening socket
    server ! ;


# there are 2 parts:
# - accept new connections
# - handle input on accepted connections

# small word to make select-input more usable in our case, since
# there's only a single stream in the list of streams (the server
# socket). this word will generate a list which contains the server
# stream if there is a new connection, or contains nothing at
# all. this allows to use 'for-each' on the return value.

: select-server 
	server @ 1 pack 0	# arguments: stream.list timeout
	select-input drop ;	# discard the idle list

: accept-new
	select-server
	      { "r" open-accept	   # accept new connection
	      	dup peers-in push  # and store it to a list
		"new connection:" p p cr }
	for-each ;


# we'll accept things, and put them in the big bag here
() variable! big-bag

# get a list of all peers that have incoming data ready
: peers-in-ready  peers-in @ 0 select-input drop ;

: handle-input 
	  peers-in-ready
	  { read-atom 
            "got:" p dup p cr 
            big-bag push } for-each ;


# ( stream thing -- )
: write-peer
	"sending:" p dup p cr 
	swap write-atom ;
  

# maybe send out a zero ( stream -- )
: seed-zero 0 write-peer ;


# try to send out what's in the big-bag  ( stream -- )
: try-big-bag
	try big-bag pop 1 + write-peer
	recover 2drop endtry ;  # ignore errors


# we'll hand out things randomly sometime
: seed? 10 rand not ;
: bag?  3 rand not ;

# perform output
: handle-output
	peers-out @ 
          { seed? if dup seed-zero then
	    bag?  if dup try-big-bag then
	    drop } for-each ;


# status
: print-status
  "bag:" p big-bag @ p cr ;	



# mainloop
# you can run this from tick or so..
: mainloop 
     try-connect	# (re) try to connect to peers
     accept-new		# accept peer inbound connections
     handle-input	# handle input data from peers
     handle-output	# handle output to peers
     print-status
     
     100. sleep ;

# startup
: go
	parse-command-line-arguments
	generate-ids
	open-server

	# debug message
	"starting" p me @ p 
	"of" p population @ p 
	"port" p me @ port p cr

	' mainloop is tick ;

go

# debug
# "going interactive" p cr
interactive