;;; inferior-jshell.el --- Run JShell in a comint buffer  -*- lexical-binding: t; -*-
;;; Commentary:
;;; Code:
(require 'comint)
(require 'cl-lib)
(require 'cc-mode)
(require 'treesit)
(eval-when-compile (require 'rx))

(defgroup inferior-jshell ()
  "Run JShell in a comint buffer."
  :group 'comint)

(defvar jshell-buffer-name "*jshell*"
  "Name to use for inferior jshell process buffer.")

(defvar jshell-program "jshell"
  "The program to default to for `run-jshell'.")

(defvar jshell-switches ()
  "List of arguments to pass to jshell in `run-jshell'.")

(defvar-local jshell--fontification-buffer nil
  "The fontification buffer for the current jshell buffer.")

(defvar-local jshell--skip-next-lines 0
  "Number of lines of output to skip.")

(defun jshell--preoutput-filter-function (output)
  "Preoutput filter function for jshell.
OUTPUT is the new text to be inserted."
  (if (<= jshell--skip-next-lines 0)
      output
    (let* ((lines (string-lines output))
           (cnt (length lines)))
      (if (> cnt jshell--skip-next-lines)
          (prog1
              (string-join (nthcdr jshell--skip-next-lines lines) "\n")
            (setq jshell--skip-next-lines 0))
        (cl-decf jshell--skip-next-lines cnt)
        (when (and (not (string-empty-p output))
                   (/= ?\n (elt output (1- (length output)))))
          (cl-incf jshell--skip-next-lines))
        ""))))

(defun jshell--get-fontification-buffer ()
  "Return the jshell fontification buffer."
  (if (buffer-live-p jshell--fontification-buffer)
      jshell--fontification-buffer
    (let ((buffer (generate-new-buffer
                   (format " %s-fontification-buffer" (buffer-name)))))
      (with-current-buffer buffer
        (unless (derived-mode-p 'c++-mode)
          (let ((delayed-mode-hooks nil))
            (delay-mode-hooks
              (java-mode))))
        (when (eq c-basic-offset 'set-from-style)
          (setq-local c-basic-offset standard-indent))
        (let ((inhibit-message t))
          (indent-tabs-mode -1))
        (unless font-lock-mode
          (font-lock-mode 1)))
      (setq-local jshell--fontification-buffer buffer))))

(defmacro jshell--with-font-lock-buffer (&rest body)
  "Execute BODY in the jshell indirect buffer.
Note that this erases the buffer before doing anything."
  `(with-current-buffer (jshell--get-fontification-buffer)
     (erase-buffer)
     ,@body))

(defun jshell--fontify-current-input ()
  "Function called from `post-command-hook' to fontify the current input."
  (when-let ((proc (get-buffer-process (current-buffer)))
             (start (process-mark proc))
             (end (point-max))
             (input (buffer-substring-no-properties start end))
             (fontified (jshell--with-font-lock-buffer
                         (insert input)
                         (font-lock-ensure)
                         (buffer-string)))
             (len (length fontified))
             (i 0))
    ;; mostly from:
    ;; `python-shell-font-lock-post-command-hook'
    (while (not (= i len))
      (let* ((props (text-properties-at i fontified))
             (change-i (or (next-property-change i fontified)
                           len)))
        (when-let ((face (plist-get props 'face)))
          (setf (plist-get props 'face) nil
                (plist-get props 'font-lock-face) face))
        (set-text-properties (+ start i) (+ start change-i) props)
        (setq i change-i)))))

(defun jshell--bounds-of-last-prompt ()
  "Return the bounds of the last jshell prompt.
This returns a cons."
  (save-excursion
    (let ((end (process-mark (get-buffer-process (current-buffer)))))
      (goto-char end)
      (cons (pos-bol) end))))

(defun jshell--remove-extra-indentation (count)
  "Remove COUNT spaces from the start of each line."
  (save-excursion
    (goto-char (point-min))
    (while (not (eobp))
      (back-to-indentation)
      (let ((indent (- (point) (pos-bol))))
        (when (> indent count)
          (delete-char (- count))))
      (forward-line))))

(defun jshell--indent-line-function ()
  "`indent-line-function' for jshell comint buffers."
  (let* ((start (process-mark (get-buffer-process (current-buffer)))))
    ;; don't indent the first line
    (unless (= (pos-bol) (save-excursion (goto-char start) (pos-bol)))
      (let* ((input (buffer-substring-no-properties start (pos-eol)))
             (prompt-size (let ((bound (jshell--bounds-of-last-prompt)))
                            (- (cdr bound) (car bound))))
             (col (jshell--with-font-lock-buffer
                   (insert input)
                   (jshell--remove-extra-indentation prompt-size)
                   (c-indent-line nil t)
                   (back-to-indentation)
                   (- (point) (pos-bol)))))
        (save-excursion
          (indent-line-to (+ prompt-size col)))
        (skip-syntax-forward "-")))))

(defun jshell-send-input ()
  "Like `comint-send-input', but with some extra stuff for jshell."
  (interactive)
  (let ((pmark (process-mark (get-buffer-process (current-buffer))))
        (end (if comint-eol-on-send (pos-eol) (point))))
    (with-restriction pmark end
      (let ((res (syntax-ppss (point-max))))
        (without-restriction
          (cond
           ;; open string
           ((cl-fourth res)
            (message "Unterminated string"))
           ;; unmatched blocks or comment
           ((or (numberp (cl-fifth res))
                (not (zerop (cl-first res)))
                ;; trailing . character
                (save-excursion
                  (end-of-line)
                  (skip-syntax-backward "-")
                  (eql (char-before) ?.)))
            (newline-and-indent))
           (t
            ;; jshell will echo out each line we send it, ignore all of them
            (setq-local jshell--skip-next-lines (count-lines pmark end))
            (when (= pmark end)
              (cl-incf jshell--skip-next-lines))
            ;; also, methods add a bunch of extra newlines
            (when (>= jshell--skip-next-lines 2)
              (cl-incf jshell--skip-next-lines (- jshell--skip-next-lines 2)))
            (comint-send-input))))))))

(defvar-keymap inferior-jshell-mode-map
  :doc "Keymap for `inferior-jshell-mode'."
  :parent comint-mode-map
  "RET" #'jshell-send-input)

(defvar inferior-jshell-mode-syntax-table (copy-syntax-table java-mode-syntax-table)
  "Syntax table for `inferior-jshell-mode'.")

(defun jshell--kill-fontification-buffer ()
  "Kill the current `jshell--fontification-buffer'."
  (kill-buffer jshell--fontification-buffer))

(define-derived-mode inferior-jshell-mode comint-mode "Inferior Jshell"
  "Major mode for buffers running inferior CERN jshell processes."
  :group 'inferior-jshell
  :syntax-table inferior-jshell-mode-syntax-table
  (setq-local comint-highlight-input nil
              indent-line-function #'jshell--indent-line-function
              electric-indent-chars '(?\n ?}))
  (add-hook 'comint-preoutput-filter-functions
            #'jshell--preoutput-filter-function
            nil t)
  (add-hook 'post-command-hook
            #'jshell--fontify-current-input
            nil t)
  (add-hook 'kill-buffer-hook
            #'jshell--kill-fontification-buffer
            nil t))

(cl-defun run-jshell (&optional (cmd jshell-program) (switches jshell-switches))
  "Run CERN jshell in a comint buffer.
CMD is the shell command to run.  It defaults to `jshell-program'.  SWITCHES are
the arguments to pass.  They default to `jshell-switches'."
  (interactive)
  (pop-to-buffer
   (with-current-buffer (get-buffer-create jshell-buffer-name)
     (prog1 (current-buffer)
       (unless (process-live-p (get-buffer-process (current-buffer)))
         (inferior-jshell-mode)
         (comint-exec (current-buffer) "Inferior Jshell" cmd nil switches))))))

(cl-defun jshell--find-buffer ()
  "Find and return a live jshell buffer."
  (dolist (buffer (buffer-list))
    (with-current-buffer buffer
      (when (and (derived-mode-p 'inferior-jshell-mode)
                 (process-live-p (get-buffer-process buffer)))
        (cl-return-from jshell--find-buffer buffer)))))

(defun jshell-eval (code)
  "Evaluate CODE in a live JShell buffer."
  (interactive "sEval: ")
  (let ((buffer (jshell--find-buffer)))
    (unless buffer
      (user-error "No live JShell buffer found"))
    (with-current-buffer buffer
      (let* ((start (process-mark (get-buffer-process buffer)))
             (end (point-max))
             (old (buffer-substring-no-properties start end)))
        (delete-region start end)
        (goto-char (point-max))
        (insert code)
        (goto-char (point-max))
        (jshell-send-input)
        (goto-char (point-max))
        (insert old)
        (goto-char (point-max))))))

(defun jshell-eval-region (start end)
  "Evaluate the current buffer from START to END in a live JShell buffer.
START and END default to the current region."
  (interactive "r")
  (jshell-eval (buffer-substring-no-properties start end)))

(defun jshell-eval-buffer ()
  "Send the current buffer to a live JShell buffer."
  (interactive)
  (jshell-eval-region (point-min) (point-max)))

(defun jshell-eval-defun ()
  "Send the defun under point to a live JShell buffer."
  (interactive)
  (let ((bounds (bounds-of-thing-at-point 'defun)))
    (unless bounds
      (user-error "No defun under point"))
    (jshell-eval-region (car bounds) (cdr bounds))))

(defun jshell-eval-expression ()
  "Send the Java expression under point to a live JShell buffer.
This only works in `java-ts-mode'."
  (interactive)
  (let ((root (treesit-buffer-root-node)))
    (let ((node (car (or (treesit-query-range
                          root '([(expression_statement)
                                  (field_declaration)
                                  (local_variable_declaration)
                                  (import_declaration)]
                                 @exp)
                          (point) (1+ (point)))
                         (treesit-query-range
                          root '([(parenthesized_expression)
                                  (binary_expression)
                                  (update_expression)
                                  (unary_expression)]
                                 @exp)
                          (point) (1+ (point)))))))
      (unless node
        (user-error "No expression found under point"))
      (let ((text (buffer-substring-no-properties (car node) (cdr node))))
        (when (string-match (rx (* (syntax whitespace))
                                ";"
                                (* (syntax whitespace)) eos)
                            text)
          (setq text (substring text 0 (match-beginning 0))))
        (when (string-match (rx bos (* (syntax whitespace)) "("
                                (group (* any))
                                ")" (* (syntax whitespace)) eos)
                            text)
          (setq text (match-string 1 text)))
        (jshell-eval text)))))


(provide 'inferior-jshell)
;;; inferior-jshell.el ends here