;; test cases

#;(in-hole a ,b)

(module core-layout mzscheme
  (require "loc-wrapper.ss"
           (lib "list.ss")
           (lib "utils.ss" "texpict")
           (lib "mrpict.ss" "texpict")
           (lib "etc.ss")
           (lib "mred.ss" "mred")
           (lib "struct.ss"))
  
  (provide find-enclosing-loc-wrapper
           lw->pict
           basic-text
           metafunction-text
           default-style
           label-style
           non-terminal-style
           non-terminal-subscript-style
           label-font-size
           default-font-size
           metafunction-font-size
           non-terminal
           literal-style
           metafunction-style
           open-white-square-bracket
           close-white-square-bracket
           just-before
           just-after
           with-unquote-rewriter
           STIX?
           
           ;; for test suite
           build-lines
           (struct token (column span))
           (struct string-token (string style))
           (struct pict-token (pict))
           (struct spacer-token ()))

  
  (define STIX? #f)
  
  ;; atomic-rewrite-table : (parameter (listof (list symbol (union string pict))))
  (define atomic-rewrite-table 
    (make-parameter 
     `((... ,(if STIX?
                 (basic-text "\u22ef" (default-style))
                 "..."))
       (hole "[]"))))
  
  ;; compound-rewrite-table : (listof lw) -> (listof (union lw pict string))
  (define compound-rewrite-table 
    (make-parameter 
     `((in-hole ,(λ (args)
                   (let ([context (list-ref args 2)]
                         [thing-in-hole (list-ref args 3)])
                     (list context
                           (if (= (loc-wrapper-line context)
                                  (loc-wrapper-line thing-in-hole))
                               (build-loc-wrapper "["
                                                  (loc-wrapper-line thing-in-hole)
                                                  0
                                                  (+ (loc-wrapper-column context)
                                                     (loc-wrapper-column-span context))
                                                  (- (loc-wrapper-column thing-in-hole)
                                                     (+ (loc-wrapper-column context)
                                                        (loc-wrapper-column-span context))))
                               (build-loc-wrapper "["
                                                  (loc-wrapper-line thing-in-hole)
                                                  0
                                                  (loc-wrapper-column thing-in-hole)
                                                  0))
                           thing-in-hole
                           (build-loc-wrapper "]"
                                              (+ (loc-wrapper-line thing-in-hole)
                                                 (loc-wrapper-line-span thing-in-hole))
                                              0
                                              (+ (loc-wrapper-column thing-in-hole)
                                                 (loc-wrapper-column-span thing-in-hole))
                                              0))))))))
  
  (define-syntax (with-unquote-rewriter stx)
    (syntax-case stx ()
      [(_ transformer e)
       #'(parameterize ([current-unquote-rewriter transformer])
           e)]))
  
  (define current-unquote-rewriter (make-parameter values))
  
  ;; token = string-token | spacer-token | pict-token | align-token
  
  (define-struct token (column span) (make-inspector))
  
  ;; string : string
  ;; style : valid third argument to mrpict.ss's `text' function
  (define-struct (string-token token) (string style) (make-inspector))
  
  ;; width : number
  ;; pict : pict
  (define-struct (pict-token token) (pict) (make-inspector))
  
  ;; spacer : number
  (define-struct (spacer-token token) () (make-inspector))

  ;; pict : pict
  ;; this token always appears at the beginning of a line and its width
  ;; is the x-coordinate of the pict inside itself (which must appear on
  ;; an earlier line)
  (define-struct align-token (pict) (make-inspector))

  (define (lw->pict nts lw)
    (lines->pict 
     (setup-lines 
      (build-lines 
       nts
       (apply-rewrites lw)))))
  
  (define (apply-rewrites orig-lw)
    (define (ar/lw lw)
      (cond
        [(eq? lw 'spring) lw]
        [else
         (let ([rewritten (if (loc-wrapper-unq? lw)
                              ((current-unquote-rewriter) lw)
                              lw)])
           (if (eq? rewritten lw)
               (copy-struct loc-wrapper
                            lw
                            [loc-wrapper-e (ar/e (loc-wrapper-e lw)
                                                 (loc-wrapper-line lw)
                                                 (loc-wrapper-column lw))])
               (ar/lw rewritten)))]))
    (define (ar/e e line col)
      (cond
        [(and (symbol? e) (assoc e (atomic-rewrite-table)))
         =>
         (λ (m)
           (when (eq? (cadr m) e)
             (error 'apply-rewrites "rewritten version of ~s is still ~s" e e))
           (cadr m))]
        [(symbol? e) e]
        [(string? e) e]
        [(pict? e) e]
        [(and (loc-wrapper? (cadr e))
              (assoc (loc-wrapper-e (cadr e)) (compound-rewrite-table)))
         =>
         (λ (m)
           (let ([rewritten ((cadr m) e)])
             (when (and (pair? rewritten)
                        (pair? (cdr rewritten))
                        (eq? (cadr rewritten) 
                             (cadr e)))
               (error 'apply-rewrites "rewritten version still has symbol of the same name as original: ~s" 
                      (cadr rewritten)))
             (adjust-spacing rewritten e line col)))]
        [(and (loc-wrapper? (cadr e))
              (loc-wrapper-metafunction-name (cadr e)))
         (map ar/lw (rewrite-metafunction-app (loc-wrapper-metafunction-name (cadr e)) e))]
        [else
         (map ar/lw e)]))
    (ar/lw orig-lw))
  
  (define (rewrite-metafunction-app kind lst)
    (case kind
      [(single-arg) 
       (list* (just-before (metafunction-text (symbol->string (loc-wrapper-e (cadr lst)))) 
                           (car lst))
              (just-before (open-white-square-bracket)
                           (car lst))
              'spring
              (let loop ([lst (cddr lst)])
                (cond
                  [(null? lst) null]
                  [(null? (cdr lst))
                   (let ([last (car lst)])
                     (list (just-before (close-white-square-bracket) last)))]
                  [else (cons (car lst) (loop (cdr lst)))])))]
      [(multi-arg)
       (list* (just-before (metafunction-text (symbol->string (loc-wrapper-e (cadr lst)))) 
                           (car lst))
              (just-before (open-white-square-bracket)
                           (car lst))
              'spring
              (let loop ([lst (cddr lst)])
                (cond
                  [(null? lst) null]
                  [(null? (cdr lst))
                   (let ([last (car lst)])
                     (list (just-before (close-white-square-bracket) last)))]
                  [(null? (cddr lst))
                   (cons (car lst) (loop (cdr lst)))]
                  [else (list* (car lst) 
                               (just-after (basic-text ", " (default-style)) (car lst))
                               'spring
                               (loop (cdr lst)))])))]))
  
  (define (just-before what lw)
    (build-loc-wrapper what
                       (loc-wrapper-line lw)
                       0
                       (loc-wrapper-column lw)
                       0))
  
  (define (just-after what lw)
    (build-loc-wrapper what
                       (+ (loc-wrapper-line lw) (loc-wrapper-line-span lw))
                       0
                       (+ (loc-wrapper-column lw) (loc-wrapper-column-span lw))
                       0))
  
  ;; adjust-spacing : (listof (union string loc-wrapper) (listof loc-wrapper) number -> (listof loc-wrapper)
  ;; builds loc-wrappers out of the strings in the rewrittens,
  ;; using the originals around the string in order to find column numbers for the strings
  ;; NB: there is still an issue with this code -- if the rewrite drops stuff that
  ;;     appears at the end of the sequence, blank space will still appear in the final output ...
  (define (adjust-spacing in-rewrittens in-originals init-line init-column)
    (let loop ([rewrittens in-rewrittens]
               [originals in-originals]
               [line init-line]
               [column init-column])
    (cond
      [(null? rewrittens) 
       null]
      [(null? originals)
       (map (λ (rw) (if (loc-wrapper? rw)
                        rw
                        (build-loc-wrapper rw line 0 column 0)))
            rewrittens)]
      [else 
       (let ([orig (car originals)]
             [rewritten (car rewrittens)])
         (cond
           [(loc-wrapper? rewritten)
            (let ([new-line (+ (loc-wrapper-line rewritten)
                               (loc-wrapper-line-span rewritten))]
                  [new-col (+ (loc-wrapper-column orig)
                              (loc-wrapper-column-span rewritten))])
              (cond
                [(memq rewritten originals)
                 ;; if this rewritten is in the list of originals, drop all the things inbetween
                 ;; (it is illegal to come back to them anyways) and continue with the column
                 ;; and line set to the place after this loc wrapper
                 ;; also, make a blank pict that takes up the space on the last line
                 (let ([rw-line (loc-wrapper-line rewritten)])
                   (let d-loop ([originals originals]
                                [first-column #f])
                     (cond
                       [(eq? (car originals) rewritten)
                        (if first-column
                            (list* (build-loc-wrapper (blank) 
                                                      rw-line
                                                      0
                                                      first-column
                                                      (- (loc-wrapper-column rewritten)
                                                         first-column))
                                   rewritten
                                   (loop (cdr rewrittens) (cdr originals) new-line new-col))
                            (cons rewritten
                                  (loop (cdr rewrittens) (cdr originals) new-line new-col)))]
                       [(and (not first-column) (= rw-line (loc-wrapper-line (car originals))))
                        (d-loop (cdr originals)
                                (loc-wrapper-column (car originals)))]
                       [else 
                        (d-loop (cdr originals)
                                first-column)])))]
                [(memq rewritten in-originals)
                 (error 'adjust-spacing "found an out of order loc-wrapper ~s" rewritten)]
                [else
                 (unless (<= line (loc-wrapper-line orig))
                   (error 'adjust-spacing
                          "new loc-wrapper's line is earlier than a loc-wrapper's line that appears earlier in the list: ~s"
                          rewritten))
                 (unless (<= column (loc-wrapper-column orig))
                   (error 'adjust-spacing
                          "new loc-wrapper's column is earlier than a loc-wrapper's column that appears earlier in the list: ~s"
                          rewritten))
                 (cons rewritten (loop (cdr rewrittens) originals new-line new-col))]))]
           [else
            (cons (build-loc-wrapper rewritten line 0 column 0)
                  (loop (cdr rewrittens) originals line column))]))])))
  
  (define (build-lines all-nts lw)
    (define initial-column (loc-wrapper-column lw))
    (define initial-line (loc-wrapper-line lw))
    (define current-line (loc-wrapper-line lw))
    (define current-column (loc-wrapper-column lw))
    (define last-token-spring? #f)
    (define tokens '())
    (define lines '())
    (define (eject line col span atom unquoted?)
      (unless (= current-line line)
        ;; make new lines
        (for-each 
         (λ (x) 
           (set! lines (cons (reverse! tokens) lines))
           (set! tokens '()))
         (build-list (max 0 (- line current-line)) (λ (x) 'whatever)))
        
        (set! tokens (cons (make-spacer-token 0 (- col initial-column))
                           tokens))
        
        (set! current-line line)
        (set! current-column col))
      (when (< current-column col)
        (let ([space-span (- col current-column)])
          (set! tokens (cons (make-blank-space-token unquoted?
                                                     (- current-column initial-column)
                                                     space-span)
                             tokens))))
      (set! last-token-spring? #f)
      (set! tokens (append 
                    (reverse
                     (atom->tokens (- col initial-column) span atom all-nts unquoted?))
                    tokens))
      (set! current-column (+ col span)))
    
    (define (make-blank-space-token unquoted? col span)
      (if last-token-spring?
          (make-pict-token col span (blank))
          (let ([str (apply string (build-list span (λ (x) #\space)))])
            (if unquoted?
                (make-pict-token col span (pink-background (text str 'modern (default-font-size))))
                (make-string-token col span str (default-style))))))
    
    (define (handle-loc-wrapped lw last-line last-column last-span)
      (cond
        [(eq? lw 'spring)
         (set! last-token-spring? #t)]
        [else
         (handle-object (loc-wrapper-e lw)
                        (loc-wrapper-line lw)
                        (loc-wrapper-column lw)
                        (loc-wrapper-column-span lw)
                        (loc-wrapper-unq? lw))]))
    
    (define (handle-object obj line col span unquoted?)
      (cond
        [(symbol? obj) (eject line col span obj unquoted?)]
        [(string? obj) (eject line col span obj unquoted?)]
        [(pict? obj) (eject line col span obj unquoted?)]
        [(not obj) (eject line col span (blank) unquoted?)]
        [else
         (for-each (λ (x) (handle-loc-wrapped x line col span))
                   obj)]))
    
    (handle-loc-wrapped lw 0 0 0)
    (set! lines (cons (reverse! tokens) lines)) ;; handle last line ejection
    lines)
  
  ;; setup-lines : (listof (listof token)) -> (listof (listof token))
  ;; removes the spacer tokens from the beginning of lines, replacing them with align tokens
  ;; expects the lines to be in reverse order
  (define (setup-lines lines)
    (let loop ([lines lines])
      (cond
        [(null? lines) null]
        [else 
         (let ([line (car lines)]
               [rst (cdr lines)])
           (if (null? line)
               (cons line (loop (cdr lines)))
               (if (spacer-token? (car line))
                   (let ([pict (blank)])
                     (if (andmap null? rst)
                         (cons (cdr line) (loop rst))
                         (let ([rst (split-out (token-span (car line))
                                               pict
                                               rst)])
                           (cons (cons (make-align-token pict) (cdr line))
                                 (loop rst)))))
                   (cons line (loop (cdr lines))))))])))
  
  (define (split-out col pict lines)
    (let ([new-token (make-pict-token col 0 pict)])
      (let loop ([lines lines])
        (cond
          [(null? lines)
           ;; this case can happen when the line in question is to the left of all other lines
           (error 'exchange-spacer "could not find matching line")]
          [else (let ([line (car lines)])
                  (if (null? line)
                      (cons line (loop (cdr lines)))
                      (let ([spacer (car line)])
                        (cond
                          [(not (spacer-token? spacer))
                           (cons (insert-new-token col new-token (token-column spacer) (car lines))
                                 (cdr lines))]
                          [(= (token-span spacer)
                              col)
                           (cons (list* spacer new-token (cdr line))
                                 (cdr lines))]
                          [(> (token-span spacer)
                              col)
                           (cons line (loop (cdr lines)))]
                          [(< (token-span spacer)
                              col)
                           (cons (insert-new-token col new-token (token-column spacer) (car lines))
                                 (cdr lines))]))))]))))
                  
  (define (insert-new-token column-to-insert new-token init-width line)
    (let loop ([line line]
               [column 0])
      (cond
        [(null? line)
         (list new-token)]
        [else
         (let ([tok (car line)])
           (unless (token? tok)
             (error 'insert-new-token "ack ~s" tok))
           (cond
             [(<= column-to-insert (token-column tok))
              (cons new-token line)]
             [(< (token-column tok)
                 column-to-insert
                 (+ (token-column tok) (token-span tok)))
              (append (split-token (- column-to-insert (token-column tok)) tok new-token)
                      (cdr line))]
             [(= column-to-insert (+ (token-column tok) (token-span tok)))
              (list* (car line) new-token (cdr line))]
             [else 
              (cons (car line)
                    (loop (cdr line)
                          (+ (token-column tok) (token-span tok))))]))])))
               
  (define (split-token offset tok new-token)
    (cond
      [(string-token? tok)
       (list (make-string-token (token-column tok)
                                offset
                                (substring (string-token-string tok)
                                           0 offset)
                                (string-token-style tok))
             new-token
             (make-string-token (+ (token-column tok) offset)
                                (- (token-span tok) offset)
                                (substring (string-token-string tok)
                                           offset 
                                           (string-length (string-token-string tok)))
                                (string-token-style tok)))]
      [(pict-token? tok)
       (list new-token)]))
  
  ;; lines->pict : (listof (listof token)) -> pict
  ;; expects the lines to be in order from bottom to top
  (define (lines->pict lines)
    (let loop ([lines lines])
      (cond
        [(null? lines) (blank)]
        [(null? (cdr lines))
         (handle-single-line (car lines) (blank))]
        [else
         (let ([rst (loop (cdr lines))])
           (vl-append rst (handle-single-line (car lines) rst)))])))
  
  (define (handle-single-line line rst)
    (cond
      [(null? line) 
       (let ([h (pict-height (token->pict (make-string-token 0 0 "x" (default-style))))])
         (blank 0 h))]
      [else
       (if (align-token? (car line))
           (let-values ([(x y) (lt-find rst (align-token-pict (car line)))])
             (apply hbl-append 
                    (blank x 0)
                    (map token->pict (cdr line))))
           (apply hbl-append (map token->pict line)))]))
  
  ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  ;;
  ;;  font specs
  ;;
  
  
  (define (token->pict tok)
    (cond
      [(string-token? tok)
       (basic-text (string-token-string tok) (string-token-style tok))]
      [(pict-token? tok) (pict-token-pict tok)]
      [else (error 'token->pict "~s" tok)]))
  
    
  (define (atom->tokens col span atom all-nts unquoted?)
    (cond
      [(pict? atom)
       (list (make-pict-token col span atom))]
      [unquoted?
       (list (make-pict-token col span 
                              (pink-background 
                               (text (if (string? atom) atom (format "~a" atom))
                                     'modern
                                     (default-font-size)))))]
      [(and (symbol? atom)
            (regexp-match #rx"^([^_]*)_(.*)$" (symbol->string atom)))
       =>
       (λ (m)
         (let* ([first-part (cadr m)]
                [second-part (caddr m)]
                [first-span (- span (string-length first-part))])
           (list 
            (make-string-token col
                               first-span
                               first-part
                               (non-terminal-style))
            (make-string-token (+ col first-span) 
                               (- span first-span)
                               second-part
                               (non-terminal-subscript-style)))))]
      [(memq atom all-nts)
       (list (make-string-token col span (format "~s" atom) (non-terminal-style)))]
      [(symbol? atom)
       (list (make-string-token col span (symbol->string atom) (literal-style)))]
      [(string? atom)
       (list (make-string-token col span atom (literal-style)))]
      [else (error 'atom->tokens "unk ~s" atom)]))
  
  (define (il-memq sym s)
    (and (pair? s)
         (or (eq? sym (car s))
             (il-memq sym (cdr s)))))
  (define (il-remq sym s)
    (if (pair? s)
        (if (eq? sym (car s))
            (cdr s)
            (cons (car s) (il-remq sym (cdr s))))
        s))
  
(define (pick-font lst fallback)
  (let ([fl (get-face-list 'all)])
    (let loop ([lst lst])
      (cond
        [(null? lst) fallback]
        [else (if (member (car lst) fl)
                  (car lst)
                  (loop (cdr lst)))]))))

  (define (basic-text str style) 
    (if (il-memq 'caps style)
        (caps-text str (il-remq 'caps style) (default-font-size))
        (text str style (default-font-size))))
  (define (non-terminal str) (text str (non-terminal-style) (default-font-size)))
  (define (unksc str) (pink-background (text str 'modern (default-font-size))))
  (define non-terminal-style (make-parameter '(italic . roman)))
  (define non-terminal-subscript-style (make-parameter `(subscript . ,(non-terminal-style))))
  (define default-style (make-parameter 'roman))
  (define metafunction-style (make-parameter 'swiss))
  (define (metafunction-text str) (text str (metafunction-style) (metafunction-font-size)))
  (define literal-style (make-parameter 'swiss))
  (define label-style (make-parameter 'swiss))
  (define default-font-size (make-parameter 14))
  (define metafunction-font-size (make-parameter (default-font-size)))
  (define label-font-size (make-parameter 14))
  
  (define (open-white-square-bracket) (white-bracket "["))
  
  (define (close-white-square-bracket) (white-bracket "]"))
  
  (define (white-bracket str)
    (hbl-append (basic-text str (default-style))
                (inset (basic-text str (default-style)) -2 0 0 0)))
  
  (define (pink-background p)
    (refocus
     (cc-superimpose 
      (colorize (filled-rectangle (pict-width p)
                                  (pict-height p))
                "pink")
      p)
     p))
  
  (define (add-between i l)
    (cond
      [(null? l) l]
      [else 
       (cons (car l)
             (apply append 
                    (map (λ (x) (list i x)) (cdr l))))]))
  
  
  ;; for use
  (define (find-enclosing-loc-wrapper lws)
    (let* ([first-line (apply min (map loc-wrapper-line lws))]
           [last-line (apply min (map (λ (x) (+ (loc-wrapper-line x) (loc-wrapper-line-span x))) lws))]
           [last-line-lws (find-lws-with-matching-last-line lws last-line)]
           [last-column (apply max (map (λ (x) (+ (loc-wrapper-column x) (loc-wrapper-column-span x))) last-line-lws))]
           [first-column (apply min (map loc-wrapper-column last-line-lws))])
      (build-loc-wrapper
       lws
       first-line
       (- last-line first-line)
       first-column
       (- last-column first-column))))
  
  (define (find-lws-with-matching-last-line in-lws line)
    (define lws '())
    
    (define (find/lw lw)
      (cond
        [(eq? lw 'spring) (void)]
        [(loc-wrapper? lw)
         (when (= line (+ (loc-wrapper-line lw) (loc-wrapper-line-span lw)))
           (set! lws (cons lw lws))
           (find/e (loc-wrapper-e lw)))]))
    
    (define (find/e e)
      (cond
        [(symbol? e) (void)]
        [(string? e) (void)]
        [(pict? e) (void)]
        [else (for-each find/lw e)]))
    
    (find/e in-lws)
    lws)

  )