#lang racket/base

(provide write-preserve/binary
         preserve->bytes)

(require racket/match)
(require (only-in racket/port call-with-output-bytes))
(require "record.rkt")
(require "float.rkt")
(require "annotation.rkt")
(require "varint.rkt")
(require "object-id.rkt")
(require racket/set)
(require racket/dict)
(require (only-in racket/list flatten))

(define (preserve->bytes v
                         #:canonicalizing? [canonicalizing? #t]
                         #:encode-embedded [encode-embedded #f]
                         #:write-annotations? [write-annotations? (not canonicalizing?)])
  (call-with-output-bytes
   (lambda (p) (write-preserve/binary v p
                                      #:canonicalizing? canonicalizing?
                                      #:encode-embedded encode-embedded
                                      #:write-annotations? write-annotations?))))

(define (write-preserve/binary v [out-port (current-output-port)]
                               #:canonicalizing? [canonicalizing? #t]
                               #:encode-embedded [encode-embedded0 #f]
                               #:write-annotations? [write-annotations? (not canonicalizing?)])
  (define encode-embedded (or encode-embedded0 object-id))

  (define (output-byte b)
    (write-byte b out-port))

  (define (output-bytes bs)
    (write-bytes bs out-port))

  (define (output-varint v)
    (write-varint v out-port))

  (define-syntax-rule (with-seq tag body ...)
    (begin (output-byte (+ tag #xB0))
           body ...
           (output-byte #x84)))

  (define (count-bytes tag bs)
    (output-byte (+ tag #xB0))
    (output-varint (bytes-length bs))
    (output-bytes bs))

  (define (prepare v) (preserve->bytes v #:canonicalizing? #t))

  (define (output-all vs)
    (for [(v (in-list vs))] (output v)))

  (define output-set
    (match* [canonicalizing? write-annotations?]
      [[#t #f] (lambda (v)
                 (for-each output-bytes
                           (sort (for/list [(e (in-set v))] (prepare e)) bytes<?)))]
      [[#t #t] (lambda (v)
                 (for-each output
                           (map cdr
                                (sort (for/list [(e (in-set v))] (cons (prepare e) e))
                                      bytes<?
                                      #:key car))))]
      [[#f _] (lambda (v) (for [(e (in-set v))] (output e)))]))

  (define (prepare-dict d)
    (sort (for/list [((k v) (in-dict d))] (list (prepare k) k v)) bytes<? #:key car))

  (define output-dict
    (match* [canonicalizing? write-annotations?]
      [[#t #f] (lambda (v)
                 (for-each (match-lambda [(list kb _ v) (output-bytes kb) (output v)])
                           (prepare-dict v)))]
      [[#t #t] (lambda (v)
                 (for-each (match-lambda [(list _ k v) (output k) (output v)])
                           (prepare-dict v)))]
      [[#f _] (lambda (v) (for [((k v) (in-dict v))] (output k) (output v)))]))

  (define (output v)
    (match v
      [#f (output-byte #x80)]
      [#t (output-byte #x81)]

      [(float v)
       (output-byte #x82)
       (output-bytes (real->floating-point-bytes v 4 #t))]
      [(? flonum?)
       (output-byte #x83)
       (output-bytes (real->floating-point-bytes v 8 #t))]

      [(annotated as _ v)
       (when write-annotations?
         (for [(a (in-list as))]
           (output-byte #x85)
           (output a)))
       (output v)]

      [(? integer?)
       (cond [(<= -3 v -1) (output-byte (+ v #xA0))]
             [(<= 0 v 12)  (output-byte (+ v #x90))]
             [else (define raw-bit-count (+ (integer-length v) 1)) ;; at least one sign bit
                   (define byte-count (quotient (+ raw-bit-count 7) 8))
                   (if (<= byte-count 16)
                       (output-byte (+ byte-count #xA0 -1))
                       (begin (output-byte #xB0)
                              (output-varint byte-count)))
                   (for [(shift (in-range (* byte-count 8) 0 -8))]
                     (output-byte (bitwise-bit-field v (- shift 8) shift)))])]

      [(? string?) (count-bytes 1 (string->bytes/utf-8 v))]
      [(? bytes?)  (count-bytes 2 v)]
      [(? symbol?) (count-bytes 3 (string->bytes/utf-8 (symbol->string v)))]

      [(record label fields) (with-seq 4 (output label) (output-all fields))]
      [(? list?)             (with-seq 5 (output-all v))]
      [(? set?)              (with-seq 6 (output-set v))]
      [(? dict?)             (with-seq 7 (output-dict v))]

      [other
       (output-byte #x86)
       (output (encode-embedded other))]))

  (output v))