(module core-layout mzscheme
  (require "loc-wrapper.ss"
           "arrow.ss"
           (lib "utils.ss" "texpict")
           (lib "mrpict.ss" "texpict")
           (lib "etc.ss")
           (lib "class.ss")
           (lib "mred.ss" "mred")
           (lib "struct.ss"))
  
  (provide lw->pict
           basic-text
           default-style
           label-style
           label-font-size
           default-font-size
           non-terminal
           set-literal-style!
           set-metafunction-style!)
  
  (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 thing-in-hole)
                                  (loc-wrapper-line context))
                               (make-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))))
                               (make-loc-wrapper "["
                                                 (loc-wrapper-line thing-in-hole)
                                                 0
                                                 (loc-wrapper-column thing-in-hole)
                                                 0))
                           thing-in-hole
                           (make-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))))))))
  
  ;; 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
       (push-down-quo
        (apply-rewrites lw))))))
  
  (define (push-down-quo lw)
    (define (pd/lw lw depth)
      (cond
        [(quo? lw) (pd/lw (quo-e lw) (+ depth 1))]
        [(unq? lw) (pd/lw (unq-e lw) (- depth 1))]
        [else
         (copy-struct loc-wrapper
                      lw
                      [loc-wrapper-e (pd/e (loc-wrapper-e lw) depth)])]))
    (define (pd/e e depth)
      (cond
        [(symbol? e)
         (if (depth . <= . 0)
             (pinkize (symbol->string e))
             e)]
        [(string? e)
         (if (depth . <= . 0)
             (pinkize e)
             e)]
        [(pict? e) e]
        [else (map (λ (x) (pd/lw x depth)) e)]))
    
    (define (pinkize str)
      (pink-background 
       (text str 'modern (default-font-size))))
    
    (pd/lw lw 1))
                                
  
  (define (apply-rewrites orig-lw)
    (define (ar/lw lw)
      (cond
        [(unq? lw) (make-unq (ar/lw (unq-e lw)))]
        [(quo? lw) (make-quo (ar/lw (quo-e lw)))]
        [else
         (copy-struct loc-wrapper
                      lw
                      [loc-wrapper-e (ar/e (loc-wrapper-e lw)
                                           (loc-wrapper-line lw)
                                           (loc-wrapper-column lw))])]))
    (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)))]
        [else
         (map ar/lw e)]))
    (ar/lw orig-lw))

  ;; 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
                        (make-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* (make-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 (make-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 tokens '())
    (define lines '())
    (define (eject line col span atom)
      (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-string-token (- current-column initial-column)
                                                space-span
                                                (apply string (build-list space-span (λ (x) #\space)))
                                                default-style)
                             tokens))))
      (set! tokens (append 
                    (reverse
                     (atom->tokens (- col initial-column) span atom all-nts))
                    tokens))
      (set! current-column (+ col span)))
    
    (define (handle-loc-wrapped lw last-line last-column last-span)
      (handle-object (loc-wrapper-e lw)
                     (loc-wrapper-line lw)
                     (loc-wrapper-column lw)
                     (loc-wrapper-column-span lw)))
    
    (define (handle-object obj line col span)
      (cond
        [(symbol? obj) (eject line col span obj)]
        [(string? obj) (eject line col span obj)]
        [(pict? obj) (eject line col span obj)]
        [(unq? obj) (handle-object (unq-e obj) line col span)]
        [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])
      (cond
        [(null? line)
         (error 'insert-new-token "not yet finished (need to go to a previous line)")]
        [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)))]))])))
               
  (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)
    (cond
      [(pict? atom)
       (list (make-pict-token col span atom))]
      [(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 (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 '(italic . roman))
  (define non-terminal-subscript-style `(subscript . ,non-terminal-style))
  (define default-style 'roman)
  (define metafunction-style 'swiss)
  (define literal-style 'swiss)
  (define label-style 'swiss)
  (define default-font-size (make-parameter 14))
  (define label-font-size (make-parameter 14))
  
  (define (set-literal-style! s) (set! literal-style s))
  (define (set-metafunction-style! s) (set! metafunction-style s))
  
  (require (lib "match.ss"))
  (define (side-condition->pict sc)
    (let loop ([sc sc])
      (match sc
        [(? pict? sc) sc]
        [`(or ,s1 ,s2)
          (hbl-append
           (loop s1)
           (basic-text " or " 'roman)
           (loop s2))]
        [`(or ,s1 ,s2 ...)
          (apply hbl-append
                 (add-between (basic-text ", or " 'roman)
                              (map loop (cons s1 s2))))]
        [else 
         (fprintf (current-error-port) "unknown Scheme code: ~s\n" sc)
         (render-side-condition sc)])))
  
  (define (render-side-condition sc)
    (cond
      [(pict? sc) sc]
      [(null? sc) (unksc "()")]
      [(pair? sc)
       (hbl-append
        (unksc "(")
        (render-side-condition (car sc))
        (if (null? (cdr sc))
            (blank)
            (render-side-condition/list (cdr sc)))
        (unksc ")"))]
      [else
       (unksc (format "~s" sc))]))
  
  ;; render-side-condition : pair -> pict
  (define (render-side-condition/list sc)
    (cond
      [(null? (cdr sc))
       (hbl-append 
        (unksc " ")
        (render-side-condition (car sc)))]
      [(pair? (cdr sc))
       (hbl-append 
        (unksc " ")
        (render-side-condition (car sc))
        (render-side-condition/list (cdr sc)))]
      [else
       (hbl-append 
        (unksc " ")
        (render-side-condition (car sc))
        (unksc " . ")
        (render-side-condition (cdr sc)))]))
  
  
  (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))))]))
  )