;;;; rfb.scm
;
;
; Notes:
;
; * conversion to pixel format does not take advantage of native endianness
; * client mode has not been tested, yet
; * needs to be tested with more clients


(module rfb (rfb-server
	     rfb-connect
	     rfb-close
	     encoding-type
	     encoding-type-name
	     read-client-message
	     read-server-message
	     framebuffer-update-rectangle
	     framebuffer-update-rectangles
	     set-colour-map-entries
	     bell
	     server-cut-text
	     register-encoding-type
	     rfb-session?
	     rfb-session-encoding-type
	     rfb-session-pixel-format
	     rfb-session-width
	     rfb-session-height
	     rfb-session-shared?
	     rfb-session-name
	     rfb-session-input-port
	     rfb-session-output-port
	     pixel-format
	     pixel-format?
	     pixel-format-bits-per-pixel
	     pixel-format-depth
	     pixel-format-big-endian?
	     pixel-format-true-color?
	     pixel-format-red-max
	     pixel-format-green-max
	     pixel-format-blue-max
	     pixel-format-red-shift
	     pixel-format-green-shift
	     pixel-format-blue-shift
	     convert-to-pixel-format
	     rectangle
	     copy-rectangle
	     set-pixel-format
	     set-encoding-type
	     framebuffer-update-request
	     key-event
	     pointer-event
	     client-cut-text)

(import scheme chicken)
(use tcp srfi-4 defstruct regex-case extras srfi-1)
(import matchable miscmacros)

(define-constant +default-protocol-version+ "RFB 003.008\n")
(define-constant +default-security-type+ 1) ; none
(define-constant +default-port+ 5900)


(define-enum encoding-type encoding-type-name
  (Raw 0)
  (CopyRect 1)
  (RRE 2)
  (Hextile 5)
  (TRLE 15)
  (ZRLE 16)
  (Cursor -239)
  (DesktopSize -223))


(defstruct rfb-session
  name
  protocol-version			; INTEGER (minor version)
  width height
  (shared? #f)				; BOOL
  pixel-format                          ; U8VECTOR
  (encoding-type 0)				; FIXNUM
  input-port
  output-port)

(define-record rfb-pixel-format data)

(define decoders '())

(define (register-encoding-type encoding-type name decoder)
  (cond ((assq encoding-type decoders) =>
	 (lambda (a) (set-cdr! a decoder)))
	(else (push! (list encoding-type name decoder) decoders))))

(set! encoding-type
  (let ((encoding-type encoding-type))
    (lambda (name)
      (cond ((find (lambda (dc) (eq? name (second dc))) decoders) => first)
	    (else (encoding-type name))))))

(set! encoding-type-name
  (let ((encoding-type-name encoding-type-name))
    (lambda (et)
      (cond ((assq et decoders) => second)
	    (else (encoding-type-name et))))))

(define pixel-format? rfb-pixel-format?)

(define (pixel-format-bits-per-pixel pf) (u8vector-ref (rfb-pixel-format-data pf) 0))
(define (pixel-format-depth pf) (u8vector-ref (rfb-pixel-format-data pf) 1))

(define (pixel-format-big-endian? pf)
  (not (zero? (u8vector-ref (rfb-pixel-format-data pf) 2))))

(define (pixel-format-true-color? pf)
  (not (zero? (u8vector-ref (rfb-pixel-format-data pf) 3))))

(define (pixel-format-red-max pf)
  (let ((pf (rfb-pixel-format-data pf)))
    (fxior (fxshl (u8vector-ref pf 4) 8) (u8vector-ref pf 5))))

(define (pixel-format-green-max pf)
  (let ((pf (rfb-pixel-format-data pf)))
    (fxior (fxshl (u8vector-ref pf 6) 8) (u8vector-ref pf 7))))

(define (pixel-format-blue-max pf)
  (let ((pf (rfb-pixel-format-data pf)))
    (fxior (fxshl (u8vector-ref pf 8) 8) (u8vector-ref pf 9))))

(define (pixel-format-red-shift pf)
  (u8vector-ref (rfb-pixel-format-data pf) 10))

(define (pixel-format-green-shift pf)
  (u8vector-ref (rfb-pixel-format-data pf) 11))

(define (pixel-format-blue-shift pf)
  (u8vector-ref (rfb-pixel-format-data pf) 12))


;; support definitions

(define (write-u32 n p)
  (write-byte (bitwise-and #xff (arithmetic-shift n -24)) p)
  (write-byte (bitwise-and #xff (arithmetic-shift n -16)) p)
  (write-byte (bitwise-and #xff (arithmetic-shift n -8)) p)
  (write-byte (bitwise-and #xff n) p))

(define (write-u16 n p)
  (write-byte (bitwise-and #xff (arithmetic-shift n -8)) p)
  (write-byte (bitwise-and #xff n) p))

(define (hi u16) (fxshr u16 8))
(define (lo u16) (fxand #xff u16))

(define (read-u16 p)
  (let ((n (read-byte p)))
    (fxior (fxshl n 8) (read-byte p))))

(define (read-u32 p)
  (let ((v (read-u8vector 4 p)))
    (bitwise-ior 
     (arithmetic-shift (u8vector-ref v 0) 24)
     (arithmetic-shift (u8vector-ref v 1) 16)
     (arithmetic-shift (u8vector-ref v 2) 8)
     (u8vector-ref v 3))))

(define (write-u16sequence s p)
  (cond ((vector? s) 
	 (let ((len (vector-length s)))
	   (do ((i 0 (fx+ i 1)))
	       ((fx>= i len))
	     (write-u16 (vector-ref s i) p))))
	((u16vector? s)
	 (let ((len (u16vector-length s)))
	   (do ((i 0 (fx+ i 1)))
	       ((fx>= i len))
	     (write-u16 (u16vector-ref s i) p))))	 
	((pair? s)
	 (for-each (cut write-u16 <> p) s))
	(else (error "not a valid sequence" s))))

(define (canonicalize-u16sequence s loc)
  (cond ((vector? s)
	 (let ((len (vector-length s)))
	   (values 
	    (cut vector-ref s <>) 
	    len
	    0
	    (lambda (i)
	      (let ((i (fx+ i 1)))
		(and (fx<= i len) i))))))
	((u16vector? s)
	 (let ((len (u16vector-length s)))
	   (values 
	    (cut u16vector-ref s <>)
	    len
	    0
	    (lambda (i)
	      (let ((i (fx+ i 1)))
		(and (fx<= i len) i))))))
	((pair? s)
	 (values 
	  car
	  (length s)
	  s
	  (lambda (s) 
	    (and (pair? (cdr s)) (cdr s)))))
	(else (error loc "invalid sequence" s))))


;; connection handshake

(define (initialize-server-connection in out w h pf name)
  (let ((v 8)
	(shared #f))
    
    (define (establish-protocol-version)
      (write-string +default-protocol-version+ #f out)
      (let ((v (read-string 12 in)))
	(regex-case v
	  ("RFB ([[:digit:]]+)\\.([[:digit:]]+)\n" (_ major minor)
	   (string->number minor))
	  (else
	   (error "invalid protocol version" v)))))

    (define (establish-security-type)
      (cond ((<= v 3)
	     (write-u32 +default-security-type+ out))
	    (else
	     (write-u8vector '#u8(1 1) out) ; none
	     (read-byte in)	      ; ignore security-type by client
	     (when (>= v 8) (write-u32 0 out))))) ; security result: OK

    (define (initialization)
      (set! shared (not (zero? (read-byte in))))
      (write-u16 w out)
      (write-u16 h out)
      (write-u8vector (rfb-pixel-format-data pf) out)
      (write-u32 (string-length name) out)
      (write-string name #f out))

    (set! v (establish-protocol-version))
    (establish-security-type)
    (initialization)
    (make-rfb-session
     name: name
     protocol-version: v
     pixel-format: pf
     width: w
     height: h
     input-port: in
     output-port: out
     shared?: shared)))


(define (initialize-client-connection in out shared)
  (let ((v 8))
    
    (define (establish-protocol-version)
      (let ((v (read-string 12 in)))
	(regex-case v
	  ("RFB ([[:digit:]]+)\\.([[:digit:]]+)\n" (_ major minor)
	   (write-string v #f out))
	  (else
	   (error "invalid protocol version" v)))))

    (define (establish-security-type)
      (cond ((<= v 3)
	     (let ((st (read-u32 in)))
	       (unless (= 1 st)		; None
		 (error "security-type not supported" st))))
	    (else
	     (let loop ((n (read-byte in)))
	       (if (zero? n)
		   (error "no supported security-type available")
		   (let ((st (read-byte in)))
		     (cond ((= 1 st)
			    (read-string (sub1 n) in) ; skip rest
			    (write-byte 1 out)
			    (when (>= v 8)
			      (let ((sr (read-u32 in)))
				(unless (zero? sr)
				  (let ((len (read-u32 in)))
				    (error "server error" (read-string len in)))))))
			   (else (loop (sub1 n))))))))))

    (define (initialization)
      (write-byte (if shared 1 0) out)
      (let* ((w (read-u16 in))
	     (h (read-u16 in))
	     (pf (read-u8vector 16 in))
	     (len (read-u32 in))
	     (name (read-string len in)))
	(make-rfb-session
	 name: name
	 width: w
	 height: h
	 protocol-version: v
	 pixel-format: (make-rfb-pixel-format pf)
	 input-port: in
	 output-port: out
	 shared?: shared)))

    (establish-protocol-version)
    (establish-security-type)
    (initialization)))


;; create pixel-format vector - defaults to 32-bit pixel values, true-color, #xRRGGBB
;; (endianness defaults to machine's endianness)
(define (pixel-format #!key
		      (bits-per-pixel 32)
		      (depth 24)
		      (big-endian (case (machine-byte-order) 
				    ((little-endian) #f)
				    (else #t)))
		      (true-color #t)
		      (red-max 255)
		      (green-max 255)
		      (blue-max 255)
		      (red-shift 16)
		      (green-shift 8)
		      (blue-shift 0))
  (make-rfb-pixel-format
   (u8vector
    bits-per-pixel
    depth
    (if big-endian 1 0)
    (if true-color 1 0)
    (hi red-max) (lo red-max)
    (hi green-max) (lo green-max)
    (hi blue-max) (lo blue-max)
    red-shift
    green-shift
    blue-shift
    0 0 0)))				; padding


;; client messages

(define (read-client-message rs)
  (let* ((in (rfb-session-input-port rs))
	 (m (read-byte in)))
    (case m
      ((#!eof) m)
      ((0) 				; SetPixelFormat
       (read-string 3 in)		; padding
       (let ((pf (make-rfb-pixel-format (read-u8vector 16 in))))
	 (rfb-session-pixel-format-set! rs pf)
	 `(SetPixelFormat ,pf)))
      ((2)				; SetEncodingType
       (read-byte in)			; padding
       (let loop ((n (read-u16 in)) (ets '()))
	 (if (zero? n)
	     `(SetEncodingType ,@(reverse ets))
	     (loop (sub1 n) (cons (read-u32 in) ets)))))
      ((3)				; FramebufferUpdateRequest
       (let* ((inc (not (zero? (read-byte in))))
	      (x (read-u16 in))
	      (y (read-u16 in))
	      (w (read-u16 in))
	      (h (read-u16 in)))
	 `(FramebufferUpdateRequest ,inc ,x ,y ,w ,h)))
      ((4)				; KeyEvent
       (let ((d (read-byte in)))
	 (read-string 2 in)		; padding
	 `(KeyEvent ,(not (zero? d)) ,(read-u32 in))))
      ((5)				; PointerEvent
       (let* ((bm (read-byte in))
	      (x (read-u16 in))
	      (y (read-u16 in)))
	 `(PointerEvent ,bm ,x ,y)))
      ((6)				; ClientCutText
       (read-string 3 in)		; padding
       (let ((len (read-u32 in)))
	 `(ClientCutText ,(read-string len in))))
      (else (error "unsupported client message" m)))))

(define (set-pixel-format rs pf)
  (let ((out (rfb-session-output-port rs)))
    (write-u8vector '#u8(0 0 0 0) out)
    (write-u8vector (rfb-pixel-format-data pf) out)
    (rfb-session-pixel-format-set! rs pf)))

(define (set-encoding-type rs et)
  (let ((out (rfb-session-output-port rs)))
    (write-u8vector '#u8(2 0) out)
    (write-u16 et out)))

(define (framebuffer-update-request rs x y w h #!optional (inc #t))
  (let ((out (rfb-session-output-port rs)))
    (write-byte 3 out)
    (write-byte (if inc 1 0) out)
    (write-u16 x out)
    (write-u16 y out)
    (write-u16 w out)
    (write-u16 h out)))

(define (key-event rs k #!optional (down #t))
  (let ((out (rfb-session-output-port rs)))
    (write-u8vector (u8vector 4 (if down 1 0) 0 0) out)
    (write-u16 k out)))

(define (pointer-event rs bm x y)
  (let ((out (rfb-session-output-port rs)))
    (write-byte 5 out)
    (write-byte bm out)
    (write-u16 x out)
    (write-u16 y out)))

(define (client-cut-text rs str)
  (let ((out (rfb-session-output-port rs)))
    (write-u8vector '#u8(6 0 0 0) out)
    (write-u32 (string-length str) out)
    (write-string str #f out)))


;;; rectangles

(defstruct rect encoding x y w h pf data rawdata)

(define (rectangle x y w h data #!optional (encoding 0)) ; Raw
  (make-rect encoding: encoding x: x y: y w: w h: h pf: #f data: data rawdata: data))

(define (copy-rectangle x y w h x2 y2)
  (let ((p2 (cons x2 y2)))
    (make-rect
     encoding: 1 x: x y: y w: w h: h pf: #f data: p2 rawdata: p2))) ; CopyRect

(define (convert-to-pixel-format rect pf)
  (let ((data (rect-data rect)))
    (cond ((eq? pf (rect-pf rect)) rect)
	  ((or (pair? data) (procedure? data)) rect)
	  ((pixel-format-true-color? pf)
	   (let* ((rm (pixel-format-red-max pf))
		  (gm (pixel-format-green-max pf))
		  (bm (pixel-format-blue-max pf))
		  (rs (pixel-format-red-shift pf))
		  (gs (pixel-format-green-shift pf))
		  (bs (pixel-format-blue-shift pf))
		  (bpp (pixel-format-bits-per-pixel pf))
		  (be (pixel-format-big-endian? pf))
		  (bytes (fx* (fx/ bpp 8) (fx* (rect-w rect) (rect-h rect))))
		  (rdata (make-u8vector bytes))
		  (len (u32vector-length data)))
	     (do ((i 0 (fx+ i 1))
		  (p 0))
		 ((fx>= i len))
	       (let* ((pixel (u32vector-ref data i))
		      (r (fxshr pixel 16))
		      (g (fxand #xff (fxshr pixel 8)))
		      (b (fxand #xff pixel)))
		 ;;XXX this can be done more efficiently
		 (case bpp
		   ((8)
		    (u8vector-set!
		     rdata p 
		     (bitwise-ior
		      (fxshl (fxand rm r) rs)
		      (fxshl (fxand gm g) gs)
		      (fxshl (fxand bm b) bs)))
		    (set! p (fx+ p 1)))
		   ((16)
		    (let ((p2 (bitwise-ior
			       (fxshl (fxand rm r) rs)
			       (fxshl (fxand gm g) gs)
			       (fxshl (fxand bm b) bs))))
		      (cond (be
			     (u8vector-set! rdata p (fxshr p2 8))
			     (u8vector-set! rdata (fx+ p 1) (fxand p2 #xff))))
			    (else
			     (u8vector-set! rdata p (fxand p2 #xff))
			     (u8vector-set! rdata (fx+ p 1) (fxshr p2 8))))
		      (set! p (fx+ p 2)))
		   ((32)
		    (let ((p2 (bitwise-ior
			       (fxshl (fxand rm r) rs)
			       (fxshl (fxand gm g) gs)
			       (fxshl (fxand bm b) bs))))
		      (cond (be
			     (u8vector-set! rdata p (bitwise-and #xff (arithmetic-shift p2 -24)))
			     (u8vector-set! rdata (fx+ p 1) 
					    (bitwise-and #xff (arithmetic-shift p2 -16)))
			     (u8vector-set! rdata (fx+ p 2)
					    (bitwise-and #xff (arithmetic-shift p2 -8)))
			     (u8vector-set! rdata (fx+ p 3) (bitwise-and p2 #xff)))
			    (else
			     (u8vector-set! rdata (fx+ p 3)
					    (bitwise-and #xff (arithmetic-shift p2 -24)))
			     (u8vector-set! rdata (fx+ p 2) 
					    (bitwise-and #xff (arithmetic-shift p2 -16)))
			     (u8vector-set! rdata (fx+ p 1)
					    (bitwise-and #xff (arithmetic-shift p2 -8)))
			     (u8vector-set! rdata p (bitwise-and p2 #xff))))
		      (set! p (fx+ p 4))))
		   (else (error "invalid bpp" bpp)))))
	     (update-rect rect pf: pf rawdata: rdata)))
	  (else				; indexed color
	   (assert 
	    (u16vector? data)
	    "client requested indexed color - rectangle-data should be u16vector" rect)
	   (let* ((bpp (pixel-format-bits-per-pixel pf))
		  (be (pixel-format-big-endian? pf))
		  (bytes (fx* (fx/ bpp 8) (fx* (rect-w rect) (rect-h rect))))
		  (rdata (make-u8vector bytes))
		  (len (u16vector-length data)))
	     (do ((i 0 (fx+ i 1))
		  (p 0))
		 ((fx>= i len))
	       (let ((index (u16vector-ref data i)))
		 ;;XXX this as well
		 (case bpp
		   ((8)
		    (u8vector-set! rdata p index)
		    (set! p (fx+ p 1)))
		   ((16)
		    (cond (be
			   (u8vector-set! rdata p (fxshr index 8))
			   (u8vector-set! rdata (fx+ index 1) (fxand index #xff)))
			  (else
			   (u8vector-set! rdata p (fxand index #xff))
			   (u8vector-set! rdata (fx+ p 1) (fxshr index 8))))
		    (set! p (fx+ p 2)))
		   ((32)
		    (cond (be
			   (u8vector-set! rdata p (fxand #xff (fxshr index 24)))
			   (u8vector-set! rdata (fx+ p 1) (fxand #xff (fxshr index 16)))
			   (u8vector-set! rdata (fx+ p 2) (fxand #xff (fxshr index 8)))
			   (u8vector-set! rdata (fx+ p 3) (fxand index #xff)))
			  (else
			   (u8vector-set! rdata (fx+ p 3) (fxand #xff (fxshr index 24)))
			   (u8vector-set! rdata (fx+ p 2) (fxand #xff (fxshr index 16)))
			   (u8vector-set! rdata (fx+ p 1) (fxand #xff (fxshr index 8)))
			   (u8vector-set! rdata p (fxand index #xff))))
		    (set! p (fx+ p 4)))
		   (else (error "invalid bpp" bpp)))))
	     (update-rect rect pf: pf rawdata: rdata))))))


;;; server messages

(define (read-server-message rs)
  (let* ((in (rfb-session-input-port rs))
	 (et (rfb-session-encoding-type rs))
	 (bpp (/ (pixel-format-bits-per-pixel (rfb-session-pixel-format rs)) 8))
	 (m (read-byte in)))
    (case m
      ((#!eof) m)
      ((0)				; FrameBufferUpdate
       (read-byte in)			; padding
       (let loop ((n (read-u16 in)) (rects '()))
	 (if (zero? n)
	     `(FramebufferUpdate ,@(reverse rects))
	     (let* ((x (read-u16 in))
		    (y (read-u16 in))
		    (w (read-u16 in))
		    (h (read-u16 in))
		    (et (read-u32 in)))
	       (case et
		 ((0)			; Raw
		  (let ((data (read-u8vector (* w h bpp) in)))
		    (loop (sub1 n) (cons `(Raw #(,x ,y ,w ,h ,data)) rects))))
		 ((1)			; CopyRect
		  (let ((x2 (read-u16 in))
			(y2 (read-u16 in)))
		    (loop (sub1 n) (cons `(CopyRect #(,x ,y ,w ,h ,x2 ,y2)) rects))))
		 (else
		  (cond ((assq et decoders) =>
			 (lambda (a)
			   (loop (sub1 n) (cons ((cdr a) in) rects))))
			((memq et '(-239 -223)) ; Cursor, DesktopSize ignored
			 (loop (sub1 n) rects))
			(else
			 (error 'read-server-message "encoding-type not supported" et)))))))))
      ((1)				; SetColourMapEntries
       (read-byte in)			; padding
       (let* ((first (read-u16 in))
	      (len (read-u16 in))
	      (v (make-u16vector len)))
	 (let loop ((i 0))
	   (if (fx>= i len)
	       `(SetColourMapEntries ,first ,v)
	       (begin
		 (u16vector-set! v i (read-u16 in)))))))
      ((2)				; Bell
       '(Bell))
      ((3)				; ServerCutText
       (read-string 3 in)		; padding
       (let ((len (read-u32 in)))
	 `(ServerCutText ,(read-string len in))))
      (else (error 'read-server-message "invalid server message" m)))))

(define (framebuffer-update-rectangles rs rects)
  (let ((out (rfb-session-output-port rs)))
    (write-u8vector '#u8(0 0) out)
    (write-u16 (length rects) out)
    (for-each
     (lambda (rect)
       (let ((w (rect-w rect))
	     (h (rect-h rect))
	     (pf (rfb-session-pixel-format rs)))
	 (write-u16 (rect-x rect) out)
	 (write-u16 (rect-y rect) out)
	 (write-u16 w out)
	 (write-u16 h out)
	 (let ((rect (convert-to-pixel-format rect pf)))
	   (match (rect-rawdata rect)
	     ((? procedure? g)
	      (g w h (rect-encoding rect) pf out))
	     ((x2 . y2)
	      (write-u32 (rect-encoding rect) out)
	      (write-u16 x2 out)
	      (write-u16 y2 out))
	     (rdata
	      (write-u32 (rect-encoding rect) out)
	      (write-u8vector rdata out))))))
     rects)
    (flush-output out)))

(define (framebuffer-update-rectangle rs rect)
  (framebuffer-update-rectangles rs (list rect)))

(define (set-colour-map-entries rs colors #!optional (first 0))
  (let ((out (rfb-session-output-port rs))
	(et (rfb-session-encoding-type rs)))
    (write-u8vector '#u8(1 0) out)
    (write-u16 first out)
    (let-values (((get len start next)
		  (canonicalize-u16sequence colors 'set-colour-map-entries)))
      (write-u16 len out)
      (do ((i start (next i)))
	  ((not i))
	(write-u16sequence (get i) out)))))

(define (bell rs)
  (write-byte 2 (rfb-session-output-port rs)))

(define (server-cut-text rs str)
  (let ((out (rfb-session-output-port rs)))
    (write-u8vector '#u8(3 0 0 0) out)
    (write-u32 (string-length str) out)
    (write-string str #f out)))


(define (rfb-server #!key (port +default-port+) (access #f) (name "rfb"))
  (let ((listener (tcp-listen port 10 access)))
    (lambda (w h #!optional (pf (pixel-format)))
      (let-values (((in out) (tcp-accept listener)))
	(initialize-server-connection in out w h pf name)))))

(define (rfb-connect host #!key (port +default-port+) shared)
  (let-values (((in out) (tcp-connect host port)))
    (initialize-client-connection in out shared)))

(define (rfb-close rfb)
  (close-input-port (rfb-session-input-port rfb))
  (close-output-port (rfb-session-output-port rfb)))


;;XXX missing
;
; - set cursor shape using cursor pseudo encoding
; - set desktop size using desktop-size pseudo encoding


)