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

(require 'cl-lib)
(require 'eieio)
(require 'sh-script)
(require 'treesit)
(require 'rx)

;;; #################
;;; # Fontification #
;;; #################
(defvar zsh-ts-mode--previous-let-arg-query
  (treesit-query-compile 'zsh '((word) @arg))
  "Query used in `zsh-ts-mode--fontify-let-arg'.")

(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)
      (cond
       ((re-search-forward (rx (or "[" "=")) ne t)
        (put-text-property ns (1- (point)) 'face
                           'font-lock-variable-name-face)
        (unless (eql (char-before) ?\[)
          (ignore-errors
            (with-restriction (point) ne
              (when (looking-at (rx (? "-") (+ (any (?0 . ?9)))))
                (put-text-property (point-min) (point-max) 'face
                                   'font-lock-number-face))))))
       ;; the parser is broken and parses:
       ;; let ans=${?}
       ;; as
       ;; command:let word:ans= word:${}
       ((null (treesit-query-capture (treesit-node-parent node)
                                     zsh-ts-mode--previous-let-arg-query
                                     (1- ns) ns))
        (font-lock-append-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 nil nil t))
    (let ((start (treesit-node-start child))
          (end (treesit-node-end 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.")

(defvar zsh-ts-mode--treesit-keywords
  (append sh-mode--treesit-keywords '("until"))
  "List of keywords defined in the grammar.")

;; 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_pattern (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
     [ ,@zsh-ts-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)))

;;; #########
;;; # IMenu #
;;; #########
(defvar zsh-ts-mode--defun-name-query
  (treesit-query-compile 'zsh '((function_definition
                                 name: ((word) @name))))
  "Query used by `zsh-ts-mode--defun-name'.")

(defun zsh-ts-mode--defun-name (node)
  "If NODE is a defun node, return its name.
This behaves according to `treesit-defun-name-function'."
  (when-let ((node)
             (res (treesit-query-capture node zsh-ts-mode--defun-name-query
                                         nil nil t)))
    (treesit-node-text (car res) t)))

(defun zsh-ts-mode--imenu-predicate (node)
  "Predicate function for `zsh-ts-mode--simple-imenu-settings'.
Return non-nil if NODE is not an anonymous function."
  (zsh-ts-mode--defun-name node))

(defvar zsh-ts-mode--simple-imenu-settings
  '(("Function" "function_definition" zsh-ts-mode--imenu-predicate nil))
  "`zsh-ts-mode' settings for `treesit-simple-imenu-settings'.")

;;; ########
;;; # xref #
;;; ########
(defun zsh-ts-mode--xref-backend ()
  "Return `zsh-ts'."
  'zsh-ts)

(defun zsh-ts-mode--not-let-predicate (node)
  "Return nil if NODE is a let node, non-nil otherwise."
  (not (equal (treesit-node-text node) "let")))

(defvar zsh-ts-mode--identifier-query
  (treesit-query-compile
   'zsh `((expansion
           name: ([(simple_variable_name)
                   (special_variable_name)]
                  @var))
          (expansion_with_modifier
           name: ([(simple_variable_name)
                   (special_variable_name)]
                  @var))
          (variable_ref ([(simple_variable_name)
                          (special_variable_name)]
                         @var))
          (expansion_pattern
           name: ([(simple_variable_name)
                   (special_variable_name)]
                  @var))
          ((word) @func
           (:pred zsh-ts-mode--not-let-predicate @func))))
  "Query for `zsh-ts-mode' `xref-backend-identifier-at-point'.")

(cl-defmethod xref-backend-identifier-at-point ((_backend (eql zsh-ts)))
  "`zsh-ts-mode' implementation for xref's identifier at point function."

  (or (when (memq 'font-lock-variable-name-face
                  (ensure-list (get-text-property (point) 'face)))
        (let ((start (previous-single-property-change (point) 'face))
              (end (next-single-property-change (point) 'face)))
          (propertize
           (buffer-substring-no-properties (if start (1+ start) (point-min))
                                           (or end (point-max)))
           'type 'variable)))
      (when-let* ((nodes (treesit-query-capture (treesit-buffer-root-node)
                                                zsh-ts-mode--identifier-query
                                                (point) (1+ (point)))))
        (if-let* ((var (alist-get 'var nodes)))
            (propertize (treesit-node-text var t) 'type 'variable)
          (when-let* ((func (alist-get 'func nodes)))
            (propertize (treesit-node-text func t) 'type 'function))))))

(cl-defmethod xref-backend-identifier-completion-table ((_backend (eql zsh-ts)))
  "Return nil."
  nil)

(cl-defmethod xref-backend-identifier-completion-ignore-case
  ((_backend (eql zsh-ts)))
  "Return `completion-ignore-case'."
  completion-ignore-case)

(defun zsh-ts-mode--treesit-node-summary (node)
  "Return the contents of the line of NODE."
  (save-excursion
    (goto-char (treesit-node-start node))
    (buffer-substring (pos-bol) (pos-eol))))

(cl-defstruct zsh-ts-mode--location
  "Xref locaton object for tree-sitter nodes."
  node)

(cl-defmethod xref-location-line ((location zsh-ts-mode--location))
  "Return the line number of LOCATION."
  (line-number-at-pos (treesit-node-start
                       (zsh-ts-mode--location-node location))))

(cl-defmethod xref-location-group ((location zsh-ts-mode--location))
  "Return the file name of the buffer of LOCATION."
  (with-slots (node) location
    (or (buffer-file-name (treesit-node-buffer node))
        "")))

(cl-defmethod xref-location-marker ((location zsh-ts-mode--location))
  "Return a marker for LOCATION."
  (let ((node (zsh-ts-mode--location-node location))
        (marker (make-marker)))
    (set-marker marker (treesit-node-start node)
                (treesit-node-buffer node))
    marker))

(defun zsh-ts-mode--make-xref-for-node (node)
  "Create an xref for a tree-sitter node NODE."
  (xref-make (zsh-ts-mode--treesit-node-summary node)
             (make-zsh-ts-mode--location :node node)))

(defvar zsh-ts-mode--xref-function-definitions-query
  (treesit-query-compile 'zsh '((function_definition
                                 name: (word) @name)))
  "Function query for `zsh-ts-mode--xref-function-definitions'.")

(defun zsh-ts-mode--xref-function-definitions (identifier)
  "Find all definitions of function IDENTIFIER in the current buffer."
  (cl-loop for node in (treesit-query-capture
                        (treesit-buffer-root-node)
                        zsh-ts-mode--xref-function-definitions-query
                        nil nil t)
           when (equal identifier (treesit-node-text node))
           collect (zsh-ts-mode--make-xref-for-node node)))

(defvar zsh-ts-mode--xref-variable-definitions-query
  (treesit-query-compile
   'zsh `((variable_assignment name: (variable_name) @non-let)
          (declaration_command argument: (word) @non-let)
          (command name: ((command_name) @let.name
                          (:equal @let.name "let"))
                   argument: ((word) @let-arg
                              (:match ,(rx-to-string '(seq bos (not "-")))
                                      @let-arg)))))
  "Variable query for `zsh-ts-mode--xref-variable-definitions'.")

(defun zsh-ts-mode--xref-variable-definitions (identifier)
  "Find all definitions of variable IDENTIFIER in the current buffer."
  (let (out)
    (dolist (ent (treesit-query-capture
                  (treesit-buffer-root-node)
                  zsh-ts-mode--xref-variable-definitions-query))
      (cl-destructuring-bind (tag . node) ent
        (cl-case tag
          (non-let
           (when (equal (treesit-node-text node) identifier)
             (push (zsh-ts-mode--make-xref-for-node node) out)))
          (let-arg
           (when (string-match-p (rx bos (literal identifier) (or eos "="))
                                 (treesit-node-text node))
             (push (zsh-ts-mode--make-xref-for-node node) out))))))
    out))

(cl-defmethod xref-backend-definitions ((_backend (eql zsh-ts)) identifier)
  "`zsh-ts-mode' implementation for finding xref definitions for IDENTIFIER."
  (when  identifier
    (let ((type (get-text-property 0 'type identifier)))
      (append
       (when (memq type '(nil function))
         (zsh-ts-mode--xref-function-definitions identifier))
       (when (memq type '(nil variable))
         (zsh-ts-mode--xref-variable-definitions identifier))))))

(defvar zsh-ts-mode--function-references-query
  (treesit-query-compile
   'zsh '((command name: (command_name (word) @name))))
  "Query used by `zsh-ts-mode--xref-function-references'.")

(defun zsh-ts-mode--xref-function-references (identifier)
  "Find all references for the function IDENTIFIER in the current buffer."
  (let ((nodes (treesit-query-capture
                (treesit-buffer-root-node)
                zsh-ts-mode--function-references-query
                nil nil t))
        (ident-copy (copy-sequence identifier))
        (quote-ident (regexp-quote identifier)))
    (set-text-properties 0 (length ident-copy) () ident-copy)
    (cl-loop for node in nodes
             when (save-excursion
                    (goto-char (treesit-node-start node))
                    (looking-at-p quote-ident))
             collect (zsh-ts-mode--make-xref-for-node node))))

(defvar zsh-ts-mode--variable-references-query
  (treesit-query-compile
   'zsh '(([(simple_variable_name) (special_variable_name)] @var)))
  "Query used by `zsh-ts-mode--xref-variable-references'.")

(defun zsh-ts-mode--xref-variable-references (identifier)
  "Find all references for the variable IDENTIFIER in the current buffer."
  (let ((ident-copy (copy-sequence identifier))
        (quote-ident (regexp-quote identifier)))
    (set-text-properties 0 (length ident-copy) () ident-copy)
    (mapcan
     #'(lambda (node)
         (when (save-excursion
                 (goto-char (treesit-node-start node))
                 (looking-at-p quote-ident))
           (list (zsh-ts-mode--make-xref-for-node node))))
     (treesit-query-capture (treesit-buffer-root-node)
                            zsh-ts-mode--variable-references-query
                            nil nil t))))

(cl-defmethod xref-backend-references ((_backend (eql zsh-ts)) identifier)
  "`zsh-ts-mode' implementation for finding xref references for IDENTIFIER."
  (let ((type (get-text-property 0 'type identifier)))
    (append
     ;; definitions are references
     (ensure-list (xref-backend-definitions 'zsh-ts identifier))
     (when (memq type '(nil function))
       (zsh-ts-mode--xref-function-references identifier))
     (when (memq type '(nil variable))
       (zsh-ts-mode--xref-variable-references identifier)))))

(defvar zsh-ts-mode--xref-apropos-query
  (treesit-query-compile
   'zsh '(((simple_variable_name) @ident)
          ((special_variable_name) @ident)
          ((variable_name) @ident)
          (command name: (command_name ( (word) @cname (:equal @cname "let")))
                   argument: (word) @ident)
          (declaration_command argument: (word) @ident)
          ((variable_name) @ident)
          (function_definition name: (word) @ident)))
  "Query used by `xref-backend-apropos'.")

(cl-defmethod xref-backend-apropos ((_backend (eql zsh-ts)) pattern)
  "Search though the buffer for PATTERN."
  (let (out)
    (dolist (ent (treesit-query-capture (treesit-buffer-root-node)
                                        zsh-ts-mode--xref-apropos-query))
      (cl-destructuring-bind (type . node) ent
        (when (and (eq type 'ident)
                   (string-match-p pattern (treesit-node-text node)))
          (push (zsh-ts-mode--make-xref-for-node node) out))))
    out))

;;;###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-settings zsh-ts-mode--treesit-settings
                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-simple-imenu-settings zsh-ts-mode--simple-imenu-settings
                treesit-defun-name-function #'zsh-ts-mode--defun-name
                imenu-create-index-function #'treesit-simple-imenu)
    (add-hook 'xref-backend-functions #'zsh-ts-mode--xref-backend nil t)
    (treesit-major-mode-setup)))

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