require unix/socket.fs

require socket-extensions.4th
require connections.4th
require sendbuffer.4th

AF_INET constant SERVER_SOCKET_DOMAIN
  SOCK_STREAM
  SOCK_NONBLOCK or
        constant SERVER_SOCKET_TYPE
      0 constant SERVER_SOCKET_PROTOCOL
      0 constant SERVER_ADDR
  50000 constant SERVER_PORT
    128 constant SERVER_LISTEN_BACKLOG

\ Listening file descriptor.
variable listenfd
0 listenfd !

\ Idle detection.
variable idle false idle !
: server-idle? ( -- flag ) idle @ ;
: server-idle! ( flag -- ) idle ! ;

\ Temporary structs.
create con /CONNECTION allot
create saddr /sockaddr_in allot
create optval /option_value allot

: (assert-setsockopt) ( result-n -- )
  0< abort" making socket reusable failed" ;
: (optval-true!) ( -- )
  1 optval l! ;
: (make-reusable) ( socket-fd-n -- )
  (optval-true!)
  SOL_SOCKET SO_REUSEADDR optval /option_value setsockopt (assert-setsockopt) ;

: (saddr!) ( protocol-n sin_addr-n port-n -- )
  htons saddr port w!
  saddr sin_addr l!
  saddr family w! ;

: (assert-socket) ( result-n -- result-n )
  dup 0< abort" socket() failed." ;
: (assert-bind) ( result-n -- )
  0< abort" bind failed." ;
: (assert-listen) ( result-n -- )
  0< abort" listen failed." ;

: (erase-saddr) ( -- )
  saddr /sockaddr_in erase ;
: (create-socket)  ( -- )
  SERVER_SOCKET_DOMAIN
  SERVER_SOCKET_TYPE
  SERVER_SOCKET_PROTOCOL
    socket (assert-socket) listenfd ! ;
: (make-socket-reusable) ( -- )
  listenfd @ (make-reusable) ;
: (set-saddr) ( -- )
  SERVER_SOCKET_DOMAIN SERVER_ADDR SERVER_PORT (saddr!) ;
: (bind-socket) ( -- )
  listenfd @ saddr /sockaddr_in bind (assert-bind) ;
: (listen-socket) ( -- )
  listenfd @ SERVER_LISTEN_BACKLOG listen() (assert-listen) ;
: (server-info) ( -- )
  cr cr ." Server listening at port: " SERVER_PORT . cr ;
: initialize-server ( -- )
  (erase-saddr)
  (create-socket)
  (make-socket-reusable)
  (set-saddr)
  (bind-socket)
  (listen-socket)
  (server-info) ;

: (perform-disconnect) ( connection-addr -- )
  dup connection.connected false swap !
      connection.fd @ close() throw ;

: (close-clients) ( -- )
  connections.count 0= IF
    EXIT
  THEN
  connections.count 0 DO
    I connections.at connection.connected @ true = IF
      I connections.at (perform-disconnect)
    THEN
  LOOP ;

: (assert-close()) ( result-n -- )
  0<> abort" close failed" ;
: (close-server) ( -- )
  listenfd @ close() (assert-close()) ;

: (close-server-info) ( -- )
  cr ." Closed server connections." cr ;
: close-server ( -- )
  (close-clients) (close-server) (close-server-info) ;

: (queue-disconnect) ( connection-addr -- )
  dup (perform-disconnect) EVENT_CONNECTION_CLOSED events.enqueue ;

: (try-accept) ( -- c-result-n )
  listenfd @ 0 0 accept() ;
: (accept-error) ( accept-result-n -- )
  errno EAGAIN <> abort" accept error" drop ;
: (erase-connection) ( connection-addr -- )
  /CONNECTION erase ;
: (set-connection-number) ( -- )
  connections.last dup connections.indexOf 1+ swap connection.number ! ;
: (enqueue-new-connection) ( -- )
  connections.last EVENT_CONNECTION_NEW events.enqueue ;
: (store-connection) ( connection-addr -- )
  dup >r connections.append IF
    (set-connection-number)
    (enqueue-new-connection)
    rdrop
  ELSE
    ." Warning: failed to store connection, disconnecting client." cr
    r> (perform-disconnect)
  THEN ;
