;;; nu-ts-mode.el --- Major mode for Nushell scripts, powered by treesitter -*- lexical-binding: t; -*-
;;
;; Copyright (C) 2026
;;
;; Author:  Oleksandr Makhmudov <oleksandr.makhmudov@proton.me>
;; Maintainer:  Oleksandr Makhmudov <oleksandr.makhmudov@proton.me>
;; Created: January 24, 2026
;; Modified: January 24, 2026
;; Package-Version: 20260721.1414
;; Package-Revision: 088c09c86bb1
;; Keywords: nushell nu languages tree-sitter
;; Homepage: https://github.com/ItsOleks/nu-ts-mode
;; Package-Requires: ((emacs "29.1"))
;;
;; This file is not part of GNU Emacs.

;; SPDX-License-Identifier: GPL-3.0-or-later
;;
;; nu-ts-mode is free software: you can redistribute it and/or modify
;; it under the terms of the GNU General Public License as published
;; by the Free Software Foundation, either version 3 of the License,
;; or (at your option) any later version.
;;
;; nu-ts-mode is distributed in the hope that it will be useful,
;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
;; GNU General Public License for more details.
;;
;; You should have received a copy of the GNU General Public License
;; along with nu-ts-mode.  If not, see <https://www.gnu.org/licenses/>.

;;; Commentary:
;; Tree-sitter powered major mode for nushell
;;
;; Current features:
;;   Syntax highlighting
;;   Indentation
;;
;;; Code:

(require 'treesit)

(unless (treesit-available-p)
  (error "`nu-ts-mode` requires Emacs to be built with tree-sitter support"))

(defgroup nu-ts nil
  "Major mode for editing Nushell scripts."
  :prefix "nu-ts-"
  :group 'languages)

(defcustom nu-ts-mode-indent-offset 2
  "Number of spaces for each indentation step in `nu-ts-mode'."
  :type 'integer
  :safe 'integerp)
(defvar nu-ts--treesit-operators
  '("+" "-" "*" "/" "//" "++" "**" "==" "!=" "<" "<=" ">" ">=" "=~" "!~"
    "=" "+=" "-=" "*=" "/=" "++=" ".." "..=" "..<" "=>" "|" "o>" "out>" "e>"
    "err>" "e+o>" "err+out>" "o+e>" "out+err>" "o>>" "out>>" "e>>" "err>>"
    "e+o>>" "err+out>>" "o+e>>" "out+err>>" "e>|" "err>|" "e+o>|" "err+out>|"
    "o+e>|" "out+err>|"
    "mod" "and" "or" "xor" "bit-or" "bit-xor" "bit-and" "bit-shl" "bit-shr"
    "in" "not-in" "has" "not-has" "starts-with" "ends-with" "not"))

(defvar nu-ts--treesit-builtins
  '("all" "ansi" "any" "append" "ast" "bits" "bytes" "cal" "cd" "char" "chunk-by" "chunks" "clear"
    "collect" "columns" "compact" "complete" "config" "cp" "date" "debug" "decode" "default"
    "detect" "drop" "du" "each" "encode" "enumerate" "every" "exec" "exit" "explain" "explore"
    "fill" "filter" "find" "first" "flatten" "format" "from" "generate" "get" "glob" "grid"
    "group-by" "hash" "headers" "histogram" "history" "http" "input" "insert" "inspect" "interleave"
    "into" "is-empty" "is-not-empty" "is-terminal" "items" "job" "join" "keybindings" "kill" "last"
    "length" "let-env" "lines" "load-env" "ls" "math" "merge" "metadata" "mkdir" "mktemp" "move"
    "mv" "nu-check" "nu-highlight" "open" "panic" "par-each" "parse" "path" "plugin" "port"
    "prepend" "print" "ps" "query" "random" "reduce" "reject" "rename" "reverse" "rm" "roll"
    "rotate" "run-external" "save" "schema" "select" "seq" "shuffle" "skip" "sleep" "slice" "sort"
    "sort-by" "split" "start" "stor" "str" "sys" "table" "take" "tee" "term" "timeit" "to" "touch"
    "transpose" "tutor" "ulimit" "uname" "uniq" "uniq-by" "update" "upsert" "url" "values" "version"
    "view" "watch" "which" "whoami" "window" "with-env" "wrap" "zip"))

(defvar nu-ts--treesit-keywords
  '("let" "mut" "const" "if" "else" "match" "loop"
    "while" "try" "catch" "error" "module" "use"
    "alias" "export-env" "export" "extern"))

