;;; 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_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))
          ((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))))))

(defun zsh-ts-mode--treesit-node-location (node)
  "Return an xref location pointing to NODE."
  (save-excursion
    (goto-char (treesit-node-start node))
    (make-xref-file-location :file (buffer-file-name)
                             :line (line-number-at-pos)
                             :column (- (point) (pos-bol)))))

(defun zsh-ts-mode--make-function-definition-xref (identifier node)
  "Make an xref object pointing to IDENTIFIER.
NODE is IDENTIFIER's tree-sitter node."
  (xref-make (format "%s %s()"
                     (propertize "function"
                                 'face 'font-lock-keyword-face)
                     (propertize identifier
                                 'face 'font-lock-function-name-face))
             (zsh-ts-mode--treesit-node-location 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-function-definition-xref identifier node)
           into out
           finally return (if (length= out 1)
                              (car out)
                            out)))

(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'.")

(defvar zsh-ts-mode--non-let-declaration-type-query
  (treesit-query-compile 'zsh '((declaration_command :anchor _ @name)))
  "Helper query for `zsh-ts-mode--get-non-let-assignment-type'.")

(defun zsh-ts-mode--get-assignment-parent (node)
  "Get the node that is a parent of NODE that is a variable assignment.
This only looks up two levels."
  (cond ((equal (treesit-node-type
                 (treesit-node-parent node))
                "declaration_command")
         (treesit-node-parent node))
        ((equal (treesit-node-type
                 (treesit-node-parent (treesit-node-parent node)))
                "declaration_command")
         (treesit-node-parent (treesit-node-parent node)))))

(defun zsh-ts-mode--get-non-let-assignment-type (node)
  "Return the assignment command for NODE."
  (when-let* ((type-node (zsh-ts-mode--get-assignment-parent node)))
    (propertize (treesit-node-text
                 (car (treesit-query-capture
                       type-node
                       zsh-ts-mode--non-let-declaration-type-query
                       nil nil t))
                 t)
                'face 'font-lock-keyword-face)))

(defun zsh-ts-mode--make-xref-for-non-let-assignment (node)
  "Create an xref for an \"other\" assignment for NODE.
An \"other\" assignment is an assignment using local, typeset, etc."
  (let ((type (zsh-ts-mode--get-non-let-assignment-type node)))
    (xref-make (format "%s%s" (if type (concat type " ") "")
                       (propertize (treesit-node-text node t)
                                   'face 'font-lock-variable-name-face))
               (zsh-ts-mode--treesit-node-location node))))

(defun zsh-ts-mode--make-xref-for-let-assignment (node)
  "Create an xref node for a list assingment.
NODE is the argument to let (possibly with a = in it)."
  (save-excursion
    (goto-char (treesit-node-start node))
    (xref-make
     (format "%s %s" (propertize "let" 'face 'font-lock-keyword-face)
             (propertize (buffer-substring-no-properties
                          (treesit-node-start node)
                          (if (re-search-forward (rx (or "=" "["))
                                                 (treesit-node-end node)
                                                 t)
                              (1- (point))
                            (treesit-node-end node)))))
     (zsh-ts-mode--treesit-node-location node))))

(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-non-let-assignment 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-let-assignment node) out))))))
    (if (length= out 1) (car out) out)))

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

;;;###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)
    (treesit-major-mode-setup)))

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