: (con!) ( fd-n connected-flag -- )
  con connection.connected !
  con connection.fd ! ;
: (accept-connection) ( fd-n -- )
  con (erase-connection)
  true (con!)
  con (store-connection) ;
: (server-idle-accept) ( -- )
  (try-accept) dup 0< IF
    (accept-error)
  ELSE
    (accept-connection)
  THEN ;

: (connected?) ( connection-addr -- flag )
  connection.connected @ ;
: (try-recv) ( connection-addr -- recv-result-n )
  dup connection.fd @
  swap connection.buffer
  CONNECTION_BUFFER_SIZE
  MSG_DONTWAIT
    recv ;
: (recv-error?) ( c-result-n -- flag ) 0< ;
: (recv) ( recv-result-n connection-addr -- )
  2dup connection.bufferlen ! swap
  0> IF
    EVENT_CONNECTION_RECV events.enqueue
  ELSE \ disconnected
    (queue-disconnect)
  THEN
;
: (recv-warning) ( -- ) ." Warning: recv failed, disconnecting client." cr ;
: (recv-error) ( recv-result-n connection-addr -- )
  errno EAGAIN <> IF
    (queue-disconnect) drop
    (recv-warning)
  ELSE
    2drop
  THEN ;
: (connection-recv) ( connection-addr -- )
  dup (try-recv) tuck (recv-error?) IF
    (recv-error)
  ELSE
    (recv)
  THEN ;

: (server-idle-recv) ( -- )
  true server-idle!
  connections.count 0= IF
    EXIT
  THEN
  connections.count 0 DO
    I connections.at dup (connected?) IF
      (connection-recv)
    ELSE
      drop
    THEN
  LOOP ;

: server-idle-accept ( eventdata-n eventid-n -- )
  2drop (server-idle-accept) ;
: server-idle-recv ( eventdata-n eventid-n -- )
  2drop (server-idle-recv) ;
: server-connection-new ( connection-addr eventid-n -- )
  2drop ." New client connected!" cr ;
: server-connection-closed ( connection-addr eventid-n -- )
  2drop ." Client disconnected." cr ;

: (to-string) ( n -- addr c )  s>d <# #s #> ;
: (connectionnumber@) ( connection-addr -- c-addr u )
  connection.number @ (to-string) ;
: (connectionbuffer@) ( connection-addr -- c-addr u )
  dup connection.buffer swap connection.bufferlen @ ;
: (format-sendbuffer) ( from-connection-addr -- )
  >r sendbuffer-reset
     s" Anon "              sendbuffer-append
     r@ (connectionnumber@) sendbuffer-append
     s" : "                 sendbuffer-append
     r> (connectionbuffer@) sendbuffer-append
  sendbuffer-sanitize ;
: (connected?) ( connection-addr -- )
  connection.connected @ ;
: (different-connection?) ( from-connection-addr to-connection-addr -- )
  <> ;
: (send?) ( from-connection-addr to-connection-addr -- )
  tuck (different-connection?) swap (connected?) and ;
: (check-send) ( result-n -- )
  0< IF ." Warning: send failed." cr THEN ;
: (send) ( from-connection-addr to-connection-addr -- )
  nip connection.fd @ sendbuffer@ 0 send (check-send) ;
: (try-send) ( from-connection-addr to-connection-addr -- )
  2dup (send?) IF
    (send)
  ELSE
    2drop
  THEN ;
: server-recv ( from-connection-addr eventid-n )
  drop dup (format-sendbuffer)
  connections.count 0 DO
    dup I connections.at (try-send)
  LOOP drop ;

' server-idle-accept       EVENT_IDLE              eventhandlers.append
' server-idle-recv         EVENT_IDLE              eventhandlers.append
' server-connection-new    EVENT_CONNECTION_NEW    eventhandlers.append
' server-connection-closed EVENT_CONNECTION_CLOSED eventhandlers.append
' server-recv              EVENT_CONNECTION_RECV   eventhandlers.append

initialize-server