(defvar nu-ts--treesit-non-node-keywords
  '("do" "source" "source-env" "hide" "hide-env" "break" "continue" "return"))


(defvar nu-ts-mode--font-lock-settings
  (treesit-font-lock-rules
   :language 'nu
   :feature 'comment
   '((comment) @font-lock-comment-face
     (shebang) @font-lock-comment-face
     ((comment) :+ @font-lock-doc-face :anchor (decl_def))
     (parameter (comment) @font-lock-doc-face))

   :language 'nu
   :feature 'definition
   '((decl_def (cmd_identifier) @font-lock-function-name-face)
     (scope_pattern (wild_card) @font-lock-function-name-face))

   :language 'nu
   :feature 'type
   :override t
   '((flat_type) @font-lock-type-face
     (list_type "list" @font-lock-builtin-face)
     (collection_type ["record" "table"] @font-lock-builtin-face)
     (collection_type key: (_) @font-lock-property-name-face)
     (composite_type "oneof" @font-lock-builtin-face))

   :language 'nu
   :feature 'keyword
   `("def" @font-lock-keyword-face
     [,@nu-ts--treesit-keywords] @font-lock-keyword-face
     (ctrl_for ["for" "in"] @font-lock-keyword-face)
     (command (cmd_identifier) @font-lock-keyword-face
              (:match
               ,(rx-to-string
                 `(seq bol (or ,@nu-ts--treesit-non-node-keywords) eol))
               @font-lock-keyword-face))
     (command (cmd_identifier) @font-lock-keyword-face
              (:match
               ,(rx-to-string
                 `(seq bol (or "error" "overlay") eol))
               @font-lock-keyword-face))
     (command (cmd_identifier) @_cmd :anchor
              (val_string) @font-lock-keyword-face
              (:match
               ,(rx-to-string
                 `(seq bol (or "error" "overlay") eol))
               @_cmd))
     (command (cmd_identifier) @_cmd
              arg_str: (val_string) @font-lock-keyword-face
              (:equal @_cmd "overlay")
              (:equal @font-lock-keyword-face "as")))

   :language 'nu
   :feature 'string
   '((val_string) @font-lock-string-face
     (inter_escape_sequence) @font-lock-escape-face
     (escape_sequence) @font-lock-escape-face
     (val_interpolated ["$\"" "$'" "\"" "'"] @font-lock-string-face)
     (unescaped_interpolated_content) @font-lock-string-face
     (escaped_interpolated_content) @font-lock-string-face)

   :language 'nu
   :feature 'literal
   :override t
   '((val_number) @font-lock-number-face
     (val_duration unit: (_) @font-lock-constant-face)
     (val_filesize unit: (_) @font-lock-constant-face)
     (val_binary digit: [(_) "," @font-lock-delimiter-face]) @font-lock-number-face
     (val_bool) @font-lock-constant-face
     (val_nothing) @font-lock-constant-face
     (val_date) @font-lock-number-face)

   :language 'nu
   :feature 'variable
   '((param_rest name: (_) @font-lock-variable-name-face)
     (param_opt name: (_) @font-lock-variable-name-face)
     (parameter param_name: (_) @font-lock-variable-name-face)
     (stmt_let (identifier) @font-lock-variable-name-face)
     (val_variable
      "$" :? @font-lock-misc-punctuation-face
      "...$" :? @font-lock-misc-punctuation-face
      [(identifier) @font-lock-variable-use-face
       "in" @font-lock-builtin-face
       "nu" @font-lock-constant-face
       "env" @font-lock-constant-face
       ]))

   :language 'nu
   :feature 'function
   :override t
   '((cmd_identifier) @font-lock-function-call-face
     (command "^" @font-lock-delimiter-face
              head: (_) @font-lock-function-call-face))

   :language 'nu
   :feature 'builtin
   :override t
   `("where" @font-lock-builtin-face
     (command (cmd_identifier) @font-lock-builtin-face
              (:match
               ,(rx-to-string
                 `(seq bol (or ,@nu-ts--treesit-builtins) eol))
               @font-lock-builtin-face)))

   :language 'nu
   :feature 'identifier
   :override t
   '(key: (identifier) @font-lock-property-use-face
     (param_completer (cmd_identifier) @font-lock-string-face)
     (attribute "@" @font-lock-preprocessor-face)
     [(attribute_identifier) (long_flag_identifier)
      (param_short_flag_identifier) (short_flag_identifier)]
     @font-lock-preprocessor-face
     (path ["." "?" "!"] @font-lock-misc-punctuation-face) @font-lock-property-use-face)

   :language 'nu
   :feature 'operator
   `([,@nu-ts--treesit-operators] @font-lock-operator-face)

   :language 'nu
   :feature 'punctuation
   :override t
   '(["," ";" ":"] @font-lock-delimiter-face
     (param_long_flag "--" @font-lock-delimiter-face)
     (long_flag ["=" "--"] @font-lock-delimiter-face)
     (short_flag ["=" "-"] @font-lock-delimiter-face)
     (param_short_flag "-" @font-lock-delimiter-face)
     (param_rest "..." @font-lock-delimiter-face)
     (param_value "=" @font-lock-delimiter-face)
     (param_completer "@" @font-lock-defimiter-face)
     (param_opt "?" @font-lock-defimiter-face)
     (returns "->" @font-lock-defimiter-face)
     ["(" ")" "{" "}" "[" "]" "<" ">" "...[" "...(" "...{"] @font-lock-bracket-face
     (parameter_pipes "|" @font-lock-bracket-face))

   :language 'nu
   :feature 'error
   :override t
   '((ERROR) @font-lock-warning-face))
  "Tree-sitter font-lock settings for `nu-ts-mode'.")

