(define-module (reflow wrap)
  #:use-module (reflow ansi)
  #:use-module (uniseg)
  #:use-module (uniseg graphemes)
  #:use-module (uniseg graphemes iterator)
  #:use-module (ice-9 textual-ports)
  #:use-module (srfi srfi-26)
  #:export (make-wrapping-port))

(define* (make-wrapping-port
          o-port
          max-width
          #:key (keep-newlines? #t)
          (preserve-space? #f)
          (tab-width 4))
  "Return a soft port which wraps the given output port such that text cleanly wraps to a new line if longer than max width."
  (define-values (ansi-port current-sequence reset-ansi restore-ansi)
    (make-ansi-port-tuple o-port))

  (define iterator (make-grapheme-iterator))
  (define %cur-width 0)
  (define %in-escape-seq? #f)
  (define %waiting-for-non-space? #f)

  (define (put-n-times char num)
    (do ((i 1 (1+ 1)))
        ((> i num))
        (put-char ansi-port char)))

  (define (new-line)
    (when (not preserve-space?)
      (set! %waiting-for-non-space? #t))
    (put-char ansi-port #\newline))

  (define (handle-char char)
    (cond
     ;; Track when we are in an escape sequence
     ((eq? %ansi-marker char)
      (set! %in-escape-seq? #t)
      (put-char ansi-port char))

     ;; Skip processing escape sequences and track when we leave escape sequence
     (%in-escape-seq?
      (put-char ansi-port char)
      (when (ansi-terminator? char)
        (set! %in-escape-seq? #f)))

     ;; Support for stripping spaces from front of newline
     ((eq? #\newline char)
      (when keep-newlines?
        (set! %cur-width 0)
        (new-line)))

     ;; Skip whitespace at beginning of line if configured
     ((and (char-set-contains? char-set:whitespace char)
           %waiting-for-non-space?)
      #f)

     ;; Convert tabs to desired number of spaces
     ((eq? #\tab char)
      (for-each handle-char (make-list tab-width #\space)))

     ;; Subtract backspaces from width
     ((eq? #\backspace char)
      (set! %cur-width (- %cur-width 1))
      (put-char ansi-port char))

     ;; Main processing area
     (else
      (set! %waiting-for-non-space? #f)

      (let* ((grapheme (iterator char))
             (new-width (+ %cur-width (grapheme-delta-width grapheme))))

        (if (> new-width max-width)
            (let* ((g-width (grapheme-width grapheme))
                   (g-modification? (grapheme-modification? grapheme))
                   (g-string (grapheme-string grapheme)))
              ;; If we go over the width, and the grapheme is a modification,
              ;; then we need to backspace to before the grapheme before we go to the next line.
              ;; But we also need to add spaces to fill in the characters already printed!
              (when g-modification?
                (let* ((num-to-delete (- g-width 1)))
                  (put-n-times #\backspace num-to-delete)
                  (put-n-times #\space num-to-delete)))

              ;; This possibly sets %waiting-for-non-space = #t
              ;; and we need to handle it.
              (new-line)
              (when %waiting-for-non-space?
                ;; Remove all whitespace chars from the grapheme
                (set! g-string (string-filter (λ (c) (not (char-set-contains? char-set:whitespace c))) g-string))
                ;; Recalculate width with remaining characters
                (set! g-width (string-width g-string))

                ;; If anything is left, we have satisfied the newline non-space check already
                (when (> (string-length g-string) 0)
                  (set! %waiting-for-non-space? #f)))
              (restore-ansi)
              (put-string ansi-port g-string)
              (set! %cur-width g-width))

            (begin
              ;; Otherwise we can just output to the ansi port as normal.
              (put-char ansi-port char)
              (set! %cur-width new-width)))))))

  (define output-port
    (make-soft-port
     (vector
      ;; Char out
      handle-char
      ;; string out, char by char
      (λ (s) (string-for-each handle-char s))
      #f #f #f)
     "w"))

  output-port)