;; https://stevelosh.com/blog/2021/03/small-common-lisp-cli-programs/

(import scheme)
(import (chicken base))
(import (chicken bitwise))
(import (chicken format))
(import (chicken io))
(import (chicken irregex))
(import (chicken process-context))
(import (srfi 1))
(import (srfi 69))
(import stack)

(define (rgb-code r g b)
  (+ (* r 36) (* g 6) b 16))

(define (make-colors exclude?)
  (let ((stack (make-stack)))
    (do ((r 0 (add1 r)))
        ((= r 6))
      (do ((g 0 (add1 g)))
          ((= g 6))
        (do ((b 0 (add1 b)))
            ((= b 6))
          (when (not (exclude? (+ r g b)))
            (stack-push! stack (rgb-code r g b))))))
    (list->vector (stack-fold stack cons '()))))

(define dark-mode? (get-environment-variable "BATCHCOLOR_DARK"))

(define dark-colors (make-colors (lambda (v) (< v 3))))
(define light-colors (make-colors (lambda (v) (> v 11))))
(define mode-colors (if dark-mode? dark-colors light-colors))

(define explicits (make-hash-table))

(define (djb2 s)
  ;; http://www.cse.yorku.ca/~oz/hash.html
  (fold (lambda (c hash)
          (bitwise-and (+ (* hash 33) (char->integer c)) (sub1 (expt 2 64))))
        5381
        (string->list s)))

(define (find-color s)
  (hash-table-ref/default explicits string
                          (vector-ref mode-colors
                                      (modulo (djb2 s)
                                              (vector-length mode-colors)))))

(define (print-colorized s)
  (printf "~c[38;5;~am" #\escape (find-color s))
  (display s)
  (printf "~c[0m" #\escape))

(define (irregex-match-submatches match)
  (define (match->indices match i)
    (and-let* ((start (irregex-match-start-index match i))
               (end (irregex-match-end-index match i)))
      (cons start end)))
  (let ((match-count (irregex-match-num-submatches match)))
    (if (zero? match-count)
        (list (match->indices match 0))
        (filter-map (lambda (i) (match->indices match i))
                    (iota match-count 1)))))

(define (colorize-line irx line)
  (define (handle-match from-index match _seed)
    (let ((submatches (irregex-match-submatches match)))
      (for-each
       (lambda (indices)
         (let ((match-start (car indices))
               (match-end (cdr indices)))
           (display (substring line from-index match-start))
           (print-colorized (substring line match-start match-end))
           (set! from-index match-end)))
       submatches)
      (display (substring line from-index (irregex-match-end-index match)))
      (null? submatches)))
  (define (handle-finish from-index seed)
    (display (substring line from-index))
    seed)
  (when (irregex-fold irx handle-match #f line handle-finish) ; no submatches?
    (display line))
  (newline))

(define (colorize-text irx port)
  (let loop ()
    (let ((line (read-line port)))
      (when (not (eof-object? line))
        (colorize-line irx line)
        (loop)))))

(define (main regex #!rest paths)
  (let ((irx (irregex regex))
        (paths (if (null? paths) '("-") paths)))
    (for-each
     (lambda (path)
       (if (equal? path "-")
           (colorize-text irx (current-input-port))
           (call-with-input-file path (lambda (in) (colorize-text irx in)))))
     paths)))

(apply main (command-line-arguments))