#lang scheme
(provide doc
         char
         double
         text
         line
         <> 
         fsep
         hcat
         punctuate
         (rename-out (doc->string compact))
         (rename-out (doc->pretty-string pretty))
         fill
         nest)

(define-struct doc (type data))
(define (doc-type=? d d-type) (and (doc? d) (eq? d-type (doc-type d))))

(define (empty? value) (doc-type=? value 'empty))
(define (empty) (make-doc 'empty null))

(define (line? value) (doc-type=? value 'line))
(define (line) (make-doc 'line null))

(define (is-char? value) (doc-type=? value 'char))
(define (char value) 
  (if (char? value)
      (make-doc 'char value)
      (error "char expects type number : " value)))

(define (text? value) (doc-type=? value 'text))
(define (text value) 
  (if (string? value)
      (make-doc 'text value)
      (error "text expects type string : " value)))

(define (concat? value) (doc-type=? value 'concat))
(define (concat l r) (make-doc 'concat (cons l r)))
(define (concat-l c) (car (doc-data c)))
(define (concat-r c) (cdr (doc-data c)))

(define (union? value) (doc-type=? value 'union))
(define (union l r) (make-doc 'union (cons l r)))
(define (union-l c) (car (doc-data c)))
(define (union-r c) (cdr (doc-data c)))

(define (double value) 
  (if (number? value)
      (text (number->string value))
      (error "double expects type number : " value)))

(define (flatten x)
  (cond [(concat? x) (concat (flatten (concat-l x))
                             (flatten (concat-r x)))]
        [(line? x)   (char #\space)]
        [(union? x)  (flatten (union-l x))]
        [else        x]))
(define (group d) (union (flatten d) d))
(define softline (group line))
(define (</> l r)
  (if (empty? r) l 
      (<> l (softline <> r))))

(define (fold f) (λ (ds) (foldr f (empty) ds)))
(define (<> l r)
  (cond [(empty? l) r]
        [(empty? r) l]
        [else (concat l r)]))
(define hcat (fold <>))
(define (fsep) (fold </>))

(define (punctuate d ds)
  (cond [(null? ds) null]
        [(null? (rest ds)) ds]
        [else (cons (<> (first ds) d)
                    (punctuate d (rest ds)))]))

(define (doc->string x)
  (list->string
   (let transform ([ds (list x)])
     (if (null? ds) 
         null
         (let ([d (first ds)]
               [ds (rest ds)])
           (cond [(empty? d)   (transform ds)]
                 [(is-char? d) (cons (doc-data d) (transform ds))]
                 [(text? d)  (append (string->list (doc-data d)) (transform ds))]
                 [(line? d)    (cons #\newline (transform ds))]
                 [(concat? d)  (transform (cons (concat-l d) (cons (concat-r d) ds)))]
                 [(union? d)   (transform (cons (union-r d) ds))]))))))

(define (doc->pretty-string width x)
  (define (nicest col s1 s2)
    (if (fits (- width (min width col)) s1)
        s1
        s2))
  
  (define (fits w s)
    (cond [(< w 0)       false]
          [(equal? s "") true]
          [(char=? (string-ref s 0) #\newline) true]
          [else (fits (sub1 w) (substring s 1))]))
  
  (list->string
   (let best ([col 0]
              [ds (list x)])
     (if (null? ds) 
         null
         (let ([d (first ds)]
               [ds (rest ds)])
           (cond [(empty? d)   (best col ds)]
                 [(is-char? d) (cons (doc-data d) (best (+ col 1) ds))]
                 [(text? d)    (append (string->list (doc-data d)) 
                                       (best (+ col 
                                                (string-length (doc-data d))) ds))]
                 [(line? d)    (cons #\newline (best 0 ds))]
                 [(concat? d)  (best col (cons (concat-l d) (cons (concat-r d) ds)))]
                 [(union? d)   (nicest col 
                                       (best col (cons (union-l d) ds))
                                       (best col (cons (union-r d) ds)))]
                 [else null]))))))

(define (fill width d)
  (define (pad-line col) (text (make-string (max 0 (- width col)) #\space)))
  (let scan-lines ([col 0]
                   [ds (list d)])
    (cond [(empty? d)   (<> d (scan-lines col ds))]
          [(is-char? d) (<> d (scan-lines (add1 col) ds))]
          [(text? d)    (<> d (scan-lines (+ (string-length (doc-data d)) col) ds))]
          [(line? d)    (<> (<> (pad-line col) (line))
                            (scan-lines 0 ds))]
          [(concat? d)  (scan-lines col (cons (concat-l d) 
                                              (cons (concat-r d) ds)))]
          [(union? d)   (union (scan-lines col (cons (union-l d) ds))
                               (scan-lines col (cons (union-r d) ds)))]
          [else (empty)])))

(define (nest indent d)
  (define (indent-line pos) (text (make-string pos #\space)))
  (define (offset c) 
    (cond [(memq c (list #\{ #\[)) indent] 
          [(memq c (list #\} #\])) (- indent)]
          [else 0]))
  (let scan-lines ([col 0]
                   [ds (list d)])
    (cond [(empty? d)   (<> d (scan-lines col ds))]
          [(is-char? d) (<> d (scan-lines (+ col (offset (doc-data d))) ds))]
          [(text? d)    (<> d (scan-lines col ds))]
          [(line? d)    (<> d (<> (indent-line col)
                                  (scan-lines col ds)))]
          [(concat? d)  (scan-lines col (cons (concat-l d) 
                                              (cons (concat-r d) ds)))]
          [(union? d)   (union (scan-lines col (cons (union-l d) ds))
                               (scan-lines col (cons (union-r d) ds)))]
          [else (empty)])))


;(define e (empty))
;(define c (char #\s))
;(define t (text "some text"))
;(define d (double pi))
;(define l (line))
;(define ds (list e c l t l d))
;(define dss (append ds (list l c t t c l d l l t e)))
;(define tcs (list c t t c t c c c t t t t c t c t))
;(define test-doc-small  (hcat ds))
;(define test-doc-bigger (hcat dss))
;(define just-text-and-chars (hcat tcs))