#lang racket
(provide new-session close-session session?
         new-message new-response message?
         message-args message-referer message-id message-flags
         read-message write-message
         protocol-version
         typed typed? untype lump-argument-type? lump-internal-type? lump-external-type?
         ;type:int8 ; M. Douglas Williams: signed byte integers not supported
         type:bool type:uint8 type:int16 type:uint16 type:int32 type:uint32
         type:int64 type:uint64 type:text type:symbol type:list type:number type:bytes
         type:vector)
(require racket/port)
(require (planet williams/packed-binary/packed-binary))

; This module defines the "lightweight unnecessary message protocol (LUMP)"
; and routines for parsing and streaming it

; LUMP Schematics
;
; Name      Size (Bytes)   Explanation
;--------------------------------------------------------------------
; HEADER
;--------------------------------------------------------------------
; version   1/2            protocol version (nibble)
; flags     1/2            4 flags (nibble)
; id        2              command id
; seqnum    4              per session sequence number
;--------------------------------------------------------------------
; If response flag is set:
; referer   4              sequence number of message this message responds to
;--------------------------------------------------------------------
; If extended flag is set:
;--------------------------------------------------------------------
; argnum    2              number of arguments to follow
;------------------------------------
; ARGUMENT (repeated argnum times)
;------------------------------------
; type      1              type of the argument
; arglen    4              length of argument
; 1...arglen argument data
;--------------------------------------------------------------------


; start a new session
; returns the session
(define (new-session)
  (session 0))

; close a session
; currently does nothing but should be called when a session is ended
(define (close-session session)
  (void))

; internal message structure
(struct message (id seqnum flags version args referer))

