(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!"))

(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) (render-sx-tokens (tokenize-sx code))))

(define
  highlight
  (fn
    (code lang)
    (if
      (or (= lang "lisp") (= lang "sx") (= lang "sexp") (= lang "scheme"))
      (highlight-sx code)
      (list (quote code) code))))