(import
  (chicken format)
  (chicken port)
  (only (chicken string) reverse-list->string)
  (srfi 34)	 ;;Exception Handling
  (srfi 35)	 ;;Exception Types
  (srfi 158)	 ;;Generators
  )

(cond-expand
  (chicken-5) (import (scheme))
  (chicken-6) (import (scheme base)))

;;; Exceptions

(define-condition-type &json-error &error
  json-error?
  (json-error-reason json-error-reason)
  (json-invalid-token json-invalid-token))

;;; Parameters

(define json-nesting-depth-limit (make-parameter +inf.0))	;; the maximum nesting depth of JSON that can be read.
(define json-number-of-character-limit (make-parameter +inf.0))	;; the maximum length of JSON input that can be read.

;;; JSON Reading

;;;;; Predicates

(define (json-null? obj) (eq? obj 'null))

(define (is-array-start? c)
  (char=? #\[ c))

(define (is-array-end? c)
  (char=? #\] c))

(define (is-object-start? c)
  (char=? #\{ c))

(define (is-object-end? c)
  (char=? #\} c))

(define (is-number-start? c)
  (or (char-numeric? c)
     (char=? #\+ c)
     (char=? #\- c)))

(define (is-string-start? c)
  (char=? #\" c))

(define (is-null-start? c)
  (char=? #\n c))

(define (is-bool-start? c)
  (or (char=? #\t c)
     (char=? #\f c)))

(define (is-whitespace? c)
  (or (char-whitespace? c)
     (char=? #\, c)
     (char=? #\: c)))

(define (is-delimiter? x)
  (or (eof-object? x)
     (is-whitespace? x)
     (is-array-start? x)
     (is-array-end? x)
     (is-object-start? x)
     (is-object-end? x)))

;;;;; Reading Logic

;; Generators are not peek-able. This wrapper fixes that.
(define (make-peekable-generator port-or-generator)
  (if (procedure? port-or-generator)
      ;; For generators...
      (let ((peeked-char #f))
        (values
         (lambda ()
            (or peeked-char
               (begin (set! peeked-char (port-or-generator))
                      peeked-char)))
         (lambda ()
            (if peeked-char
                (let ((peeked-char* peeked-char))
                  (set! peeked-char #f)
                  peeked-char*)
                (port-or-generator)))))
      ;; ...and for ports
      (values (cut peek-char port-or-generator)
              (cut read-char port-or-generator))))

(define (determine-reader-proc peek-char)
  (cond
   ((is-array-start? peek-char) read-array-start)
   ((is-array-end? peek-char) read-array-end)
   ((is-object-start? peek-char) read-object-start)
   ((is-object-end? peek-char) read-object-end)
   ((is-null-start? peek-char) read-null-sym)
   ((is-bool-start? peek-char) read-boolean)
   ((is-number-start? peek-char) read-number)
   ((is-string-start? peek-char) read-string)
   ((is-whitespace? peek-char) read-whitespace)
   (else (raise (make-condition &json-error 'json-error-reason "Invalid token" 'json-invalid-token peek-char)))))

(define (json-generator #!optional (port-or-generator (current-input-port)))
  (let-values (((generator-peek generator-pop)
                (make-peekable-generator port-or-generator)))
    (let* ((nesting-limit (json-nesting-depth-limit))
           (character-limit (json-number-of-character-limit)))
      (make-coroutine-generator
       (lambda (yield)
         (let loop ((next-char (generator-peek))
                    (json-nesting-depth #f)
                    (json-number-of-characters 0))
           (cond
            ((> (or json-nesting-depth 0) nesting-limit)
             (raise (make-condition &json-error
                                    'json-error-reason "Nesting depth exceeded"
                                    'json-invalid-token next-char)))
            ((> json-number-of-characters character-limit)
             (raise (make-condition &json-error
                                    'json-error-reason "Character limit exceeded"
                                    'json-invalid-token next-char)))
            ((and (eof-object? next-char)
                (< 0 json-nesting-depth))
             (raise (make-condition &json-error
                                    'json-error-reason "Unfinished JSON expression"
                                    'json-invalid-token next-char)))
            ((or (eof-object? next-char)
                (eq? 0 json-nesting-depth))
             #!eof)
            (else
             (let-values (((token new-charcount nesting-delta)
                           ((determine-reader-proc next-char)
                            json-number-of-characters generator-peek generator-pop)))
               (let ((json-nesting-depth* (+ (or json-nesting-depth 0) nesting-delta)))
                 (unless (null? token)
                   (yield token))
                 (unless (= 0 json-nesting-depth*)
                   (loop (generator-peek) json-nesting-depth* new-charcount))))))))))))

(define (read-whitespace charcount generator-peek generator-pop)
  (generator-pop)
  (values '() (+ charcount 1) 0))

(define (read-array-start charcount generator-peek generator-pop)
  (generator-pop)
  (values 'array-start (+ charcount 1) +1))

(define (read-array-end charcount generator-peek generator-pop)
  (generator-pop)
  (values 'array-end (+ charcount 1) -1))

(define (read-object-start charcount generator-peek generator-pop)
  (generator-pop)
  (values 'object-start (+ charcount 1) +1))

(define (read-object-end charcount generator-peek generator-pop)
  (generator-pop)
  (values 'object-end (+ charcount 1) -1))

(define (read-null-sym charcount generator-peek generator-pop)
  (generator-pop)
  (if (not (is-delimiter? (generator-peek)))
      (read-null-sym (+ charcount 1) generator-peek generator-pop)
      (values 'null (+ charcount 1) 0)))

(define (read-boolean charcount generator-peek generator-pop #!optional (accu '()))
  (set! accu (cons (generator-pop) accu))
  (let ((accu-str (reverse-list->string accu)))
    (cond
     ((string=? "true" accu-str) (values #t (+ charcount 1) 0))
     ((string=? "false" accu-str) (values #f (+ charcount 1) 0))
     (else (let ((next-char* (generator-peek)))
             (if (is-delimiter? next-char*)
                 (values accu charcount 0) ;; TODO: Throw error instead, this should never happen (no partial "true" or "false" possible)
                (read-boolean (+ charcount 1) generator-peek generator-pop accu)))))))

(define (read-number charcount generator-peek generator-pop #!optional (accu '()))
  (set! accu (cons (generator-pop) accu))
  (let ((next-char* (generator-peek)))
    (if (is-delimiter? next-char*)
        (values (string->number (reverse-list->string accu))
                (+ charcount 1) 0)
        (read-number (+ charcount 1) generator-peek generator-pop accu))))

(define (translate-escape char input-proc)
  (case char
    ((#\") #\")
    ((#\') #\')
    ((#\\) #\\)
    ((#\n) #\newline)
    ((#\t) #\tab)
    ((#\u) (read-unicode-escape input-proc))
    ((#\x) (read-hex-escape input-proc))
    ((#\O) #\null)
    ((#\r) #\return)
    ((#\|) #\|)
    ((#\v) #\vtab)
    ((#\a) #\alarm)
    ((#\b) #\backspace)))

(define (read-hex-escape generator-pop)
  (let ((pos1 (generator-pop))
        (pos2 (generator-pop)))
    (integer->char
     (string->number (list->string (list pos1 pos2)) 16))))

(define (read-unicode-escape generator-pop)
  (let ((pos1 (generator-pop))
        (pos2 (generator-pop))
        (pos3 (generator-pop))
        (pos4 (generator-pop)))
    (integer->char
     (string->number (list->string (list pos1 pos2 pos3 pos4)) 16))))

(define (read-string charcount generator-peek generator-pop #!optional (beginning? #t) (accu '()) (esc? #f))
  (cond
   (beginning?
    (begin
      (generator-pop) ; Get rid of the opening quote
      (read-string (+ charcount 1)
                   generator-peek generator-pop
                   #f '() #f)))
   ((and (not esc?) (char=? (generator-peek) #\"))
    (begin
      (generator-pop) ; Get rid of the closing quote
      (values (reverse-list->string accu)
              (+ charcount 1) 0)))
   ((and (not esc?) (char=? (generator-peek) #\\))
    (begin
      (generator-pop)
      (read-string (+ charcount 1) generator-peek generator-pop #f accu #t)))
   (else (let* ((current-pre-char (generator-pop))
                (current-char (if esc?
                                  (translate-escape current-pre-char generator-pop)
                                  current-pre-char)))
           (read-string (+ charcount 1)
                        generator-peek generator-pop
                        #f
                        (cons current-char accu)
                        #f)))))

;;;;; JSON folding

(define-record json-foldstate mode cache accumulator)

(define (json-proc obj foldstate)
  (if (json-foldstate? foldstate)
      (case (json-foldstate-mode foldstate)
        ((%array) (begin
                    (json-foldstate-accumulator-set!
                     foldstate
                     (cons obj (json-foldstate-accumulator foldstate)))
                    foldstate))
        ((%object) (begin
                     (if (null? (json-foldstate-cache foldstate))
                         (begin
                           (json-foldstate-cache-set! foldstate obj))
                         (begin
                           (json-foldstate-accumulator-set!
                            foldstate
                            (cons (cons (string->symbol (json-foldstate-cache foldstate)) obj)
                                  (json-foldstate-accumulator foldstate)))
                           (json-foldstate-cache-set! foldstate '())))
                     foldstate)))
      obj))

(define (json-array-start seed)
  (make-json-foldstate '%array '() '()))

(define (json-array-end seed)
  (list->vector (reverse (json-foldstate-accumulator seed))))

(define (json-object-start seed)
  (make-json-foldstate '%object '() '()))

(define (json-object-end seed)
  (reverse (json-foldstate-accumulator seed)))

(define (json-fold proc array-start array-end object-start object-end seed #!optional (port-or-generator (current-input-port)))
  (let ((generator (json-generator port-or-generator)))
    (let recurse ((seed seed)
                  (jump #f))
      (generator-fold
       (lambda (token seed)
         (case token 
           ((array-start) (proc
                           (call-with-current-continuation
                            (lambda (jump)
                              (recurse (array-start seed) jump)))
                           seed))
           ((array-end) (if jump
                            (jump (array-end seed))
                            (array-end seed)))
           ((object-start)  (proc
                             (call-with-current-continuation
                              (lambda (jump)
                               (recurse (object-start seed) jump)))
                             seed))
           ((object-end) (if jump
                             (jump (object-end seed))
                             (object-end seed)))
           (else (proc token seed))))
       seed generator))))

(define (json-read #!optional (port-or-generator (current-input-port)))
  (json-fold json-proc
             json-array-start json-array-end
             json-object-start json-object-end
             '() port-or-generator))

(define json-lines-read json-read)

(define json-sequence-read json-read)

;;; JSON Writing

(define (accumulate-boolean accumulator bool)
  (if bool (accumulator 'true) (accumulator 'false)))

(define (accumulate-null accumulator)
  (accumulator 'null))

(define (accumulate-number accumulator num)
  (accumulator num))

(define (accumulate-string accumulator str)
  (accumulator str))

(define (accumulate-vector accumulator vec)
  (accumulator #\[)
  (let ((len (vector-length vec)))
    (let loop ((index 0))
      (when (< index len)
        (when (> index 0)
          (accumulator #\,))
        (accumulate-dispatch accumulator
                             (vector-ref vec index))
        (loop (+ index 1)))))
  (accumulator #\]))

(define (accumulate-alist accumulator alist)
  (accumulator #\{)
  (when (> (length alist) 0)
    (let loop ((alist alist))
      (let ((kv-pair (car alist)))
        (if (not (pair? kv-pair))
            (raise (make-condition &json-error
                                   'json-error-reason "Unbalanced alist"
                                   'json-invalid-token kv-pair)))
        (accumulate-dispatch accumulator
                             (symbol->string (car kv-pair)))
        (accumulator #\:)
        (accumulate-dispatch accumulator (cdr kv-pair))
        (if (not (eq? '() (cdr alist)))
            (begin
              (accumulator #\,)
              (loop (cdr alist)))))))
  (accumulator #\}))

(define (accumulate-dispatch accumulator obj)
  (cond
   ((number? obj) (accumulate-number accumulator obj))
   ((string? obj) (accumulate-string accumulator obj))
   ((boolean? obj) (accumulate-boolean accumulator obj))
   ((eq? 'null obj) (accumulate-null accumulator))
   ((vector? obj) (accumulate-vector accumulator obj))
   ((list? obj) (accumulate-alist accumulator obj))))

(define (json-accumulator #!optional (port-or-accumulator (current-output-port)))
  (let ((accumulator (if (procedure? port-or-accumulator)
                         port-or-accumulator
                         (lambda (txt)
                           (if (char? txt)
                               (display txt port-or-accumulator)
                               (write txt port-or-accumulator)))))
        (leading-space? #f))
    (lambda (obj)
      (if leading-space? (accumulator #\space) (set! leading-space? #t))
      (accumulate-dispatch accumulator obj))))

(define (json-write obj #!optional (port-or-accumulator (current-output-port)))
  (let ((black-hole (make-output-port (lambda (poor-soul) #t) (lambda () #t))))
    ((json-accumulator black-hole) obj))
  ((json-accumulator port-or-accumulator) obj))