;;; haskell-ts-mode.el --- A treesit based major mode for haskell -*- lexical-binding:t -*-

;; Copyright (C) 2024  Pranshu Sharma


;; Author: Pranshu Sharma <pranshusharma366 at gmail>
;; URL: https://codeberg.org/pranshu/haskell-ts-mode
;; Package-Requires: ((emacs "29.3"))
;; Version: 1
;; Keywords: languages, haskell

;; 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 <https://www.gnu.org/licenses/>.

;;; Commentary:

;; This is a major mode that uses treesitter to provide all the basic
;; major mode stuff, like indentation, font lock, etc...

;;; Code:

(require 'comint)
(require 'treesit)

(declare-function treesit-parser-create "treesit.c")
(declare-function treesit-node-start "treesit.c")
(declare-function treesit-node-parent "treesit.c")
(declare-function treesit-node-prev-sibling "treesit.c")
(declare-function treesit-node-next-sibling "treesit.c")
(declare-function treesit-node-end "treesit.c")
(declare-function treesit-node-child "treesit.c")
(declare-function treesit-node-type "treesit.c")

(defgroup haskell-ts-mode nil
  "Group that contains haskell-ts-mode variables"
  :group 'langs)

(defvar haskell-ts-font-lock-feature-list
  `((comment str pragma parens)
    (type definition function args)
    (match keyword)
    (otherwise signature type-sig)))

(defcustom haskell-ts-use-indent t
  "Set to non-nil to use the indentation provided by haskell-ts-mode"
  :group 'haskell-ts-mode
  :type 'boolean)

(defcustom haskell-ts-font-lock-level 4
  "Level of font lock, 1 for minimum highlghting and 4 for maximum."
  :group 'haskell-ts-mode
  :type 'integer)

(defvar haskell-ts-prettify-symbols-alist
  '(("\\" . "λ")
    ("/=" . "≠")
    ("->" . "→")
    ("=>" . "⇒")
    ("<-" . "←")
    ("<=" . "≥")
    (">=" . "≤")))

(defvar haskell-ts-font-lock
  (treesit-font-lock-rules
   :language 'haskell
   :feature 'keyword
   `(["module" "import" "data" "let" "where" "case" "type"
      "if" "then" "else" "of" "do" "in" "instance" "class"]
     @font-lock-keyword-face)
   :language 'haskell
   :feature 'otherwise
   :override t
   `(((match (guards guard: (boolean (variable) @font-lock-keyword-face)))
      (:match "otherwise" @font-lock-keyword-face)))
   :language 'haskell
   :feature 'type-sig
   "(signature (binding_list (variable) @font-lock-doc-markup-face))
    (signature (variable) @font-lock-doc-markup-face)"
   :language 'haskell
   :feature 'args
   :override 'keep
   (concat
    "(function (infix left_operand: (_) @haskell-ts--fontify-arg))"
    "(function (infix right_operand: (_) @haskell-ts--fontify-arg))"
    "(generator . (_) @haskell-ts--fontify-arg)"
    "(bind (as (variable) . (_) @haskell-ts--fontify-arg))"
    "(patterns) @haskell-ts--fontify-arg")
   :language 'haskell
   :feature 'type
   `((type) @font-lock-type-face
     (constructor) @font-lock-type-face)
   :language 'haskell
   :override t
   :feature 'signature
   `((signature (function) @haskell-ts--fontify-type)
     (context (function) @haskell-ts--fontify-type))
   :language 'haskell
   :feature 'match
   `((match ("|" @font-lock-doc-face) ("=" @font-lock-doc-face))
     (list_comprehension ("|" @font-lock-doc-face
			  (qualifiers (generator "<-" @font-lock-doc-face))))
     (match ("->" @font-lock-doc-face)))
   :language 'haskell
   :feature 'comment
   `(((comment) @font-lock-comment-face)
     ((haddock) @font-lock-doc-face))
   :language 'haskell
   :feature 'pragma
   `((pragma) @font-lock-preprocessor-face
     (cpp) @font-lock-preprocessor-face)
   :language 'haskell
   :feature 'str
   :override t
   `((char) @font-lock-string-face
     (string) @font-lock-string-face
     (quasiquote (quoter) @font-lock-type-face)
     (quasiquote (quasiquote_body) @font-lock-preprocessor-face))
   :language 'haskell
   :feature 'parens
   :override t
   `(["(" ")" "[" "]"] @font-lock-operator-face
     (infix operator: (_) @font-lock-operator-face))
   :language 'haskell
   :feature 'function
   :override t
   `((function name: (variable) @font-lock-function-name-face)
     (function (infix (operator)  @font-lock-function-name-face))
     (declarations (type_synomym (name) @font-lock-function-name-face))
     (bind (variable) @font-lock-function-name-face)
     (function (infix (infix_id (variable) @font-lock-function-name-face)))
     (bind (as (variable) @font-lock-function-name-face))))
  "A function that returns the treesit font lock lock settings for haskell.")

(defun haskell-ts--stand-alone-parent (_ parent bol)
  (save-excursion
    (goto-char (treesit-node-start parent))
    (let ((type (treesit-node-type parent)))
      (if (and (not bol)
	       (or (looking-back "^[ \t]*" (line-beginning-position))
		   (seq-some
		    (lambda (kw) 
		      (string= type kw))
		    '("when" "where" "do" "let" "local_binds" "function"))))
	  (treesit-node-start parent)
	(haskell-ts--stand-alone-parent 1 (funcall
					   (if bol 'treesit-node-parent 'identity)
					   (treesit-node-parent parent))
					nil)))))

(defvar haskell-ts--ignore-types
  (regexp-opt '("comment" "cpp" "haddock"))
  "Node types that will be ignored by indentation.")

(defvar haskell-ts-indent-rules
  (let* ((p-sib
	  (lambda (node arg)
	    (let* ((func (if arg
			     'treesit-node-prev-sibling
			   'treesit-node-next-sibling))
		   (n (funcall func	 node)))
	      (while (and n (string-match haskell-ts--ignore-types
					  (treesit-node-type n)))
		(setq n (funcall func n)))
	      n)))
	 (p-prev-sib
	  (lambda (node _ _) (treesit-node-start (funcall p-sib node t))))
	 (p-n-prev (lambda (node) (funcall p-sib node t)))
	 (parent-first-child (lambda (_ parent _)
			       (treesit-node-start (treesit-node-child parent 0)))))
    `((haskell
       ((node-is "^cpp$") column-0 0)
       ((parent-is "^comment$") column-0 0)
       ((parent-is "^haddock$") column-0 0)
       ((parent-is "^imports$") column-0 0)
       ;; Infix
       ((n-p-gp nil "infix" "infix")
	(lambda (_ node _)
	  (let ((first-inf nil))
	    (while (string= "infix"
			    (treesit-node-type
			     (setq node (treesit-node-parent node))))
	      (setq first-inf node))
	    (funcall ,parent-first-child nil first-inf nil)))
	0)
       ((node-is "^infix$") ,parent-first-child 0)
       
       ;; Lambda
       ((parent-is "^lambda\\(_case\\)?$") standalone-parent 2)

       ((parent-is "^class_declarations$") prev-sibling 0)

       ((node-is "^where$") parent 2)
       
       ;; in
       ((node-is "^in$") parent 0)
       
       ((parent-is "qualifiers") parent 0)
       
       ;; list
       ((node-is "^]$") parent 0)
       ((parent-is "^list$") standalone-parent 2)
       
       ;; If then else
       ((node-is "^then$") parent 2)
       ((node-is "^else$") parent 2)

       ((parent-is "^apply$") haskell-ts--stand-alone-parent 1)
       ((node-is "^quasiquote$") grand-parent 2)
       ((parent-is "^quasiquote_body$") (lambda (_ _ c) c) 0)
       ((lambda (node parent bol)
	  (when-let ((n (treesit-node-prev-sibling node)))
	    (while (string= "comment" (treesit-node-type n))
	      (setq n (treesit-node-prev-sibling n)))
	    (string= "do" (treesit-node-type n))))
	haskell-ts--stand-alone-parent
	3)
       ((parent-is "^do$") ,p-prev-sib 0)

       ((parent-is "^alternatives$") ,p-prev-sib 0)

       ;; prev-adaptive-prefix is broken sometimes
       (no-node
	(lambda (_ _ _)
	  (save-excursion
	    (goto-char (line-beginning-position 0))
	    (back-to-indentation)
	    (point)))
	0)
       
       ((parent-is "^data_constructors$") parent 0)

       ;; where
       ((lambda (node _ _)
	  (let ((n (treesit-node-prev-sibling node)))
	    (while (string= "comment" (treesit-node-type n))
	      (setq n (treesit-node-prev-sibling n)))
	    (string= "where" (treesit-node-type n))))
	
	(lambda (_ b _)
	  (+ 1 (treesit-node-start (treesit-node-prev-sibling b))))
	3)
       ((parent-is "local_binds\\|instance_declarations") ,p-prev-sib 0)
       
       ;; Match
       ((lambda (node _ _)
	  (and (string= "match" (treesit-node-type node))
	       (string-match (regexp-opt '("patterns" "variable"))
			     (treesit-node-type (funcall ,p-n-prev node)))))
	standalone-parent 2)
       
       ((node-is "match") ,p-prev-sib 0)
       ((parent-is "match") standalone-parent 2)
       
       ((parent-is "^haskell$") column-0 0)
       ((parent-is "^declarations$") column-0 0)
       
       ((parent-is "^record$") standalone-parent 2)

       ((parent-is "^exports$")
	(lambda (_ b _) (treesit-node-start (treesit-node-prev-sibling b)))
	0)
       ((n-p-gp nil "signature" "foreign_import") grand-parent 3)
       ((parent-is "^case$") standalone-parent 4)
       ((node-is "^alternatives$")
	(lambda (_ b _)
	  (treesit-node-start (treesit-node-child b 0)))
	2)
       ((node-is "^comment$")
	;; Indenting comments by priorites:
	;; 1. next relevent sibling if exists
	;; 2. previous relevent sibling if exists
	;; 3. parent
	;; (relevent means type not it haskell-ts--ignore-types)
	(lambda (node parent _)
	  (if-let ((next-sib (funcall ,p-sib node nil)))
	      (treesit-node-start next-sib)
	    (if-let ((prev-sib (funcall ,p-prev-sib node nil nil)))
		prev-sib
	      (treesit-node-start parent))))
	0)
       ;; Backup
       (catch-all parent 2)))))

;; Copied from haskell-tng-mode, changed a bit

(defvar haskell-ts-mode-syntax-table
      (let ((table (make-syntax-table)))
	;; The defaults are mostly fine
	(mapc
	 (lambda (ls)
	   (mapc
	    (lambda (char)
	      (modify-syntax-entry char (car ls) table))
	    (cdr ls)))
	 '(("_" ?! ?_)
	   ("w" ?')
	   ;; Haskell has some goofy comment enders like C-q C-l
	   (">" 13 10 12 11)
	   ("_ 123" ?-)
	   ("(}1nb" ?\{)
	   ("){4nb" ?\})
	   ("<" ?#)
	   (">" ?\n)
	   ;; Special operaters
	   ("." ?\, ?\; ?@)
	   ("\"" ?\")
	   ("$`"  ?\`)))
	table))

(defmacro haskell-ts-imenu-name-function (check-func)
  `(lambda (node)
     (let ((nn (treesit-node-child node 0 node)))
	 (if (funcall ,check-func node)
	 (if (string= (treesit-node-type nn) "infix")
	     (treesit-node-text (treesit-node-child nn 1))
	     (haskell-ts-defun-name node))
       nil))))

(defun haskell-ts-indent-defun (pos)
  "Indent the current function."
  (interactive "d")
  (let ((node (treesit-node-at pos)))
    (while (not (string-match
		 "^declarations$\\|haskell"
		 (treesit-node-type (treesit-node-parent node))))
      (setq node (treesit-node-parent node)))
    (indent-region (treesit-node-start node) (treesit-node-end node))))

(defvar haskell-ts-mode-map
  (define-keymap
    "C-c C-c" 'haskell-ts-compile-region-and-go
    "C-c C-r" 'haskell-ts-run-haskell
    "C-M-q" 'haskell-ts-indent-defun)
  "Map for haskell-ts-mode.")

;;;###autoload
(define-derived-mode haskell-ts-mode prog-mode "haskell ts mode"
  "Major mode for Haskell files using tree-sitter."
  (unless (treesit-ready-p 'haskell)
    (error "Tree-sitter for Haskell is not available"))
  (treesit-parser-create 'haskell)
  (setq-local treesit-defun-type-regexp "\\(?:\\(?:function\\|struct\\)_definition\\)")
  ;; Indent
  (when haskell-ts-use-indent
    (setq-local treesit-simple-indent-rules haskell-ts-indent-rules)
    (setq-local indent-tabs-mode nil))
  ;; Comment
  (setq-local comment-start "-- ")
  (setq-local comment-use-syntax t)
  (setq-local comment-start-skip "\\(?: \\|^\\)-+")
  ;; Electric
  (setq-local electric-pair-pairs
	      '((?` . ?`) (?\( . ?\)) (?{ . ?}) (?\" . ?\") (?\[ . ?\])))
  ;; Navigation
  (setq-local treesit-defun-name-function 'haskell-ts-defun-name)
  (setq-local treesit-defun-type-regexp
	      ;; Since haskell is strict functional, any 2nd level
	      ;; entity is defintion
	      (cons ".+"
		    (lambda (node)
		      (and (not (string-match haskell-ts--ignore-types (treesit-node-type node)))
			   (string= "declarations" (treesit-node-type (treesit-node-parent node)))))))
  (setq-local prettify-symbols-alist haskell-ts-prettify-symbols-alist)
  ;; Imenu
  (setq-local treesit-simple-imenu-settings
	      `((nil haskell-ts-imenu-func-node-p nil
		     ,(haskell-ts-imenu-name-function #'haskell-ts-imenu-func-node-p))
		("Signatures.." haskell-ts-imenu-sig-node-p nil
		 ,(haskell-ts-imenu-name-function #'haskell-ts-imenu-sig-node-p))
		("Data..." haskell-ts-imenu-data-type-p nil
		 (lambda (node)
		   (treesit-node-text (treesit-node-child node 1))))))
  ;; font-lock
  (setq-local treesit-font-lock-level haskell-ts-font-lock-level)
  (setq-local treesit-font-lock-settings haskell-ts-font-lock)
  (setq-local treesit-font-lock-feature-list
	      haskell-ts-font-lock-feature-list)
  (treesit-major-mode-setup))

(defun haskell-ts--fontify-arg (node &optional _ _ _)
  (if (string= "variable" (treesit-node-type node))
      (put-text-property
       (treesit-node-start node)
       (treesit-node-end node)
       'face font-lock-variable-name-face)
    (mapc 'haskell-ts--fontify-arg (treesit-node-children node))))

(defun haskell-ts--fontify-type (node &optional _ _ _)
  (let ((last-child (treesit-node-child node -1)))
    (if (string= (treesit-node-type last-child) "function")
	(haskell-ts--fontify-type last-child)
      (put-text-property
       (treesit-node-start last-child)
       (treesit-node-end last-child)
       'face font-lock-variable-name-face))))

(defun haskell-ts-imenu-node-p (regex node)
  (and (string-match-p regex (treesit-node-type node))
       (string= (treesit-node-type (treesit-node-parent node)) "declarations")))

(defun haskell-ts-imenu-func-node-p (node)
  (haskell-ts-imenu-node-p "function\\|bind" node))

(defun haskell-ts-imenu-sig-node-p (node)
  (haskell-ts-imenu-node-p "signature" node))

(defun haskell-ts-imenu-data-type-p (node)
  (haskell-ts-imenu-node-p "data_type" node))

(defun haskell-ts-defun-name (node)
  (treesit-node-text (treesit-node-child node 0)))

(defun haskell-ts-compile-region-and-go (start end)
  "Compile the text from START to END in the haskell proc."
  (interactive "r")
  (let ((hs (haskell-ts-haskell-session)))
    (comint-send-string hs ":{\n")
    (comint-send-region hs start end)
    (comint-send-string hs "\n:}\n")))

(defun haskell-ts-run-haskell()
  (interactive)
  (pop-to-buffer-same-window
   (if (comint-check-proc "*haskell*")
       "*haskell*"
     (make-comint "haskell" "ghci" nil buffer-file-name))))

(defun haskell-ts-haskell-session ()
  (get-buffer-process "*haskell*"))

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

(provide 'haskell-ts-mode)

;;; haskell-ts-mode.el ends here