(module loc-wrapper mzscheme
  (require (lib "kw.ss")
           (lib "etc.ss")
           "term.ss")
  
  (provide to-loc-wrapper
           to-loc-wrapper/sc
           (struct loc-wrapper (e line line-span column span last-column))
           (struct grouper (content))
           (struct unq-pict (arg))
           (struct term-pict (arg)))
  
  (define (build-loc-wrapper e line column span)
    (make-loc-wrapper e line #f column span #f))
  
  (define-struct loc-wrapper (e line line-span column span last-column) (make-inspector))

  ;; a grouper puts together multiple tokens that shouldn't have parens around them
  ;; used by hole-proc or others
  ;; content : (cons X (listof X))
  (define-struct grouper (content) (make-inspector))

  (define-struct hole-pict (ctxt exp) (make-inspector))
  (define-struct unq-pict (arg) (make-inspector))
  (define-struct term-pict (arg) (make-inspector))
  
  (define-syntax-set (to-loc-wrapper to-loc-wrapper/sc)
    (define (process-arg stx)
      (syntax-case* stx (unquote in-hole side-condition) (λ (x y) (eq? (syntax-e x) (syntax-e y)))
        [,a
          #`(build-loc-wrapper
             (make-unq-pict #,(process-arg/sc #'a))
             #,(syntax-line stx) 
             #,(syntax-column stx)
             #,(syntax-span stx))]
        [(in-hole a b) #`(build-loc-wrapper (make-hole-pict #,(process-arg #'a)
                                                            #,(process-arg #'b))
                                            #,(syntax-line stx) 
                                            #,(syntax-column stx)
                                            #,(syntax-span stx))]
        
        [(a ...)
         #`(build-loc-wrapper 
            (list #,@(map process-arg (syntax->list (syntax (a ...)))))
            #,(syntax-line stx) 
            #,(syntax-column stx)
            #,(syntax-span stx))]
        [(a b ... . c)
         #`(build-loc-wrapper 
            (append (list #,@(map process-arg (syntax->list (syntax (a b ...)))))
                    #,(process-arg #'c))
            #,(syntax-line stx) 
            #,(syntax-column stx)
            #,(syntax-span stx))]
        [x
         (identifier? #'x)
         #`(build-loc-wrapper 
            'x
            #,(syntax-line stx) 
            #,(syntax-column stx)
            #,(syntax-span stx))]
        [x 
         #`(build-loc-wrapper 
            (syntax-e #'x)
            #,(syntax-line stx) 
            #,(syntax-column stx)
            #,(syntax-span stx))]))
    
    (define (process-arg/sc stx) 
      (syntax-case stx (term)
        [(term x)
         #`(make-term-pict #,(process-arg #'x))]
        [(a ...)
         #`(build-loc-wrapper 
            (list #,@(map process-arg/sc (syntax->list (syntax (a ...)))))
            #,(syntax-line stx) 
            #,(syntax-column stx)
            #,(syntax-span stx))]
        [(a b ... . c)
         #`(build-loc-wrapper 
            (append (list #,@(map process-arg/sc (syntax->list (syntax (a b ...)))))
                    #,(process-arg/sc #'c))
            #,(syntax-line stx) 
            #,(syntax-column stx)
            #,(syntax-span stx))]
        [x
         (identifier? #'x)
         #`(build-loc-wrapper 
            'x
            #,(syntax-line stx) 
            #,(syntax-column stx)
            #,(syntax-span stx))]
        [x 
         #`(build-loc-wrapper 
            (syntax-e #'x)
            #,(syntax-line stx) 
            #,(syntax-column stx)
            #,(syntax-span stx))]))

    (define (to-loc-wrapper/proc stx)
      (syntax-case stx ()
        [(_ stx)
         #`(lambda/kw (#:key hole-pict)
             (process-holes
              hole-pict
              (add-line-spans #,(process-arg #'stx))))]))
    (define (to-loc-wrapper/sc/proc stx)
      (syntax-case stx ()
        [(_ stx)
         (with-syntax ([stx (datum->syntax-object #f (list 'unquote #'stx) #'stx)])
           #`(lambda/kw (#:key hole-pict)
               (process-holes
                hole-pict
                (add-line-spans #,(process-arg #'stx)))))])))

  (define (process-holes hole-pict exp)
    (define (main-loop exp)
      (cond
        [(pair? exp)
         (cons (main-loop (car exp))
               (main-loop (cdr exp)))]
        [(loc-wrapper? exp)
         (let ([e (loc-wrapper-e exp)]
               [line (loc-wrapper-line exp)]
               [line-span (loc-wrapper-line-span exp)]
               [column (loc-wrapper-column exp)]
               [span (loc-wrapper-span exp)]
               [last-column (loc-wrapper-last-column exp)])
           (make-loc-wrapper 
            (cond
              [(hole-pict? e)
               (hole-pict (main-loop (hole-pict-ctxt e))
                          (main-loop (hole-pict-exp e))
                          line
                          line-span
                          column
                          span
                          last-column)]
              [(unq-pict? e)
               (make-unq-pict (strip-loop (unq-pict-arg e)))]
              [else (main-loop e)])
            line
            line-span
            column
            span
            last-column))]
        [else exp]))
    (define (strip-loop exp)
      (cond
        [(pair? exp)
         (cons (strip-loop (car exp))
               (strip-loop (cdr exp)))]
        [(loc-wrapper? exp) (strip-loop (loc-wrapper-e exp))]
        [(term-pict? exp) (make-term-pict (main-loop (term-pict-arg exp)))]
        [else exp]))
    
    (main-loop exp))
  
  (define (add-line-spans lw)
    (define (add-spans/lw lw)
      (let-values ([(last-line last-column) (add-spans/obj (loc-wrapper-line lw) 
                                                           (loc-wrapper-column lw)
                                                           (loc-wrapper-e lw))])
        (set-loc-wrapper-line-span! lw (- last-line (loc-wrapper-line lw)))
        (set-loc-wrapper-last-column! lw last-column)
        (values last-line
                last-column)))
    (define (add-spans/obj line col e)
      (cond
        [(null? e) (values line (+ col 2))]
        [(list? e)
         (let loop ([fst (car e)]
                    [rst (cdr e)])
           (cond
             [(null? rst)
              (let-values ([(last-line last-col) (add-spans/obj line col fst)])
                (values last-line
                        (+ last-col 1)))]
             [else
              (add-spans/obj line col fst)
              (loop (car rst) (cdr rst))]))]
        [(pair? e)
         (add-spans/obj line (car e))
         (let-values ([(last-line last-col) (add-spans/obj line (cdr e))])
           (values last-line
                   (+ last-col 1)))]
        [(loc-wrapper? e) (add-spans/lw e)]
        [(grouper? e) 
         (let loop ([fst (car (grouper-content e))]
                    [rst (cdr (grouper-content e))])
           (cond
             [(null? rst)
              (add-spans/obj line col fst)]
             [else
              (add-spans/obj line col fst)
              (loop (car rst) (cdr rst))]))]
        [(hole-pict? e) 
         (add-spans/obj line col (hole-pict-ctxt e))
         (add-spans/obj line col (hole-pict-exp e))]
        [(term-pict? e)
         (add-spans/obj line col (term-pict-arg e))]
        [(unq-pict? e)
         (add-spans/obj line col (unq-pict-arg e))]
        [else (values line 
                      (+ col (string-length (format "~s" e))))]))
    (add-spans/lw lw)
    lw))