#lang racket/base

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

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

(define (preserve->iolist v
                          #:canonicalizing? [canonicalizing? #t]
                          #:encode-embedded [encode-embedded0 #f]
                          #:write-annotations? [write-annotations? (not canonicalizing?)])
  (define encode-embedded (or encode-embedded0 object-id))

  (define (prepare v)
    (iolist->bytes (preserve->iolist v #:canonicalizing? #t #:encode-embedded encode-embedded0)))

  (define (length-prefixed-iolist i)
    (define c (count-iolist i))
    (cons (varint->iolist (counted-iolist-length c)) c))

  (define (length-prefixed v) (length-prefixed-iolist (encode v)))

  (define (encode v)
    (match v
      [#f                          #xA0]
      [#t                          #xA1]
      [(float v)             (cons #xA2 (real->floating-point-bytes v 4 #t))]
      [(? flonum?)           (cons #xA2 (real->floating-point-bytes v 8 #t))]
      [(? integer?)          (cons #xA3 (when (not (zero? v))
                                          (define nbits (bitwise-and (+ (integer-length v) 8) -8))
                                          (for/list [(shift (in-range nbits 0 -8))]
                                            (bitwise-bit-field v (- shift 8) shift))))]
      [(? string?)           (list #xA4 (string->bytes/utf-8 v) 0)]
      [(? bytes?)            (cons #xA5 v)]
      [(? symbol?)           (cons #xA6 (string->bytes/utf-8 (symbol->string v)))]
      [(record label fields) (list #xA7 (length-prefixed label) (map length-prefixed fields))]
      [(? list?)             (cons #xA8 (map length-prefixed v))]
      [(? set?)              (cons #xA9 (set->iolist v))]
      [(? dict?)             (cons #xAA (dict->iolist v))]
      [(embedded value)      (cons #xAB (encode (encode-embedded value)))]

      [(annotated as _ v)
       (if (and write-annotations? (pair? as))
           (list #xBF (length-prefixed v) (map length-prefixed as))
           (encode v))]

      [other (error 'preserve->iolist "Attempt to serialize non-preserve: ~v" other)]))

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

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

  (define dict->iolist
    (match* [canonicalizing? write-annotations?]
      [[#t #f] (lambda (v)
                 (map (match-lambda [(list kb _ v) (cons (length-prefixed-iolist kb)
                                                         (length-prefixed v))])
                      (prepare-dict v)))]
      [[#t #t] (lambda (v)
                 (map (match-lambda [(list _ k v) (cons (length-prefixed k)
                                                        (length-prefixed v))])
                      (prepare-dict v)))]
      [[#f _] (lambda (v)
                (for/list [((k v) (in-dict v))] (cons (length-prefixed k) (length-prefixed v))))]))

  (encode v))

(define (preserve->bytes v
                         #:canonicalizing? [canonicalizing? #t]
                         #:encode-embedded [encode-embedded #f]
                         #:write-annotations? [write-annotations? (not canonicalizing?)])
  (iolist->bytes (preserve->iolist v
                                   #: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-embedded #f]
                               #:write-annotations? [write-annotations? (not canonicalizing?)])
  (write-iolist (preserve->iolist v
                                  #:canonicalizing? canonicalizing?
                                  #:encode-embedded encode-embedded
                                  #:write-annotations? write-annotations?)
                out-port))