;; Indentation
(defvar nu-ts-mode-indent-rules
  '((nu
     ((node-is "]") parent-bol 0)
     ((node-is ")") parent-bol 0)
     ((node-is "}") parent-bol 0)

     ((parent-is "expr_parenthesized") parent-bol nu-ts-mode-indent-offset)
     ((parent-is "parameter_bracks") parent-bol nu-ts-mode-indent-offset)
     ((parent-is "ctrl_match") parent-bol nu-ts-mode-indent-offset)
     ((parent-is "val_record") parent-bol nu-ts-mode-indent-offset)
     ((parent-is "val_list") parent-bol nu-ts-mode-indent-offset)
     ((parent-is "val_closure") parent-bol nu-ts-mode-indent-offset)
     ((parent-is "val_table") parent-bol nu-ts-mode-indent-offset)
     ((parent-is "block") parent-bol nu-ts-mode-indent-offset)

     ((parent-is "pipeline") parent 0)))
  "Tree-sitter indent rules for `nu-ts-mode'.")


;; Keymap
(defvar nu-ts-mode-map
  (let ((map (make-sparse-keymap)))
    map)
  "Keymap for `nu-ts-mode'.")

;; Syntax map
(defvar nu-ts-mode--syntax-table
  (let ((table (make-syntax-table)))
    (modify-syntax-entry ?# "<" table)
    (modify-syntax-entry ?\n ">#" table)
    table)
  "Syntax table for `nu-ts-mode'.")

;;;###autoload
(define-derived-mode nu-ts-mode prog-mode "Nushell"
  "Major mode for editing Nushell scripts, powered by treesitter.

\\{nu-ts-mode-map}"
  :syntax-table nu-ts-mode--syntax-table

  (when (treesit-ready-p 'nu)
    (treesit-parser-create 'nu)

    ;; Font locking
    (setq-local treesit-font-lock-settings nu-ts-mode--font-lock-settings)

    (setq-local treesit-font-lock-feature-list
                '((comment definition)
                  (keyword type string)
                  (variable literal)
                  (function builtin identifier operator punctuation error)))

    ;; Comments
    (setq-local comment-start "# ")
    (setq-local comment-start-skip "#+\\s-*")

    ;; Indentation
    (setq-local treesit-simple-indent-rules nu-ts-mode-indent-rules)

    (treesit-major-mode-setup)))

(when (treesit-ready-p 'nu)
  (add-to-list 'auto-mode-alist '("\\.nu\\'" . nu-ts-mode))
  (add-to-list 'interpreter-mode-alist '("nu" . nu-ts-mode)))

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