

(define-library
  (sx highlight)
  (export
    sx-specials
    sx-special?
    hl-digit?
    hl-alpha?
    hl-sym-char?
    hl-ws?
    hl-escape
    hl-span
    tokenize-sx
    sx-token-classes
    render-sx-tokens
    highlight-sx
    highlight)
  (begin
    (define
      sx-specials
      (list
        "defcomp"
        "defrelation"
        "defisland"
        "defpage"
        "defhelper"
        "define"
        "defmacro"
        "defconfig"
        "deftest"
        "if"
        "when"
        "cond"
        "case"
        "and"
        "or"
        "not"
        "let"
        "let*"
        "lambda"
        "fn"
        "do"
        "begin"
        "quote"
        "quasiquote"
        "->"
        "map"
        "filter"
        "reduce"
        "some"
        "every?"
        "map-indexed"
        "for-each"
        "&key"
        "&rest"
        "set!"
        "satisfies?"
        "match"
        "let-match"
        "define-protocol"
        "implement"
        "->>"
        "|>"
        "as->"
        "define-library"
        "import"
        "perform"
        "guard"
        "call/cc"
        "raise"
        "define-syntax"
        "syntax-rules"
        "make-parameter"
        "parameterize"))
    (define sx-special? (fn (s) (some (fn (x) (= x s)) sx-specials)))
    (define hl-digit? (fn (c) (and (>= c "0") (<= c "9"))))
    (define
      hl-alpha?
      (fn
        (c)
        (or (and (>= c "a") (<= c "z")) (and (>= c "A") (<= c "Z")))))
    (define
      hl-sym-char?
      (fn
        (c)
        (or
          (hl-alpha? c)
          (hl-digit? c)
          (= c "_")
          (= c "-")
          (= c "?")
          (= c "!")
          (= c "+")
          (= c "*")
          (= c "/")
          (= c "<")
          (= c ">")
          (= c "=")
          (= c "&")
          (= c "."))))
    (define
      hl-ws?
      (fn (c) (or (= c " ") (= c "\n") (= c "\t") (= c "\r"))))
    (define hl-escape (fn (s) s))
    (define
      hl-span
      (fn
        (class text)
        (if
          (= class "")
          (list (quote span) text)
          (list (quote span) (make-keyword "class") class text))))
    (define
      tokenize-sx
      (fn
        (code)
        (let
          ((tokens (list)) (i 0) (len (string-length code)))
          (let
            loop
            ()
            (when
              (< i len)
              (let
                ((c (substring code i (+ i 1))))
                (if
                  (= c ";")
                  (let
                    ((start i))
                    (set! i (+ i 1))
                    (let
                      scan
                      ()
                      (when
                        (and
                          (< i len)
                          (not (= (substring code i (+ i 1)) "\n")))
                        (set! i (+ i 1))
                        (scan)))
                    (set!
                      tokens
                      (append
                        tokens
                        (list (list "comment" (substring code start i))))))
                  (if
                    (= c "\"")
                    (let
                      ((start i))
                      (set! i (+ i 1))
                      (let
                        sloop
                        ()
                        (when
                          (< i len)
                          (let
                            ((sc (substring code i (+ i 1))))
                            (if
                              (= sc "\\")
                              (do (set! i (+ i 2)) (sloop))
                              (if
                                (= sc "\"")
                                (set! i (+ i 1))
                                (do (set! i (+ i 1)) (sloop)))))))
                      (set!
                        tokens
                        (append
                          tokens
                          (list (list "string" (substring code start i))))))
                    (if
                      (= c ":")
                      (let
                        ((start i))
                        (set! i (+ i 1))
                        (when
                          (and
                            (< i len)
                            (hl-alpha? (substring code i (+ i 1))))
                          (let
                            scan
                            ()
                            (when
                              (and
                                (< i len)
                                (hl-sym-char? (substring code i (+ i 1))))
                              (set! i (+ i 1))
                              (scan))))
                        (set!
                          tokens
                          (append
                            tokens
                            (list (list "keyword" (substring code start i))))))
                      (if
                        (= c "~")
                        (let
                          ((start i))
                          (set! i (+ i 1))
                          (let
                            scan
                            ()
                            (when
                              (and
                                (< i len)
                                (let
                                  ((x (substring code i (+ i 1))))
                                  (or (hl-sym-char? x) (= x "/"))))
                              (set! i (+ i 1))
                              (scan)))
                          (set!
                            tokens
                            (append
                              tokens
                              (list
                                (list "component" (substring code start i))))))
                        (if
                          (or
                            (= c "(")
                            (= c ")")
                            (= c "[")
                            (= c "]")
                            (= c "{")
                            (= c "}"))
                          (do
                            (set!
                              tokens
                              (append tokens (list (list "paren" c))))
                            (set! i (+ i 1)))
                          (if
                            (hl-digit? c)
                            (let
                              ((start i))
                              (let
                                scan
                                ()
                                (when
                                  (and
                                    (< i len)
                                    (let
                                      ((x (substring code i (+ i 1))))
                                      (or (hl-digit? x) (= x "."))))
                                  (set! i (+ i 1))
                                  (scan)))
                              (set!
                                tokens
                                (append
                                  tokens
                                  (list
                                    (list "number" (substring code start i))))))
                            (if
                              (hl-sym-char? c)
                              (let
                                ((start i))
                                (let
                                  scan
                                  ()
                                  (when
                                    (and
                                      (< i len)
                                      (hl-sym-char?
                                        (substring code i (+ i 1))))
                                    (set! i (+ i 1))
                                    (scan)))
                                (let
                                  ((text (substring code start i)))
                                  (if
                                    (or
                                      (= text "true")
                                      (= text "false")
                                      (= text "nil"))
                                    (set!
                                      tokens
                                      (append
                                        tokens
                                        (list (list "boolean" text))))
                                    (if
                                      (sx-special? text)
                                      (set!
                                        tokens
                                        (append
                                          tokens
                                          (list (list "special" text))))
                                      (set!
                                        tokens
                                        (append
                                          tokens
                                          (list (list "symbol" text))))))))
                              (if
                                (hl-ws? c)
                                (let
                                  ((start i))
                                  (let
                                    scan
                                    ()
                                    (when
                                      (and
                                        (< i len)
                                        (hl-ws? (substring code i (+ i 1))))
                                      (set! i (+ i 1))
                                      (scan)))
                                  (set!
                                    tokens
                                    (append
                                      tokens
                                      (list
                                        (list "ws" (substring code start i))))))
                                (do
                                  (set!
                                    tokens
                                    (append tokens (list (list "other" c))))
                                  (set! i (+ i 1))))))))))))
              (loop)))
          tokens)))
    (define sx-token-classes {:boolean "text-orange-600" :component "text-rose-600 font-semibold" :number "text-amber-700" :string "text-emerald-700" :special "text-sky-700 font-semibold" :paren "text-stone-400" :keyword "text-violet-600" :comment "text-stone-400 italic"})
    (define
      render-sx-tokens
      (fn
        (tokens)
        (map
          (fn
            (tok)
            (let
              ((cls (or (dict-get sx-token-classes (first tok)) "")))
              (hl-span cls (nth tok 1))))
          tokens)))
    (define highlight-sx (fn (code) (-> code tokenize-sx render-sx-tokens)))
    (define
      highlight
      (fn
        (code lang)
        (if
          (or
            (= lang "lisp")
            (= lang "sx")
            (= lang "sexp")
            (= lang "scheme"))
          (highlight-sx code)
          (list (quote code) code)))))) ;; end define-library

;; Re-export to global namespace for backward compatibility
(import (sx highlight))
