#lang racket/base
(require syntax/strip-context
         syntax-color/module-lexer
         syntax-color/lexer-contract
         "../racket.rkt"
         "../base.rkt"
         "manual-scheme.rkt"
         "manual-style.rkt"
         scribble/core
         (for-syntax racket/base
                     syntax/parse))

(provide codeblock
         codeblock0
         typeset-code
         code)

(define-for-syntax (do-codeblock stx)
  (syntax-parse stx
    [(_ (~seq (~or (~optional (~seq #:expand expand-expr:expr)
                              #:defaults ([expand-expr #'#f])
                              #:name "#:expand keyword")
                   (~optional (~seq #:indent indent-expr:expr)
                              #:defaults ([indent-expr #'0])
                              #:name "#:expand keyword")
                   (~optional (~seq #:keep-lang-line? keep-lang-line?-expr:expr)
                              #:defaults ([keep-lang-line?-expr #'#t])
                              #:name "#:keep-lang-line? keyword")
                   (~optional (~seq #:context context-expr:expr)
                              #:name "#:context keyword")
                   (~optional (~seq #:line-numbers line-numbers:expr)
                              #:defaults ([line-numbers #'#f])
                              #:name "#:line-numbers keyword")
                   (~optional (~seq #:line-number-sep line-number-sep:expr)
                              #:defaults ([line-number-sep #'1])
                              #:name "#:line-number-sep keyword"))
              ...)
        str ...)
     #`(typeset-code str ...
                     #:expand expand-expr
                     #:keep-lang-line? keep-lang-line?-expr
                     #:indent indent-expr
                     #:context #,(if (attribute context-expr)
                                     #'context-expr
                                     (or
                                      (let ([v #'(str ...)])
                                        (and (pair? (syntax-e v))
                                             #`#'#,(car (syntax-e v))))
                                      #'#f))
                     #:line-numbers line-numbers
                     #:line-number-sep line-number-sep)]))

(define-syntax (codeblock stx) #`(code-inset #,(do-codeblock stx)))
(define-syntax (codeblock0 stx) (do-codeblock stx))

(define (typeset-code #:context [context #f]
                      #:expand [expand #f]
                      #:indent [indent 2]
                      #:keep-lang-line? [keep-lang-line? #t]
                      #:line-numbers [line-numbers #f]
                      #:line-number-sep [line-number-sep 1]
                      #:block? [block? #t]
                      . strs)
  (define-values (tokens bstr) (get-tokens strs context expand))
  (define default-color meta-color)
  ((if block? table (lambda (style lines) (make-element #f lines)))
   block-color
   ((if keep-lang-line? values cdr) ; FIXME: #lang can span lines
    (list->lines
     indent
     #:line-numbers line-numbers
     #:line-number-sep line-number-sep
     #:block? block?
     (let loop ([pos 0]
                [tokens tokens])
       (cond
         [(null? tokens) (split-lines default-color (substring bstr pos))]
         [(eq? (caar tokens) 'white-space) (loop pos (cdr tokens))]
         [(= pos (cadar tokens))
          (append (let ([style (caar tokens)]
                        [get-str (lambda ()
                                   (substring bstr (cadar tokens) (caddar tokens)))])
                    (cond
                      [(symbol? style)
                       (define scribble-style
                         (case style
                           [(symbol) symbol-color]
                           [(parenthesis hash-colon-keyword) paren-color]
                           [(constant string) value-color]
                           [(comment) comment-color]
                           [else default-color]))
                       (split-lines scribble-style (get-str))]
                      [(procedure? style)
                       (list (style (get-str)))]
                      [else (list style)]))
                  (loop (caddar tokens) (cdr tokens)))]
         [(> pos (cadar tokens))
          (loop pos (cdr tokens))]
         [else (append
                (split-lines default-color (substring bstr pos (cadar tokens)))
                (loop (cadar tokens) tokens))]))))))

;; (listof string) boolean boolean -> tokens string
;; tokens is a
;; (listof (list T natural natural natural))
;; T being a symbol returned as a token type from the languages lexer
;;   OR a function created by get-tokens
;; the first number being the start position
;; the second being the end position
;; the third 0 if T is a symbol, and 1 or greater if its a function or element
;; the tokens are sorted by the start end end positions
(define (get-tokens strs context expand)
  (define xstr (apply string-append strs))
  (define bstr (regexp-replace* #rx"(?m:^$)" xstr "\xA0"))
  (define in (open-input-string bstr))
  (port-count-lines! in)
  (let* ([tokens
          (let loop ([mode #f])
            (let-values ([(lexeme type data start end backup-delta mode)
                          (module-lexer in 0 mode)])
              (if (equal? type 'eof)
                  null
                  (cons (list type (sub1 start) (sub1 end) 0)
                        (loop (if (dont-stop? mode)
                                  (dont-stop-val mode)
                                  mode))))))]
         ;; use a source that both identifies the original code
         ;; and is unique wrt eq? as used below
         [program-source (or context bstr)]
         [e (parameterize ([read-accept-reader #t])
              ((or expand 
                   (lambda (stx) 
                     (if context
                         (replace-context context stx)
                         stx)))
               (let ([p (open-input-string bstr)])
                 (port-count-lines! p)
                 (let loop ()
                   (define v (read-syntax program-source p))
                   (cond
                     [expand v]
                     [(eof-object? v) null]
                     [else (datum->syntax #f (cons v (loop)) v v)])))))]
         [ids (let loop ([e e])
                (cond
                  [(and (identifier? e)
                        (syntax-original? e)
                        (syntax-position e)
                        (eq? program-source (syntax-source e)))
                   (define pos (sub1 (syntax-position e)))
                   (list (list (lambda (str)
                                 (to-element (syntax-property
                                              e
                                              'display-string
                                              str)
                                             #:escapes? #f))
                               pos
                               (+ pos (syntax-span e))
                               1))]
                  [(syntax? e) (append (loop (syntax-e e))
                                       (loop (or (syntax-property e 'origin)
                                                 null))
                                       (loop (or (syntax-property e 'disappeared-use)
                                                 null)))]
                  [(pair? e) (append (loop (car e)) (loop (cdr e)))]
                  [else null]))]
         [link-mod (lambda (mp-stx priority #:orig? [always-orig? #f])
                     (cond
                       [(or always-orig?
                            (syntax-original? mp-stx))
                        (define mp (syntax->datum mp-stx))
                        (define pos (sub1 (syntax-position mp-stx)))
                        (list (list (racketmodname #,mp)
                                    pos
                                    (+ pos (syntax-span mp-stx))
                                    priority))]
                       [else null]))]
         ;; This makes sense when `expand' actually expands, and
         ;; probably not otherwise:
         [mods (let loop ([e e])
                 (syntax-case e (module #%require begin)
                   [(module name lang (mod-beg form ...))
                    (apply append
                           (link-mod #'lang 2)
                           (map loop (syntax->list #'(form ...))))]
                   [(#%require spec ...)
                    (apply append
                           (map (lambda (spec)
                                  ;; Need to add support for renaming forms, etc.:
                                  (if (module-path? (syntax->datum spec))
                                      (link-mod spec 2)
                                      null))
                                (syntax->list #'(spec ...))))]
                   [(begin form ...)
                    (apply append
                           (map loop (syntax->list #'(form ...))))]
                   [else null]))]
         [has-hash-lang? (regexp-match? #rx"^#lang " bstr)]
         [hash-lang (if has-hash-lang?
                        (list (list (hash-lang)
                                    0
                                    5
                                    1)
                              (list 'white-space 5 6 0))
                        null)]
         [language (cond
                     [has-hash-lang?
                      (define m (regexp-match #rx"^#lang ([-0-9a-zA-Z/._+]+)" bstr))
                      (if m
                          (link-mod
                           #:orig? #t
                           (datum->syntax #f
                                          (string->symbol (cadr m))
                                          (vector 'in 1 6 7 (string-length (cadr m))))
                           3)
                          null)]
                     [else null])]
         [tokens (sort (append ids
                               mods
                               hash-lang
                               language
                               (filter (lambda (x) (not (eq? (car x) 'symbol)))
                                       (if has-hash-lang?
                                           ;; Drop #lang entry:
                                           (cdr tokens)
                                           tokens)))
                       (lambda (a b)
                         (or (< (cadr a) (cadr b))
                             (and (= (cadr a) (cadr b))
                                  (> (cadddr a) (cadddr b))))))])
    (values tokens bstr)))

(define (typeset-code-line context expand lang-line . strs)
  (typeset-code
   #:context context
   #:expand expand
   #:keep-lang-line? (not lang-line)
   #:block? #f
   #:indent 0
   (let ([s (regexp-replace* #px"(?:\\s*(?:\r|\n|\r\n)\\s*)+" (apply string-append strs) " ")])
     (if lang-line
         (string-append "#lang " lang-line "\n" s)
         s))))

(define-syntax (code stx)
  (syntax-parse stx
    [(_ (~seq (~or (~optional (~seq #:expand expand-expr:expr)
                              #:defaults ([expand-expr #'#f])
                              #:name "#:expand keyword")
                   (~optional (~seq #:context context-expr:expr)
                              #:name "#:context keyword")
                   (~optional (~seq #:lang lang-line-expr:expr)
                              #:defaults ([lang-line-expr #'#f])
                              #:name "#:lang-line keyword"))
              ...)
        str ...)
     #`(typeset-code-line #,(if (attribute context-expr)
                                #'context-expr
                                (or
                                 (let ([v #'(str ...)])
                                   (and (pair? (syntax-e v))
                                        #`#'#,(car (syntax-e v))))
                                 #'#f))
                          expand-expr
                          lang-line-expr
                          str ...)]))

(define (split-lines style s)
  (cond
   [(regexp-match-positions #rx"(?:\r\n|\r|\n)" s)
    => (lambda (m)
         (append (split-lines style (substring s 0 (caar m)))
                 (list 'newline)
                 (split-lines style (substring s (cdar m)))))]
   [(regexp-match-positions #rx" +" s)
    => (lambda (m)
         (append (split-lines style (substring s 0 (caar m)))
                 (list (hspace (- (cdar m) (caar m))))
                 (split-lines style (substring s (cdar m)))))]
   [else (list (element style s))]))

(define omitable (make-style #f '(omitable)))

(define (list->lines indent-amt l 
                     #:line-numbers line-numbers
                     #:line-number-sep line-number-sep
                     #:block? block?)
  (define indent-elem (if (zero? indent-amt)
                          ""
                          (hspace indent-amt)))
  ;(list of any) delim -> (list of (list of any))
  (define (break-list lst delim)
    (let loop ([l lst] [n null] [c null])
      (cond
       [(null? l) (reverse (if (null? c) n (cons (reverse c) n)))]
       [(eq? delim (car l)) (loop (cdr l) (cons (reverse c) n) null)]
       [else (loop (cdr l) n (cons (car l) c) )])))

  (define lines (break-list l 'newline))
  (define line-cnt (length lines))
  (define line-cntl (string-length (format "~a" (+ line-cnt (or line-numbers 0)))))

  (define (prepend-line-number n r)
    (define ln (format "~a" n))
    (define lnl (string-length ln))
    (define diff (- line-cntl lnl))
    (define l1 (list (tt ln) (hspace line-number-sep)))
    (cons (make-element 'smaller 
                        (make-element 'smaller  
                                      (if (not (zero? diff))
                                          (cons (hspace diff) l1)
                                          l1)))
          r))

  (define (make-line accum-line line-number)
    (define rest (cons indent-elem accum-line))
    (list ((if block? paragraph (lambda (s e) e))
           omitable 
           (if line-numbers
               (prepend-line-number line-number rest)
               rest))))

  (for/list ([l (break-list l 'newline)]
             [i (in-naturals (or line-numbers 1))])
    (make-line l i)))


;; ----------------------------------------

(module+ test
  (require racket/list
           racket/match
           rackunit)

  (define (tokens strs)
    (define-values (toks _) (get-tokens strs #f #f))
    (for/list ([tok (in-list toks)])
      (match tok
        [(list _ start end (or 1 2 3))
         (list 'function start end 1)] ; this looses information
        [_ tok])))

  (define (make-test-result  lst)
    (define-values (res _)
      (for/fold ([result null] [count 12])
                ([p lst])
        (define next (+ count (second p)))
        (define r (if (eq? (first p) 'function) 1 0))
        (values
         (cons (list (first p) count next r) result)
         next)))
    (list* `(function 0 5 1) `(white-space 5 6 0) `(function 6 12 1) `(function 6 12 1)
           (reverse res)))

  (check-equal?
   (tokens (list "#lang racket\n1"))
   `((function 0 5 1) (white-space 5 6 0) ;"#lang "
     (function 6 12 1) (function 6 12 1) (white-space 12 13 0) ;"racket\n"
     (constant 13 14 0))) ; "1"
  (check-equal?
   (tokens (list "#lang racket\n" "(+ 1 2)"))
   (make-test-result
    '((white-space 1)
      (parenthesis 1) (function 1)
      (white-space 1) (constant 1) (white-space 1) (constant 1)
      (parenthesis 1))))
  (check-equal?
   (tokens (list "#lang racket\n(apply x (list y))"))
   (make-test-result
    '((white-space 1)
      (parenthesis 1)
      (function 5) (white-space 1);apply
      (function 1) (white-space 1);x
      (parenthesis 1)
      (function 4) (white-space 1) (function 1);list y
      (parenthesis 1)
      (parenthesis 1)))))
