;;; zsh-ts-mode.el --- TreeSitter based mode for editing Zsh sripts  -*- lexical-binding: t; -*-
;;; Commentary:
;;; Code:

(require 'sh-script)
(require 'treesit)
(require 'rx)

(defun zsh-ts-mode--fontify-let-arg (node _override _start _end &rest _r)
  "Internal function used from `zsh-ts-mode--treesit-settings'.
Fontify NODE, an argument to a Zsh \"let\" statement."
  (let ((ns (treesit-node-start node))
        (ne (treesit-node-end node)))
    (save-excursion
      (goto-char ns)
      (if (re-search-forward (rx (or "=" "[")) ne t)
          (put-text-property ns (1- (point)) 'face
                             'font-lock-variable-name-face)
        (put-text-property ns ne 'face 'font-lock-variable-name-face)))))

(defvar zsh-ts-mode--arithmetic-variable-query
  (treesit-query-compile 'zsh '((word) @name))
  "Query used in `zsh-ts-mode--fontify-arithmetic-variables'.")

(defun zsh-ts-mode--fontify-arithmetic-variables
    (node _override _start _end &rest _r)
  "Fontify all arithmetic variables below NODE.
NODE should be an `arithmetic_expansion' node."
  (dolist (child (treesit-query-capture
                  node zsh-ts-mode--arithmetic-variable-query))
    (let ((start (treesit-node-start (cdr child)))
          (end (treesit-node-end (cdr child))))
      (font-lock-append-text-property start end 'face
                                      'font-lock-variable-use-face))))

(defun zsh-ts-mode--other-keywords ()
  "Return a list of `other' keywords for Zsh.
This just appends some more stuff to `sh-mode--treesit-other-keywords'."
  (append (sh-mode--treesit-other-keywords) '("let")))

(defvar zsh-ts-mode--declaration-commands
  '("typeset" "local" "integer" "float" "readonly" "export" "declare")
  "Return a list of keywords that declare a parameter in Zsh.
Note that \"let\" isn't included here as the parser treats it as a normal
command.")

;; from `bash-ts-mode'
(defvar zsh-ts-mode--treesit-settings
  (treesit-font-lock-rules
   :feature 'comment
   :language 'zsh
   '((comment) @font-lock-comment-face)

   :feature 'function
   :language 'zsh
   '((function_definition name: (word) @font-lock-function-name-face))

   :feature 'string
   :language 'zsh
   '(((raw_string) @font-lock-string-face)
     ((string :anchor ("\"" @font-lock-string-face)
              [(string_content)
               (arithmetic_expansion :anchor "((")]
              :* @font-lock-string-face
              ("\"" @font-lock-string-face) :anchor)))

   :feature 'string-interpolation
   :language 'zsh
   :override t
   '((command_substitution :anchor ("$(" @sh-quoted-exec)
                           (")" @sh-quoted-exec) :anchor)
     ;; work around parser error of "(( ))" (no $) being an arithmetic expansion
     (string (arithmetic_expansion :anchor ("$((" @sh-quoted-exec)
                                   ("))" @sh-quoted-exec) :anchor)))

   :feature 'variable-reference
   :language 'zsh
   :override t
   '((expansion (simple_variable_name) @font-lock-variable-use-face)
     (special_variable_name) @font-lock-variable-use-face
     (expansion_with_modifier (simple_variable_name)
                              @font-lock-variable-use-face)
     (expansion_substring (simple_variable_name) @font-lock-variable-use-face)
     (variable_ref (simple_variable_name) @font-lock-variable-use-face)
     ;; variables in math expansions
     ((arithmetic_expansion :anchor "$((")
      @zsh-ts-mode--fontify-arithmetic-variables))

   :feature 'math-function-call
   :language 'zsh
   :override t
   `((arithmetic_call name: (word) @font-lock-function-name-face))

   :feature 'heredoc
   :language 'zsh
   '([(heredoc_start) (heredoc_body)] @sh-heredoc)

   :feature 'keyword
   :language 'zsh
   `(;; keywords
     [ ,@sh-mode--treesit-keywords ] @font-lock-keyword-face
     ;; reserved words
     (command_name
      ((word) @font-lock-keyword-face
       (:match ,(rx-to-string
                 `(seq bos
                       (or ,@(zsh-ts-mode--other-keywords))
                       eos))
               @font-lock-keyword-face))))

   :feature 'command
   :language 'zsh
   `(;; builtin commands
     (command_name
      ((word) @font-lock-builtin-face
       (:match ,(let ((builtins
                       (sh-feature sh-builtins)))
                  (rx-to-string
                   `(seq bos
                         (or ,@builtins)
                         eos)))
               @font-lock-builtin-face)))
     ;; function/non-builtin command calls
     (command_name (word) @font-lock-function-name-face))

   :feature 'declaration-command
   :language 'zsh
   ;; the parser doesn't treat "let" as a keyword
   `([,@zsh-ts-mode--declaration-commands] @font-lock-keyword-face)

   ;; variable declarations in the form of a=b
   ;; e.g. local x=2
   :feature 'decl-equal-variable
   :language 'zsh
   '((variable_name) @font-lock-variable-name-face)

   ;; variable declarations in the form of a
   ;; e.g. local y z
   :feature 'decl-solo-variable
   :language 'zsh
   :override 'prepend
   `((declaration_command
      argument: ((_) @font-lock-variable-name-face
                 (:match ,(rx-to-string '(seq bos (not "-")))
                         @font-lock-variable-name-face))))


   ;; let variable declarations (handled differently by the parser)
   :feature 'decl-let-variable
   :language 'zsh
   :override 'prepend
   `((command
      (command_name ((word) @ignore
                     (:equal "let" @ignore)))
      argument: ((_) @zsh-ts-mode--fontify-let-arg
                 (:match ,(rx-to-string '(seq bos (not "-")))
                         @zsh-ts-mode--fontify-let-arg))))

   :feature 'for-variable
   :language 'zsh
   `("for" variable: ((simple_variable_name) @font-lock-variable-name-face))

   :feature 'constant
   :language 'zsh
   '((case_item value: (word) @font-lock-constant-face)
     (file_descriptor) @font-lock-constant-face)

   :feature 'operator
   :language 'zsh
   `([,@sh-mode--treesit-operators] @font-lock-operator-face)

   :feature 'number
   :language 'zsh
   `(((word) @font-lock-number-face
      (:match "\\`[0-9]+\\'" @font-lock-number-face))
     ((number) @font-lock-number-face))

   :feature 'bracket
   :language 'zsh
   '((["(" ")" "((" "))" "[" "]" "[[" "]]" "{" "}"]) @font-lock-bracket-face)

   :feature 'delimiter
   :language 'zsh
   '(([";" ";;"]) @font-lock-delimiter-face)

   :feature 'misc-punctuation
   :language 'zsh
   '((["$"]) @font-lock-misc-punctuation-face)))

;;;###autoload
(define-derived-mode zsh-ts-mode sh-base-mode "Zsh"
  ;; This function is based mostly on `bash-ts-mode'.
  "Major mode for editing Zsh shell scripts.
This mode automatically falls back to `sh-mode' if the buffer is
not written in Zsh."
  :syntax-table sh-mode-syntax-table
  (when (treesit-ready-p 'zsh)
    (sh-set-shell "zsh")
    (add-hook 'hack-local-variables-hook
              #'sh-after-hack-local-variables nil t)
    (treesit-parser-create 'zsh)
    (setq-local treesit-font-lock-feature-list
                '((comment function)
                  (command declaration-command keyword string math-function-call)
                  (constant heredoc number string-interpolation
                   decl-equal-variable decl-solo-variable decl-let-variable
                   for-variable variable-reference)
                  (bracket delimiter misc-punctuation operator)))
    (setq-local treesit-font-lock-settings
                zsh-ts-mode--treesit-settings)
    (treesit-major-mode-setup)))

(provide 'zsh-ts-mode)
;;; zsh-ts-mode.el ends here