(import
  (scheme)
  (chicken format)
  (chicken port)
  (chicken string)
  (srfi-34)	 ;;Exception Handling
  (srfi-35)	 ;;Exception Types
  (srfi-158)	 ;;Generators
  )

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

(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.

(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)))

(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* ((input-generator (if (procedure? port-or-generator)
                              port-or-generator
                              (lambda () (read-char port-or-generator))))
         (nesting-limit (json-nesting-depth-limit))
         (character-limit (json-number-of-character-limit)))
    (make-coroutine-generator
     (lambda (yield)
       (let loop ((next-char (input-generator))
                  (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 next-char* new-charcount nesting-delta)
                         ((determine-reader-proc next-char)
                          json-number-of-characters next-char input-generator)))
             (unless (null? token)
                 (yield token))
             (loop next-char* (+ (or json-nesting-depth 0) nesting-delta) new-charcount)))))))))

(define (read-whitespace charcount next-char input-proc)
  (values '() (input-proc) (+ charcount 1) 0))

(define (read-array-start charcount next-char input-proc)
  (values 'array-start (input-proc) (+ charcount 1) +1))

(define (read-array-end charcount next-char input-proc)
  (values 'array-end (input-proc) (+ charcount 1) -1))

(define (read-object-start charcount next-char input-proc)
  (values 'object-start (input-proc) (+ charcount 1) +1))

(define (read-object-end charcount next-char input-proc)
  (values 'object-end (input-proc) (+ charcount 1) -1))

(define (read-null-sym charcount next-char input-proc)
  (if (not (is-delimiter? next-char))
      (read-null-sym (+ charcount 1) (input-proc) input-proc)
      (values 'null next-char charcount 0)))

(define (read-boolean charcount next-char input-proc #!optional (accu '()))
  (set! accu (cons next-char accu))
  (let ((accu-str (reverse-list->string accu)))
    (cond
     ((string=? "true" accu-str) (values #t (input-proc) (+ charcount 1) 0))
     ((string=? "false" accu-str) (values #f (input-proc) (+ charcount 1) 0))
     (else (let ((next-char* (input-proc)))
             (if (is-delimiter? next-char*)
                 (values accu next-char* charcount 0) ;; TODO: Throw error instead
                (read-boolean (+ charcount 1) next-char* input-proc accu)))))))

(define (read-number charcount next-char input-proc #!optional (accu '()))
  (set! accu (cons next-char accu))
  (let ((next-char* (input-proc)))
    (if (is-delimiter? next-char*)
        (values (string->number (reverse-list->string accu))
                next-char* (+ charcount 1) 0)
        (read-number (+ charcount 1) next-char* input-proc 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 input-proc)
  (let ((pos1 (input-proc))
        (pos2 (input-proc)))
    (integer->char
     (string->number (list->string (list pos1 pos2)) 16))))

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

(define (read-string charcount next-char input-proc #!optional (beginning? #t) (accu '()) (esc? #f))
  (cond
   (beginning?
    (read-string (+ charcount 1)
                 (input-proc)
                 input-proc
                 #f '() #f))
   ((and (not esc?) (char=? next-char #\"))
    (values (reverse-list->string accu)
            (input-proc) (+ charcount 1) 0))
   ((and (not esc?) (char=? next-char #\\))
    (read-string (+ charcount 1) (input-proc) input-proc #f accu #t))
   (else (let ((current-char (if esc?
                                 (translate-escape next-char input-proc)
                                 next-char)))
           (read-string (+ charcount 1)
                        (input-proc) input-proc
                        #f
                        (cons current-char accu)
                        #f)))))

(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 (json-foldstate-cache foldstate) obj)
                                  (json-foldstate-accumulator foldstate)))
                           (json-foldstate-cache-set! foldstate '())))
                     foldstate)))
      obj))

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

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

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

(define (array-end seed)
  (list->vector (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 array-start array-end object-start object-end '() port-or-generator))

(define json-lines-read json-read)

(define json-sequence-read json-read)

(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 ((max-index (- (vector-length vec) 1)))
    (let loop ((index 0))
      (accumulate-dispatch accumulator
                           (vector-ref vec index))
      (if (< index max-index)
          (begin (accumulator #\,) (accumulator #\space)
                 (loop (+ index 1))))))
  (accumulator #\]))

(define (accumulate-alist accumulator alist)
  (accumulator #\{)
  (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 #\:) (accumulator #\space)
      (accumulate-dispatch accumulator (cdr kv-pair))
      (if (not (eq? '() (cdr alist)))
          (begin
            (accumulator #\,) (accumulator #\space)
            (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))