;; test cases:
#;()
#;(a)
#;(a
   b)
#;(a
   (b c)
   d)
#;(abcdefghijkl
   c)
#;((a b)
   c)
#;([{}])
#;,@,'(x)
#;(term a)

(module loc-wrapper mzscheme
  (require (lib "kw.ss")
           (lib "etc.ss")
           "term.ss"
           (lib "contract.ss"))
  (require-for-syntax "term-fn.ss")
  
  (define (init-loc-wrapper e line column quoted?)
    (make-loc-wrapper e line #f column #f (not quoted?) #f))
  
  ;; lw = (union 'spring loc-wrapper)
  
  ;; e : (union string symbol #f (listof lw))
  ;; line, line-span, column, column-span : number
  (define-struct loc-wrapper (e line line-span column column-span unq? metafunction-name) (make-inspector))
  
  (define (lw? x) (or (eq? 'spring x) (loc-wrapper? x)))
  
  ;; build-loc-wrapper is designed for external consumption
  (define (build-loc-wrapper e line line-span column column-span)
    (make-loc-wrapper e line line-span column column-span #f #f))
  
  (define-syntax-set (to-loc-wrapper to-loc-wrapper/uq)
    (define (process-arg stx quote-depth)
      (define quoted? (quote-depth . > . 0))
      (define (reader-shorthand arg qd-delta)
        #`(init-loc-wrapper 
           (list (init-loc-wrapper ""
                                    #,(syntax-line stx)
                                    #,(syntax-column stx)
                                    #,quoted?)
                 'spring
                 #,(process-arg arg (+ quote-depth qd-delta)))
           #,(syntax-line stx) 
           #,(syntax-column stx)
           #,quoted?))
      (let-values ([(op cl)
                    (if (syntax? stx)
                        (case (syntax-property stx 'paren-shape)
                          [(#\{) (values "{" "}")]
                          [(#\[) (values "[" "]")]
                          [else (values "(" ")")])
                        (values #f #f))])
        (syntax-case* stx (unquote quote unquote-splicing term) (λ (x y) (eq? (syntax-e x) (syntax-e y)))
          ['a (reader-shorthand #'a +1)]
          [,a (reader-shorthand #'a -1)]
          [,@a (reader-shorthand #'a -1)]
          [(term a) 
           #`(init-loc-wrapper 
              (list (init-loc-wrapper "" #,(syntax-line stx) #,(syntax-column stx) #,quoted?)
                    'spring
                    #,(process-arg (cadr (syntax->list stx)) (+ quote-depth 1))
                    'spring)
              #,(syntax-line stx) 
              #,(syntax-column stx)
              #,quoted?)]
          [(a ...)
           #`(init-loc-wrapper 
              (list (init-loc-wrapper #,op #,(syntax-line stx) #,(syntax-column stx) #,quoted?)
                    #,@(map (λ (x) (process-arg x quote-depth)) (syntax->list (syntax (a ...))))
                    (init-loc-wrapper #,cl #f #f #,quoted?))
              #,(syntax-line stx) 
              #,(syntax-column stx)
              #,quoted?)]
          [(a b ... . c)
           #`(init-loc-wrapper 
              (list (init-loc-wrapper #,op #,(syntax-line stx) #,(syntax-column stx) #,quoted?)
                    #,@(map (λ (x) (process-arg x quote-depth)) (syntax->list (syntax (a b ...))))
                    (init-loc-wrapper #," . " #f #f #,quoted?)
                    #,(process-arg #'c quote-depth)
                    (init-loc-wrapper #,cl #f #f #,quoted?))
              #,(syntax-line stx) 
              #,(syntax-column stx)
              #,quoted?)]
          [x 
           (and (identifier? #'x)
                (term-fn? (syntax-local-value #'x (λ () #f))))
           #`(make-loc-wrapper
              '#,(syntax-e #'x)
              #,(syntax-line stx) 
              #f
              #,(syntax-column stx)
              #f
              #f
              #,(if (term-fn-multi-arg? (syntax-local-value #'x))
                    #''multi-arg
                    #''single-arg))]
          [x 
           (identifier? #'x)
           #`(init-loc-wrapper 
              '#,(syntax-e #'x)
              #,(syntax-line stx) 
              #,(syntax-column stx)
              #,quoted?)]
          [x 
           #`(init-loc-wrapper 
              #,(format "~s" (syntax-e #'x))
              #,(syntax-line stx) 
              #,(syntax-column stx)
              #,quoted?)])))
    
    (define (to-loc-wrapper/proc stx)
      (syntax-case stx ()
        [(_ stx)
         #`(add-spans #,(process-arg #'stx 1))]))
    (define (to-loc-wrapper/uq/proc stx)
      (syntax-case stx ()
        [(_ stx)
         #`(add-spans #,(process-arg #'stx 0))])))

  (define (add-spans lw)
    (define (add-spans/lw lw line col)
      (cond
        [(eq? lw 'spring) (values line col)]
        [else
         (let ([start-line (or (loc-wrapper-line lw) line)]
               [start-column (or (loc-wrapper-column lw) col)])
           (let-values ([(last-line last-column) (add-spans/obj (loc-wrapper-e lw) start-line start-column)])
             (unless (loc-wrapper-line lw)
               (set-loc-wrapper-line! lw line))
             (unless (loc-wrapper-column lw)
               (set-loc-wrapper-column! lw col))
             (set-loc-wrapper-line-span! lw (- last-line start-line))
             (set-loc-wrapper-column-span! lw (- last-column start-column))
             (values last-line last-column)))]))
    (define (add-spans/obj e line col)
      (cond
        [(string? e) 
         (values line (+ col (string-length e)))]
        [(symbol? e)
         (values line (+ col (string-length (symbol->string e))))]
        [(not e) (values line col)]
        [else 
         (let loop ([lws e]
                    [line line]
                    [current-col col])
           (cond
             [(null? lws) (values line current-col)]
             [else 
              (let-values ([(last-line last-column) (add-spans/lw (car lws) line current-col)])
                (loop (cdr lws)
                      last-line
                      last-column))]))]))
    (add-spans/lw lw #f #f)
    lw)
  
  (define pnum (and/c number? (or/c zero? positive?)))
  
  (provide/contract
   (struct loc-wrapper ((e any/c)
                        (line pnum)
                        (line-span pnum)
                        (column pnum)
                        (column-span pnum)
                        (unq? boolean?)
                        (metafunction-name (or/c (symbols 'multi-arg 'single-arg) false/c))))
   [build-loc-wrapper (-> any/c pnum pnum pnum pnum loc-wrapper?)])
  
  (provide to-loc-wrapper
           to-loc-wrapper/uq))