#lang scheme/base

(require scribble/manual
         scribble/struct
         scribble/scheme)

(provide code-block
         code-elem)

(define (read-all-syntaxes str read)
  (with-handlers ([void (lambda (e)
                          (fprintf (current-error-port)
                                   "READ ERROR: ~a\nwhile reading:\n~a\n"
                                   (exn-message e) str)
                          '())])
    (parameterize ([current-input-port (open-input-string str)])
      (port-count-lines! (current-input-port))
      (let loop ()
        (let ([stx (read)])
          (if (eof-object? stx) '() (cons stx (loop))))))))

;; collect all identifiers, ignores un-original ids and non-identifiers
(define (get-identifiers x)
  (let loop ([x x] [acc null])
    (cond [(identifier? x) (if (syntax-original? x) (cons x acc) acc)]
          [(syntax? x) 
           (let* ([v (syntax-e x)]
                  [acc
                   (cond
                    [(string? v) 
                     (if (syntax-property x 'scribble)
                         acc
                         (cons x acc))]
                    [(number? v) (cons x acc)]
                    [(boolean? v) (cons x acc)]
                    [else acc])])
             (loop (syntax-e x) acc))]
          [(pair? x) (loop (car x) (loop (cdr x) acc))]
          [(null? x) acc]
          [else acc])))

(define (decorate-id id str)
  (to-element (make-just-context (datum->syntax #f (string->symbol str))
                                 (datum->syntax #'here (syntax-e id)))))

(define (decorate-identifier id str)
  (cond
   [(equal? str "'") str]
   [(identifier? id)
    (decorate-id id str)]
   [(string? (syntax-e id))
    (schemevalfont str)]
   [else
    (to-element id)]))

(define (read-syntax*)
  (parameterize ([read-accept-reader #t]) (read-syntax)))

(define (expr-decorate strs #:reader [reader read-syntax*])
  (let* ([str (apply string-append strs)]
         [ids (sort (get-identifiers (read-all-syntaxes str reader))
                    < #:key syntax-position)]
         [len (string-length str)])
    (let loop ([i 0] [ids ids])
      (if (null? ids)
        (if (= i len) '() (list (substring str i)))
        (let* ([id (car ids)]
               [pos (sub1 (syntax-position id))]
               [span (syntax-span id)])
          (cond [(pos . < . i)
                 (loop i (cdr ids))
                 #;
                 (error 'expr-decorate "nested identifiers found in: ~e at: ~e" str (syntax->datum id))]
                [(pos . > . i)
                 (cons (substring str i pos) (loop pos ids))]
                ;; pos = i
                [else 
                 (cons (decorate-identifier
                        id (substring str pos (+ pos span)))
                       (loop (+ pos span) (cdr ids)))]))))))

(define (split-lines l)
  (let loop ([l l][so-far null])
    (cond
     [(null? l) (if (null? so-far)
                    null
                    (cons (reverse so-far) null))]
     [(equal? (car l) "\n")
      (cons (reverse so-far) (loop (cdr l) null))]
     [(and (string? (car l))
           (regexp-match #rx"(.*)\n(.*)" (car l)))
      => (lambda (m)
           (loop (list* (cadr m)
                        "\n"
                        (caddr m)
                        (cdr l))
                 so-far))]
     [(and (string? (car l))
           (regexp-match #rx"(.*)(#lang [a-z/]+)(.*)" (car l)))
      => (lambda (m)
           (loop (list* (cadr m)
                        (tt (caddr m))
                        (cadddr m)
                        (cdr l))
                 so-far))]
     [(and (string? (car l))
           (regexp-match #rx"(.*)( +)(.*)" (car l)))
      => (lambda (m)
           (loop (list* (cadr m)
                        (let ([len (string-length (caddr m))])
                          (if (= len 1)
                              (make-element 'tt (list " ")) ; to allow line breaks
                              (hspace len)))
                        (cadddr m)
                        (cdr l))
                 so-far))]
     [(equal? (car l) "")
      (loop (cdr l) so-far)]
     [else (loop (cdr l) (cons (car l) so-far))])))

(define (maybe-tt s)
  (if (string? s)
      (schemeparenfont s)
      s))

(define (code-block #:reader [reader read-syntax*] . strs)
  (make-table
   #f
   (map
    (lambda (l)
      (list (make-flow (list (make-paragraph (cons (hspace 1) (map maybe-tt l)))))))
    (split-lines (expr-decorate #:reader reader strs)))))

(define (code-elem #:reader [reader read-syntax*] . strs)
  (make-element
   #f
   (map maybe-tt 
        (car
         (split-lines
          (expr-decorate #:reader reader
                         (map (lambda (str) (regexp-replace #rx"\n" str " "))
                              strs)))))))