(use bokbok)
(use matchable)

(define *server* #f)

(define (open-handler con)
  (printf "OPEN from ~a @ ~s\n"
          (connection-user con)
          (connection-addr con)))


(define (close-handler con)
  (printf "CLOSE from ~a @ ~s\n"
          (connection-user con)
          (connection-addr con)))


(define (request-handler con request)
  (printf "REQUEST from ~a @ ~s: ~s\n"
          (connection-user con)
          (connection-addr con)

          request)

  (match request
         (("error")
          (error "Error raised"))
         (("callback")
          (let ((response
                 (request! con '("called you back"))))
            (list "did it" (sprintf "~s" response))))
         (("kill")
          (if *server*
              (begin
                (stop-server! *server*)
                (list "ok"))
              (error "I'm not a server")))
         (else
          (cons "echo" request))))

(define (expect exp actual)
  (unless
   (equal? exp actual)
   (printf "Expected ~s but got ~s\n" exp actual)))

(define (run-client! con)
  (expect '(ok ("echo" "hello"))
          (request! con '("hello")))
  (expect '(error ("\"Error raised\" in (#f)"))
          (request! con '("error")))
  (expect '(ok ("did it"  "(ok (\"echo\" \"called you back\"))"))
          (request! con '("callback")))



#;  (expect '(ok ("ok"))
          (request! con '("kill")))
  (close-connection! con))

(define (run-server! server)
  (set! *server* server)
  (printf "Running server...\n")
  (wait-until-server-stopped server))

(match
 (command-line-arguments)
 (("client" "unix" path)
  (run-client! (open-connection `(unix ,path) #f #f request-handler close-handler)))
 (("client" "unix" path user pass)
  (run-client! (open-connection `(unix ,path) user (passphrase->key pass) request-handler close-handler)))

 (("server" "unix" path)
  (run-server! (start-server `(unix ,path) 10 #f open-handler request-handler close-handler)))
 (("server" "unix" path pass)
  (run-server! (start-server `(unix ,path) 10 (lambda _ (passphrase->key pass)) open-handler request-handler close-handler)))

 (("client" "tcp" addr port)
  (run-client! (open-connection `(tcp ,addr ,(string->number port)) #f #f request-handler close-handler)))
 (("client" "tcp" addr port user pass)
  (run-client! (open-connection `(tcp ,addr ,(string->number port)) user (passphrase->key pass) request-handler close-handler)))

 (("server" "tcp" port)
  (run-server! (start-server `(tcp #f ,(string->number port)) 10 #f open-handler request-handler close-handler)))
 (("server" "tcp" port pass)
  (run-server! (start-server `(tcp #f ,(string->number port)) 10 (lambda _ (passphrase->key pass)) open-handler request-handler close-handler)))

 (else
  (printf "Imagine a comprehensive usage screen here!\n")))

(module
 bokbok
 (passphrase->key

  open-connection
  close-connection!
  request!

  connection?
  connection-user
  connection-addr


  start-server
  stop-server!
  wait-until-server-stopped
  )

 (import chicken scheme)

 ;; FIXME: Reduce timeouts to detect gone servers quickly

 ;; FIXME: Auto reconnect on request!

 ;; FIXME: Configurable maximum message size (and use 4 bytes, so it can be >64KiB)

#; (define (debug . args)
  (let ((str (apply sprintf args)))
    (printf "DEBUG: ~a\n" str)))

 (define (debug . args)
   (void))

 (define *header* "BOKBOKBOK")

 (define (read-packet! i)
   (let ((l1 (read-byte i)))
     (if (eof-object? l1)

          (nonce-u8vector (blob->u8vector/shared (string->blob nonce)))
          (plaintext ((cdr key) c-wo-n nonce-u8vector)))
     (if plaintext
         plaintext
         (error "Invalid cyphertext"))))

 (define-record-type connection
   (make-connection* addr user
                    socket input output
                    key
                    mutex
                    thread
                    waiters
                    request-handler
                    close-handler
                    counter)
   connection?
   (addr connection-addr)

   (user connection-user)
   (socket connection-socket)
   (input connection-input)
   (output connection-output)
   (key connection-key)
   (mutex connection-mutex)
   (thread connection-thread)

                     (append
                      (list "ok" id)
                      ((connection-request-handler con)
                       con
                       body)))))
               (debug "Sending response id:~a ~s" id response)
               (connection-send! con response)))
           `(bokbok-request-thread ,(connection-addr con) ,id))))
     (thread-start! thread)))

 (define (handle-response! con id body)
   (debug "Handling response id:~a ~s" id body)
   (mutex-lock! (connection-mutex con))
   (let ((waiter (hash-table-ref/default (connection-waiters con)
                                         id #f)))

               (("ok" id . body)
                (handle-response! con id (list 'ok body)))
               (("err" id . body)
                (handle-response! con id (list 'error body))))
              ;; Loop for next request
              (loop)))))))

 (define (make-connection addr user socket input output key request-handler close-handler)
   (let ((con
          (make-connection*
           addr user
           socket input output
           key
           (make-mutex `(bokbok-connection-mutex ,addr))
           (make-thread handle-connection-thread!
                        `(bokbok-connection-thread ,addr))
           (make-hash-table)
           request-handler
           close-handler
           0)))

     (thread-specific-set!
      (connection-thread con)
      con)
     (thread-start!
      (connection-thread con))
     con))

 (define (open-connection addr username key request-handler close-handler)
   (let ((s
          (match
           addr
           (('unix path)
            ;; UNIX domain
            (let ((s (socket af/unix sock/stream)))
              (socket-connect s (unix-address path))
              s))
           (('tcp host port)
            ;; TCP domain
            (let ((s (socket af/inet sock/stream))
                  (ai (address-information host port)))
              (socket-connect s (addrinfo-address (car ai))#;(inet-address host port))
              (set! (tcp-no-delay? s) #t)
              s)))))
     (receive
      (input output)
      (parameterize
       ((socket-send-buffer-size 4096)
        (socket-send-size 16384))
       (socket-i/o-ports s))
      (if (and username key)
          (begin ;; Encrypted connection
            (receive
             (session-key session-key-bytes)
             (generate-session-key)
             (debug "Client writing username and session key")
             (write-packet-no-flush! output username)
             (write-packet! output (encrypt key session-key-bytes))
             (debug "Client waiting for header")
             (let ((header (read-string (string-length *header*) input)))
               (debug "Got header bytes: ~s" header)
               (if (string=? header *header*)
                   (make-connection addr username s input output session-key request-handler close-handler)
                   (error "Invalid hello from server" header)))))
          (begin ;; Plaintext connection
            (debug "Client waiting for header")
            (let ((header (read-string (string-length *header*) input)))
              (debug "Got header bytes: ~s" header)
              (if (string=? header *header*)
                  (make-connection addr #f s input output #f request-handler close-handler)
                  (error "Invalid hello from server" header))))))))

 (define (close-connection! con)
   ;; FIXME: Set flags to suppress any further responses from still-running handlers
   ;; FIXME: Return an error to all pending waiters
   (thread-terminate! (connection-thread con))
   (close-output-port (connection-output con))
   (close-input-port (connection-input con)))

 (define (request! con packet-parts)
   (mutex-lock! (connection-mutex con))
   (let* ((id (number->string (connection-counter con)))
          (waiter (cons (make-mutex `(bokbok-request-mutex ,(connection-addr con) id))
                        #f)))

     ;; Mutex starts life locked by the connection thread
     (mutex-lock! (car waiter)
                  #f
                  (connection-thread con))

     (hash-table-set! (connection-waiters con) id waiter)
     (set! (connection-counter con)
           (+ (connection-counter con) 1))
     (mutex-unlock! (connection-mutex con))

     (connection-send! con (cons "req" (cons id packet-parts)))

     ;; Wait for response, when connection thread unlocks the mutex
     (mutex-lock! (car waiter))

     ;; Return response
     (cdr waiter)))

 (define (start-server bind-addr backlog user->key open-handler request-handler close-handler)
   (let* ((s (match
              bind-addr
              (('unix path)
                 ;; UNIX domain
               (let ((s (socket af/unix sock/stream)))
                 (socket-bind s (unix-address path))
                 (socket-listen s backlog)
                 s))
              (('tcp addr port)
               ;; TCP domain
               (let ((s (socket af/inet sock/stream)))
                 (socket-bind s (inet-address addr port))
                 (set! (so-reuse-address? s) #t)
                 (socket-listen s backlog)
                 s))
              (else (error "Unknown bind address ~s" bind-addr))))
          (tcp? (match bind-addr (('tcp . any) #t) (else #f)))
          (thread
           (make-thread
            (lambda ()
              (let loop ()
                (debug "Listener thread calling accept")
                (let* ((cs
                        (socket-accept s))
                       (handler-thread
                        (make-thread
                         (lambda ()
                           ;; Connection setup handler thread
                           ;; Hands over to connection thread when make-connection is called
                           (when tcp?
                                 (set! (tcp-no-delay? cs) #t))
                           (receive
                            (input output)
                            (parameterize
                             ((socket-send-buffer-size 4096)
                              (socket-send-size 16384))
                             (socket-i/o-ports cs))
                            (let* ((peer (socket-peer-name cs))
                                   (peer-addr (if tcp?
                                                  (list 'tcp (sockaddr-address peer) (sockaddr-port peer))


                                                  '(unix (sockaddr-path peer)))))

                              (debug "Handshake thread started for ~s" peer-addr)
                              (if user->key
                                  ;; Encrypted connection
                                  (let* ((user-name (read-packet! input))
                                         (encrypted-session-key (read-packet! input)))
                                    (if (and (not (eof-object? user-name))
                                             (not (eof-object? encrypted-session-key)))
                                        (let ((user-key (user->key user-name)))
                                          (if user-key
                                              (let* ((session-key-bytes (decrypt user-key encrypted-session-key))
                                                     (session-key (bytes->key session-key-bytes)))
                                                (if session-key
                                                    (begin
                                                      (debug "Writing header")
                                                      (write-string *header* #f output)
                                                      (flush-output output)
                                                      (open-handler
                                                       (make-connection
                                                        peer-addr user-name
                                                        cs input output
                                                        session-key
                                                        request-handler close-handler)))
                                                    (begin
                                                      (debug "Rejecting connection due to invalid session key")
                                                      (close-output-port output)
                                                      (close-input-port input))))

                                  ;; Plaintext connection
                                  (begin
                                    (debug "Writing header")
                                    (write-string *header* #f output)
                                    (flush-output output)
                                    (open-handler
                                     (make-connection
                                      peer-addr #f
                                      cs input output
                                      #f
                                      request-handler close-handler))))))
                           cs)
                         `(bokbok-handshake-handler ,bind-addr))))
                  (debug "Listener thread starting handshake thread")
                  (thread-start! handler-thread)
                  (loop))))
            `(bokbok-listen-handler ,bind-addr))))
     (thread-specific-set! thread s)
     (thread-start! thread)
     thread))

 (define (stop-server! server)
   (thread-terminate! server))