;;; ts-mode.el --- tree sitter support for TypeScript  -*- lexical-binding: t; -*-

;; Copyright (C) 2022 Free Software Foundation, Inc.

;; Author     : Theodor Thornhill <theo@thornhill.no>
;; Maintainer : Theodor Thornhill <theo@thornhill.no>
;; Created    : October 2022
;; Keywords   : typescript tsx languages tree-sitter

;; This file is part of GNU Emacs.

;; This program 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.

;; This program 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 this program.  If not, see <http://www.gnu.org/licenses/>.

(require 'treesit)
(require 'rx)
(require 'js)

(defcustom ts-mode-indent-offset 2
  "Number of spaces for each indentation step in `ts-mode'."
  :type 'integer
  :safe 'integerp
  :group 'typescript)

(defvar ts-mode--syntax-table
  (let ((table (make-syntax-table)))
    ;; Taken from the cc-langs version
    (modify-syntax-entry ?_  "_"     table)
    (modify-syntax-entry ?$ "_"      table)
    (modify-syntax-entry ?\\ "\\"    table)
    (modify-syntax-entry ?+  "."     table)
    (modify-syntax-entry ?-  "."     table)
    (modify-syntax-entry ?=  "."     table)
    (modify-syntax-entry ?%  "."     table)
    (modify-syntax-entry ?<  "."     table)
    (modify-syntax-entry ?>  "."     table)
    (modify-syntax-entry ?&  "."     table)
    (modify-syntax-entry ?|  "."     table)
    (modify-syntax-entry ?` "\""     table)
    (modify-syntax-entry ?\240 "."   table)
    table)
  "Syntax table for `ts-mode'.")

(defvar ts-mode--indent-rules
  `((tsx
     ((node-is "}") parent-bol 0)
     ((node-is ")") parent-bol 0)
     ((node-is "]") parent-bol 0)
     ((node-is ">") parent-bol 0)
     ((node-is ".")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "ternary_expression")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "named_imports")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "statement_block")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "type_arguments")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "variable_declarator")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "arguments")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "array")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "formal_parameters")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "template_substitution")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "object_pattern")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "object")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "object_type")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "enum_body")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "arrow_function")
      parent-bol ,ts-mode-indent-offset)
     ((parent-is "parenthesized_expression")
      parent-bol ,ts-mode-indent-offset)

     ;; TSX
     ((parent-is "jsx_opening_element")
      parent ,ts-mode-indent-offset)
     ((node-is "jsx_closing_element") parent 0)
     ((parent-is "jsx_element")
      parent ,ts-mode-indent-offset)
     ((node-is "/") parent 0)
     ((parent-is "jsx_self_closing_element")
      parent ,ts-mode-indent-offset)
     (no-node parent-bol 0))))

(defvar ts-mode--settings
  (treesit-font-lock-rules
   :language 'tsx
   :override t
   '(
     (template_string) @font-lock-string-face

     ((identifier) @font-lock-constant-face
      (:match "^[A-Z_][A-Z_\\d]*$" @font-lock-constant-face))

     (nested_type_identifier
      module: (identifier) @font-lock-type-face)
     (type_identifier) @font-lock-type-face
     (predefined_type) @font-lock-type-face

     (new_expression
      constructor: (identifier) @font-lock-type-face)

     (function
      name: (identifier) @font-lock-function-name-face)

     (function_declaration
      name: (identifier) @font-lock-function-name-face)

     (method_definition
      name: (property_identifier) @font-lock-function-name-face)

     (variable_declarator
      name: (identifier) @font-lock-function-name-face
      value: [(function) (arrow_function)])

     (variable_declarator
      name: (array_pattern
             (identifier)
             (identifier) @font-lock-function-name-face)
      value: (array (number) (function)))

     (assignment_expression
      left: [(identifier) @font-lock-function-name-face
             (member_expression
              property: (property_identifier) @font-lock-function-name-face)]
      right: [(function) (arrow_function)])

     (call_expression
      function:
      [(identifier) @font-lock-function-name-face
       (member_expression
        property: (property_identifier) @font-lock-function-name-face)])

     (variable_declarator
      name: (identifier) @font-lock-variable-name-face)

     (enum_declaration (identifier) @font-lock-type-face)

     (enum_body (property_identifier) @font-lock-type-face)

     (enum_assignment name: (property_identifier) @font-lock-type-face)

     (assignment_expression
      left: [(identifier) @font-lock-variable-name-face
             (member_expression
              property: (property_identifier) @font-lock-variable-name-face)])

     (for_in_statement
      left: (identifier) @font-lock-variable-name-face)

     (arrow_function
      parameter: (identifier) @font-lock-variable-name-face)

     (arrow_function
      parameters:
      [(_ (identifier) @font-lock-variable-name-face)
       (_ (_ (identifier) @font-lock-variable-name-face))
       (_ (_ (_ (identifier) @font-lock-variable-name-face)))])


     (pair key: (property_identifier) @font-lock-variable-name-face)

     (pair value: (identifier) @font-lock-variable-name-face)

     (pair
      key: (property_identifier) @font-lock-function-name-face
      value: [(function) (arrow_function)])

     (property_signature
      name: (property_identifier) @font-lock-variable-name-face)

     ((shorthand_property_identifier) @font-lock-variable-name-face)

     (pair_pattern
      key: (property_identifier) @font-lock-variable-name-face)

     ((shorthand_property_identifier_pattern)
      @font-lock-variable-name-face)

     (array_pattern (identifier) @font-lock-variable-name-face)

     (jsx_opening_element
      [(nested_identifier (identifier)) (identifier)]
      @font-lock-function-name-face)

     (jsx_closing_element
      [(nested_identifier (identifier)) (identifier)]
      @font-lock-function-name-face)

     (jsx_self_closing_element
      [(nested_identifier (identifier)) (identifier)]
      @font-lock-function-name-face)

     (jsx_attribute (property_identifier) @font-lock-constant-face)

     [(this) (super)] @font-lock-keyword-face

     [(true) (false) (null)] @font-lock-constant-face
     (regex pattern: (regex_pattern)) @font-lock-string-face
     (number) @font-lock-constant-face

     (string) @font-lock-string-face
     (template_string) @font-lock-string-face

     (template_substitution
      ["${" "}"] @font-lock-constant-face)

     ["!"
      "abstract"
      "as"
      "async"
      "await"
      "break"
      "case"
      "catch"
      "class"
      "const"
      "continue"
      "debugger"
      "declare"
      "default"
      "delete"
      "do"
      "else"
      "enum"
      "export"
      "extends"
      "finally"
      "for"
      "from"
      "function"
      "get"
      "if"
      "implements"
      "import"
      "in"
      "instanceof"
      "interface"
      "keyof"
      "let"
      "namespace"
      "new"
      "of"
      "private"
      "protected"
      "public"
      "readonly"
      "return"
      "set"
      "static"
      "switch"
      "target"
      "throw"
      "try"
      "type"
      "typeof"
      "var"
      "void"
      "while"
      "with"
      "yield"
      ] @font-lock-keyword-face

     (comment) @font-lock-comment-face
     )))

(defvar ts-mode--defun-type-regexp
  (rx (or "class_declaration"
          "method_definition"
          "function_declaration"
          "lexical_declaration"))
  "Regular expression that matches type of defun nodes.
Used in `ts-mode--beginning-of-defun' and friends.")

(defun ts-mode--beginning-of-defun (&optional arg)
  "Tree-sitter `beginning-of-defun' function.
ARG is the same as in `beginning-of-defun."
  (let ((arg (or arg 1)))
    (if (> arg 0)
        ;; Go backward.
        (while (and (> arg 0)
                    (treesit-search-forward-goto
                     ts-mode--defun-type-regexp 'start nil t))
          (setq arg (1- arg)))
      ;; Go forward.
      (while (and (< arg 0)
                  (treesit-search-forward-goto
                   ts-mode--defun-type-regexp 'start))
        (setq arg (1+ arg))))))

(defun ts-mode--end-of-defun (&optional arg)
  "Tree-sitter `end-of-defun' function.
ARG is the same as in `end-of-defun."
  (let ((arg (or arg 1)))
    (if (< arg 0)
        ;; Go backward.
        (while (and (< arg 0)
                    (treesit-search-forward-goto
                     ts-mode--defun-type-regexp 'end nil t))
          (setq arg (1+ arg)))
      ;; Go forward.
      (while (and (> arg 0)
                  (treesit-search-forward-goto
                   ts-mode--defun-type-regexp 'end))
        (setq arg (1- arg))))))

;;;###autoload
(add-to-list 'auto-mode-alist '("\\.ts\\'" . ts-mode))

;;;###autoload
(add-to-list 'auto-mode-alist '("\\.tsx\\'" . ts-mode))

;;;###autoload
(define-derived-mode ts-mode prog-mode "TypeScript"
  "Major mode for editing TypeScript."
  :group 'typescript
  :syntax-table ts-mode--syntax-table

  (cond
   ((and (treesit-can-enable-p)
         (treesit-language-available-p 'tsx))
    ;; Comments
    (setq-local comment-start "// ")
    (setq-local comment-start-skip "\\(?://+\\|/\\*+\\)\\s *")
    (setq-local comment-end "")

    (setq-local treesit-simple-indent-rules ts-mode--indent-rules)
    (setq-local indent-line-function #'treesit-indent)

    (setq-local beginning-of-defun-function #'ts-mode--beginning-of-defun)
    (setq-local end-of-defun-function #'ts-mode--end-of-defun)

    (unless font-lock-defaults
      (setq font-lock-defaults '(nil t)))

    (setq-local treesit-font-lock-settings ts-mode--settings)

    (treesit-font-lock-enable))
   (t
    (message "Tree sitter for TypeScript isn't available, defaulting to js-mode")
    (js-mode))))

(provide 'ts-mode)

;;; ts-mode.el ends here