;;; Scheme indentation

(import (chicken irregex))
(import (chicken plist))
(import (chicken string))
(import (chicken process-context))
(import (chicken io))
(import (chicken file))
(import (chicken condition))
(import srfi-1)

(define tabs-to-spaces #t)
(define tab-stops 8)
(define space-closing-parens #f)
(define leader-restart #f)
(define detect-data #t)
(define use-dir-locals #t)

(define block-forms
  '(and-let*
    append-map
    begin
    call-with-current-continuation
    call-with-input-file
    call-with-output-file
    call-with-values
    call/cc
    case
    case-lambda
    condition-case
    defstruct
    dynamic-wind
    else
    er-macro-transformer
    filter-map
    fluid-let
    for-each
    getter-with-setter
    guard
    handle-exceptions
    ir-macro-transformer
    lambda
    let
    let*
    let*-values
    let-values
    letrec
    letrec*
    letrec-values
    map
    match
    match-lambda
    match-lambda*
    match-let
    match-let*
    match-letrec
    module
    parameterize
    receive
    set!
    string-for-each
    string-map
    syntax-rules
    syntax-case
    unless
    vector-for-each
    vector-map
    when
    with-exception-handler
    with-input-from-file
    with-input-from-port
    with-input-from-string
    with-output-to-file
    with-output-to-port
    with-output-to-string))

(for-each
  (lambda (p)
    (put! p 'scheme-indent 2))
  block-forms)

(define *data* #f)
(define *quasi* 0)

(define separators '(#\(#\) #\' #\; #\[ #\] #\{ #\} #\| #\"))

(define (spaces n)
  (let loop ((n n))
    (cond ((<= n 0))
          ((and (not tabs-to-spaces) (>= n tab-stops))
           (write-char #\tab)
           (loop (- n tab-stops)))
          (else
            (write-char #\space)
            (loop (sub1 n))))))

(define (expand-tab col)
  (let ((n (modulo col tab-stops)))
    (spaces (if (zero? n)
                tab-stops
                (- tab-stops (modulo col tab-stops))))))

(define (skip-line align)
  (let ((c (read-char)))
    (cond ((eof-object? c) 0)
          ((char=? #\newline c)
           (write-char c)
           (indent 0 align #t))
          (else
            (write-char c)
            (skip-line align)))))

(define (quoted qc col align)
  (let ((c (read-char)))
    (cond ((eof-object? c) col)
          ((char=? #\\ c)
           (write-char #\\)
           (let ((c2 (read-char)))
             (cond ((eof-object? c2) (add1 col))
                   (else
                     (write-char c2)
                     (quoted qc (+ col 2) align)))))
          ((char=? qc c)
           (write-char qc)
           (indent (add1 col) align #f))
          ((char=? #\tab c)
           (write-char c)
           (let ((n (modulo col tab-stops)))
             (if (zero? n)
                 (set! col (+ col tab-stops))
                 (set! col (+ col (- tab-stops n))))
             (quoted qc col align)))
          (else
            (write-char c)
            (quoted qc (add1 col) align)))))

(define (special col align)
  (let ((c (read-char)))
    (case c
      ((#\|)
       (write-char c)
       (skip-block-comment col align))
      ((#\<)
       (write-char c)
       (let ((c (read-char)))
         (write-char c)
         (if (memv c '(#\# #\<))
             (let ((t (read-line)))
               (print t)
               (let loop ()
                 (let ((ln (read-line)))
                   (print ln)
                   (if (or (eof-object? ln) (string=? ln t))
                       (indent 0 align #t)
                       (loop)))))
             (skip-word (+ col 2) align))))
      ((#\\)
       (write-char c)
       (let ((c (read-char)))
         (if (eof-object? c)
             col
             (begin
               (write-char c)
               (skip-word (add1 col) align)))))
      (else
        (write-char c)
        (skip-word (add1 col) align)))))

(define (skip-word col align)
  (let ((c (peek-char)))
    (cond ((eof-object? c) col)
          ((or (char-whitespace? c)
               (memv c separators))
           (indent col align #f))
          (else
            (write-char (read-char))
            (skip-word (add1 col) align)))))

(define (skip-block-comment col align)
  (let loop ((n 0) (col col))
    (let ((c (read-char)))
      (cond ((eof-object? c) col)
            ((char=? #\| c)
             (write-char c)
             (let ((c (read-char)))
               (cond ((eof-object? c) (add1 col))
                     ((char=? #\# c)
                      (write-char c)
                      (if (zero? n)
                          (indent (+ col 2) align #f)
                          (loop (sub1 n) (+ col 2))))
                     (else
                       (write-char c)
                       (loop n (+ col 2))))))
            ((char=? #\newline c)
             (write-char c)
             (loop n 0))
            ((char=? #\# c)
             (let ((c (read-char)))
               (cond ((eof-object? c) (add1 col))
                     ((char=? #\| c)
                      (write-char c)
                      (loop (add1 n) (+ col 2)))
                     (else
                       (write-char c)
                       (loop n (+ col 2))))))
            (else
              (write-char c)
              (loop n (add1 col)))))))

(define (peek-open)
  (let ((c (peek-char)))
    (cond ((eof-object? c) #f)
          ((memv c '(#\( #\[ #\{)))
          ((char-whitespace? c)
           (read-char)
           (peek-open))
          (else #f))))

(define (indent col align start)
  (define (normal c col)
    (when start
      (spaces align)
      (set! col align))
    (set! start #f)
    (write-char c)
    (set! col (add1 col))
    (case c
      ((#\')
       (cond ((peek-open) =>
              (lambda (c)
                (write-char (read-char))
                (set! col (add1 col))
                (indent (fluid-let ((*data* detect-data))
                          (block col align))
                        align #f)))
             (else (indent col align #f))))
      ((#\`)
       (cond ((peek-open) =>
              (lambda (c)
                (write-char (read-char))
                (set! col (add1 col))
                (indent (fluid-let ((*data* detect-data)
                                    (*quasi* (add1 *quasi*)))
                          (block col align))
                        align #f)))
             (else (indent col align #f))))
      ((#\,) ; implies ",@"
       (cond ((and *data* (positive? *quasi*) (peek-open)) =>
              (lambda (c)
                (write-char (read-char))
                (set! col (add1 col))
                (indent (fluid-let ((*data* #f)
                                    (*quasi* (sub1 *quasi*)))
                          (block col align))
                        align #f)))
             (else (indent col align #f))))
      ((#\;) (skip-line align))
      ((#\#) (special col align))
      ((#\|) (quoted #\| col align))
      ((#\") (quoted #\" col align))
      ((#\) #\] #\})
       (cond ((and space-closing-parens
                   (memv (peek-nonspace) '(#\) #\] #\})))
              (write-char #\space)
              (add1 col))
             (else col)))
      ((#\( #\[ #\{)
       (indent (block col align) align #f))
      (else (word col align #f))))
  (let loop ((col col))
    (let ((c (read-char)))
      (cond ((eof-object? c)
             (when start (newline))
             col)
            ((char=? #\tab c)
             (if start
                 (loop (+ col (- tab-stops (modulo col tab-stops))))
                 (begin
                   (expand-tab col)
                   (loop (add1 col)))))
            ((char=? #\space c)
             (if start
                 (loop (add1 col))
                 (begin
                   (write-char #\space)
                   (loop (add1 col)))))
            ((char=? #\newline c)
             (write-char c)
             (indent 0 align #t))
            ((and leader-restart (zero? col))
             (set! align 0)
             (normal c 0))
            (else (normal c col))))))

(define (peek-nonspace)
  (let ((c (peek-char)))
    (if (and (not (eof-object? c))
             (not (char=? #\newline c))
             (char-whitespace? c))
        (begin
          (read-char)
          (peek-nonspace))
        c)))

(define (interpret str col align)
  (let ((n (if (irregex-match "^define" str)
               2
               (get (string->symbol str) 'scheme-indent)))
        (c2 (peek-nonspace)))
    (cond ((eof-object? c2) col)
          ((memv c2 '(#\) #\] #\}))
           (indent col align #f))
          ((eqv? c2 #\newline)
           (write-char (read-char))
           (indent 0 (add1 align) #t))
          (else
            (write-char #\space)
            (set! col (add1 col))
            (cond ((string=? "quote" str)
                   (fluid-let ((*data* detect-data))
                     (indent col (+ align 1) #f)))
                  ((string=? "quasiquote" str)
                   (fluid-let ((*data* detect-data)
                               (*quasi* (add1 *quasi*)))
                     (indent col (+ align 1) #f)))
                  ((and *data* (positive? *quasi*)
                        (or (string=? "unquote" str)
                            (string=? "unquote-splicing" str)))
                   (fluid-let ((*data* #f)
                               (*quasi* (sub1 *quasi*)))
                     (indent col (+ align 1) #f)))
                  (n (indent col (+ align n -1) #f))
                  (else (indent col col #f)))))))


(define (word col align open)
  (let loop ((col col) (w '()))
    (let ((c (peek-char)))
      (cond ((eof-object? c) col)
            ((or (char-whitespace? c)
                 (memv c separators))
             (if open
                 (interpret (reverse-list->string w) col align)
                 (indent col align #f)))
            (else
              (write-char (read-char))
              (loop (add1 col) (cons c w)))))))

(define (block col align)
  (let loop ()
    (let ((c (peek-char)))
      (cond ((eof-object? c) col)
            ((char=? #\newline c)
             (indent col col #f))
            ((char-whitespace? c)
             (read-char)
             (loop))
            ((memv c '(#\) #\] #\}))
             (write-char (read-char))
             (add1 col))
            ((memv c '(#\( #\[ #\{))
             (write-char (read-char))
             (let ((col2 (block (add1 col) (add1 col))))
               (indent col2 col #f)))
            ((char=? #\# c)
             (write-char (read-char))
             (special (add1 col) col))
            ((memv c separators)
             (indent col align #f))
            (else (word col col (not *data*)))))))

(define (load-dir-locals)
  (and (file-exists? ".dir-locals.el")
       (let loop ((form (handle-exceptions ex
                          (display "error while loading .dir-locals.el"
                                   (current-error-port))
                          (with-input-from-file ".dir-locals.el"
                            read))))
         (when (pair? form)
           (when (and (pair? (car form))
                      (memq (caar form) '(nil scheme-mode)))
             (for-each
               (lambda (p)
                 (when (pair? p)
                   (case (car p)
                     ((tab-width)
                      (when (number? (cdr p))
                        (set! tab-stops (cdr p))))
                     ((indent-tabs-mode)
                      (set! tabs-to-spaces (not (eq? (cdr p) 't)))))))
               (cdar form)))
           (loop (cdr form)))
         #t)))

(define (load-editorconfig)
  (and (file-exists? ".editorconfig")
       (with-input-from-file ".editorconfig"
         (lambda ()
           (let loop ()
             (let ((ln (read-line)))
               (cond ((eof-object? ln) #t)
                     ((irregex-match "^ *indent_style *= *(tab|space)"
                                     ln) =>
                      (lambda (m)
                        (let ((s (irregex-match-substring m 1)))
                          (set! tabs-to-spaces (string=? s "space"))
                          (loop))))
                     ((irregex-match "^ *tab_width *= *([0-9 ]+)"
                                     ln) =>
                      (lambda (m)
                        (set! tab-stops (string->number (irregex-match-substring m 1)))
                        (loop)))
                     (else (loop)))))))))

(define start-align 0)

(define (main)
  (let loop ((args (append (string-split (or (get-environment-variable "SCHEME_INDENT")
                                             ""))
                           (command-line-arguments))))
    (if (null? args)
        (begin
          (when use-dir-locals
            (or (load-editorconfig) (load-dir-locals)))
          (let loop2 ((f #f))
            (unless (eof-object? (peek-char))
              (indent 0 start-align f)
              (loop2 #f))))
        (let* ((arg (car args))
               (args (cdr args)))
          (cond ((string=? "-t" arg)
                 (set! tabs-to-spaces #f)
                 (loop args))
                ((string=? "-n" arg)
                 (set! use-dir-locals #f)
                 (loop args))
                ((string=? "-p" arg)
                 (set! space-closing-parens #t)
                 (loop args))
                ((string=? "-N" arg)
                 (cond ((null? args) (usage 1))
                       ((string->number (car args)) =>
                        (lambda (n)
                          (set! start-align n)
                          (loop (cdr args))))
                       (else (usage 1))))
                ((string=? "-T" arg)
                 (cond ((null? args) (usage 1))
                       ((string->number (car args)) =>
                        (lambda (n)
                          (set! tab-stops n)
                          (loop (cdr args))))
                       (else (usage 1))))
                ((string=? "-I" arg)
                 (when (or (null? args) (null? (cdr args)))
                   (usage 1))
                 (put! (string->symbol (car args))
                       'scheme-indent
                       (or (string->number (cadr args))
                           (usage 1)))
                 (loop (cddr args)))
                ((string=? "-d" arg)
                 (set! detect-data #f)
                 (loop args))
                ((string=? "-l" arg)
                 (set! leader-restart #t)
                 (loop args))
                ((or (string=? "-h" arg)
                     (string=? "-help" arg)
                     (string=? "--help" arg))
                 (usage 0))
                ((and (> (string-length arg) 1)
                      (char=? #\- (string-ref arg 0))
                      (null? (lset-difference char=?
                                              (cdr (string->list arg))
                                              '(#\h #\t #\p #\n #\d #\l))))
                 (loop (append (map (lambda (c) (string #\- c))
                                 (cdr (string->list arg)))
                               args)))
                (else (usage 1)))))))

(define (usage code)
  (display #<<EOF
usage: scheme-indent [-h] [-t] [-p] [-l] [-d] [-n] [-N COLUMN] [-T NUM] [-I SYM NUM] ...
  -h -help --help    show this message
  -t                 do not convert tabs to spaces
  -p                 separate closing parens by space
  -T NUM             set tab-width
  -I SYMBOL NUM      set indentation for form
  -N COLUMN          start with alignment at given COLUMN (0-based)
  -n                 ignore .dir-locals.el
  -d                 disable special treatment of (quasi)quoted data
  -l                 restart indentation on non-whitespace at first column
EOF
           (current-error-port))
  (exit code))

(main)