; holds the state of a session
; some procedures mutate this to increase the counter
; for obtaining a monotonically increasing sequence number
(struct session (counter) #:mutable)

; create a new message with given id, referer flags and optional arguments
(define new-message 
  (lambda (id #:referer [referer #f] #:flags [flags 0] #:protocol-version [version protocol-version] . args)
    (message id 0 flags version args referer)))

; create a new response to an existing message
; automatically fills in the sequence number of the referer
(define (new-response id referer-message #:flags [flags 0] #:protocol-version [version protocol-version] . args)
  (message id 0 flags version args (message-seqnum referer-message)))

; used to encapsulate values with type information
; only needed for fixed-size numerical types
; e.g. (type type:int16 200) will ensure that 200 is written as unsigned 2 byte integer
(struct typed (type value))

; remove the type information from a value
; returns the original value
(define (untype v)
  (if (typed? v)
      (typed-value v)
      v))

; current protocol version
(define protocol-version 1) 

; these data types are exported, to be used with typed above when needed
; a number is written as string, so some of the fixed-size numeric
; data types might be useful for performance reasons
(define type:bool 0)
(define type:int8 1)
(define type:uint8 2)
(define type:int16 3)
(define type:uint16 4)
(define type:int32 5)
(define type:uint32 6)
(define type:int64 7)
(define type:uint64 8)
(define type:text 9)
(define type:symbol 10)
(define type:list 11)
(define type:number 12)
(define type:bytes 13)
(define type:vector 14)

(define (lump-internal-type? n)
  (and (number? n)
       (>= n 0)
       (<= n 14)))

(define (lump-external-type? d)
  (or (number? d)
      (string? d)
      (list? d)
      (vector? d)
      (symbol? d)
      (bytes? d)))

(define (lump-argument-type? d)
  (or (typed? d)
      (lump-external-type? d)))

; write a datum with given type to the output stream
; use write-datum instead
(define (write-datum/typed type datum out-stream)
  (call-with-exception-handler 
   (lambda (exn)
     (raise-type-error 'write-datum "LUMP data type" type))
   (lambda ()
     ; procedure dispatcher
     ((vector-ref 
       (vector 
        (lambda ()
          (write-packed "<BB" out-stream type:bool (if datum 1 0)))
        (lambda ()
          (write-packed "<Bb" out-stream type:int8 datum))
        (lambda ()
          (write-packed "<BB" out-stream type:uint8 datum))
        (lambda ()
          (write-packed "<Bh" out-stream type:int16 datum))
        (lambda ()
          (write-packed "<BH" out-stream type:uint16 datum))
        (lambda ()
          (write-packed "<Bl" out-stream type:int32 datum))
        (lambda ()
          (write-packed "<BL" out-stream type:uint32 datum))
        (lambda ()
          (write-packed "<Bq" out-stream type:int64 datum))
        (lambda ()
          (write-packed "<BQ" out-stream type:uint64 datum))
        (lambda ()
          (write-packed "<BL" out-stream type:text (string-length datum))
          (write-string datum out-stream))
        (lambda () 
          (let ((s (symbol->string datum)))
            (write-packed "<BL" out-stream type:symbol (string-length s))
            (write-string s out-stream)))
        (lambda () 
          (write-packed "<BL" out-stream type:list (length datum))
          (for-each
           (lambda (element)
             (write-datum element out-stream))))
        (lambda ()
          (let ((s (number->string datum)))
            (write-packed "<BL" out-stream type:number (string-length s))
            (write-string s out-stream)))
        (lambda ()
          (write-packed "<BL" out-stream type:bytes (bytes-length datum))
          (write-bytes datum out-stream))
        (lambda ()
          (write-packed "<BL" out-stream type:vector (vector-length datum))
          (for ([element datum])
            (write-datum element out-stream))))
       type)))))

; write a given datum in LUMP format to out-stream
(define (write-datum datum out-stream)
  (cond
    ((typed? datum) (write-datum/typed (typed-type datum) (typed-value datum) out-stream))
    ((boolean? datum) (write-datum/typed type:bool datum out-stream))
    ((symbol? datum) (write-datum/typed type:symbol datum out-stream))
    ((number? datum) (write-datum/typed type:number datum out-stream))
    ((string? datum) (write-datum/typed type:text datum out-stream))
    ((bytes? datum) (write-datum/typed type:bytes datum out-stream))
    ((list? datum) (write-datum/typed type:list datum out-stream))
    ((vector? datum) (write-datum/typed type:vector datum out-stream))
    (else (raise-type-error 'write-datum "LUMP-supported value" datum))))

; read a given datum in LUMP format from in-stream
(define (read-datum in-stream)
  (define (read-length in-stream)
    (car (read-packed "<L" in-stream)))
  (define (read-in pack-type)
    (car (read-packed pack-type in-stream)))
  (define type (car (read-packed "<B" in-stream)))
  (call-with-exception-handler 
   (lambda (exn)
     (raise-type-error 'read-datum "LUMP data type" type))
   (lambda ()
     ; procedure dispatcher
     ((vector-ref 
       (vector 
        (lambda () 
          (if (> (read-in "<B") 0) #t #f))
        (lambda ()
          (read-in "<b"))
        (lambda ()
          (read-in "<B"))
        (lambda ()
          (read-in "<h"))
        (lambda ()
          (read-in "<H"))
        (lambda ()
          (read-in "<l"))
        (lambda ()
          (read-in "<L"))
        (lambda ()
          (read-in "<q"))
        (lambda ()
          (read-in "<Q"))
        (lambda ()
          (read-string (read-length in-stream) in-stream))
        (lambda ()
          (string->symbol (read-string (read-length in-stream) in-stream)))
        (lambda ()
          (for/list ([i (in-range (read-length in-stream))])
            (read-datum in-stream)))
        (lambda ()
          (string->number (read-string (read-length in-stream) in-stream)))
        (lambda ()
          (read-bytes (read-length in-stream) in-stream))
        (lambda ()
          (for/vector ([i (in-range (read-length in-stream))])
                      (read-datum in-stream))))
       type)))))

; the sequence number returned from a session's counter value
; is used for generating response messages (for the referer field)
(define (increase-seqnum! s)
  (set-session-counter! s (add1 (session-counter s))))

; write a message to an optional output stream
(define (write-message session message [out-stream (current-output-port)])
  ; automatically set the referer or extended flags if necessary
  (define (auto-set-flags flags referer? args?)
    (if args? 
        (bitwise-ior (if referer? (bitwise-ior flags 2) flags) 1)
        (if referer? (bitwise-ior flags 2) flags)))
  ; convert the arguments to a byte string in the above LUMP protocol format
  (define (write-args args out)
    (write-packed "<H" out (length args))
    (for-each 
     (lambda (datum)
       (write-datum datum out))
     args))
  ; main write
  (increase-seqnum! session)
  (write-packed "<BHL" out-stream 
                (pack-version-and-flags 
                 (message-version message)
                 (auto-set-flags (message-flags message) 
                                 (message-referer message)
                                 (not (null? (message-args message)))))
                (message-id message) 
                (session-counter session))
  (when (message-referer message)
    (write-packed "<L" out-stream (message-referer message)))
  (when (not (null? (message-args message)))
    (write-args (message-args message) out-stream))
  (session-counter session))

; read a message from an optional input stream
; returns the message
(define (read-message 
         [in-stream (current-input-port)] 
         [version-check-proc (lambda (version)
                                (when (> version protocol-version)
                                  (raise-argument-error 'read-message
                                                        (format "LUMP protocol version ~a or lower" protocol-version)
                                                        version)))])
  (define header (read-packed "<BHL" in-stream))
  (define flags (unpack-flags (first header)))
  (version-check-proc (unpack-version (first header)))
  (define referer (if (bitwise-bit-set? flags 1)
                      (car (read-packed "<L" in-stream))
                      #f))
  (if (bitwise-bit-set? flags 0)
      (message (second header) (third header) flags version (read-args in-stream) referer)
      (message (second header) (third header) flags version null referer)))

; read the arguments of a message in LUMP format from in-stream
(define (read-args in-stream)
  (for/list ([i (in-range (car (read-packed "<H" in-stream)))])
    (read-datum in-stream)))


; pack two nibbles into a number
; first nibble is most significant, second nibble least significant
(define (pack-nibbles b1 b2)
  (bitwise-ior (arithmetic-shift b1 4) b2))

; retrieve the first nibble from a number
(define (unpack-first-nibble n)
  (arithmetic-shift n -4))

; retrieve the second nibble from a number
(define (unpack-second-nibble n)
  (bitwise-bit-field n 0 4))

; pack protocol version (nibble) and flags (nibble) into one byte-sized number
(define (pack-version-and-flags version flags)
  (pack-nibbles (sub1 version) flags))

; unpack the version number from a packed byte
(define (unpack-version b)
  (add1 (unpack-first-nibble b)))

; unpack the flags from a packed byte
(define (unpack-flags b)
  (unpack-second-nibble b))