;;; eshell-starship.el --- Starship-like (https://starship.rs) prompt for eshell -*- lexical-binding: t; -*-
;;; Commentary:
;;; Code:

(require 'vc)
(require 'vc-git)
(require 'eshell)
(require 'cl-lib)

(defvar eshell-starship-modules nil
  "List of modules used by eshell-starship.
This should be an alist of (name function).  The macro
`eshell-starship-defmodule' can help modify this list.")

(defvar eshell-starship--module-cache (make-hash-table :test 'eq)
  "Hash table to hold module cache for eshell-starship.")

(cl-defmacro eshell-starship-defmodule (name &key pred files dirs exts icon color
                                             allow-remote action reload-on)
  "Define an eshell-starship module called NAME.
The module will be added to `eshell-starship-modules'.
  :PRED - a function that should return t if the module should be run
  :FILES - a list of file names and wildcard expansions that will be used to
           determine if the module should be run
  :EXTS - save as FILES but mach any file with the given extensions
  :DIRS - same as FILES, but for directories
  :ICON - this is the string to print before the modules text
  :COLOR - the color to print the modules text in
  :ALLOW-REMOTE - weather to allow the module to run on remote machines
  :ACTION - a function that will return the module text, or nil
  :RELOAD-ON - when to re-run the module.  Can be one of `'cwd' or `'always'
               (same as nil)"
  (declare (indent defun))
  `(setf (alist-get ',name eshell-starship-modules)
         ,(list 'list pred files dirs exts icon color allow-remote action reload-on)))

(cl-defmacro eshell-starship-find-version-function (command pattern
                                                    &rest format)
  "Return a lambda that calls COMMAND.
COMMAND is in the form of (exec args...).  The temp buffer that was used to run
COMMAND will then have `re-search-forward' run with PATTERN and FORMAT should
arguments to pass to `concat' to format the output."
  (declare (indent defun))
  `(lambda ()
     (with-temp-buffer
       (when (zerop (process-file ,(car command) nil t nil ,@(cdr command)))
         (goto-char (point-min))
         (when (re-search-forward ,pattern nil t)
           (concat ,@format))))))

(eshell-starship-defmodule cc
  :exts '("c" "h")
  :icon "C"
  :color "green yellow"
  :allow-remote nil
  :reload-on 'cwd
  :action (eshell-starship-find-version-function
            ("cc" "-v")
            "^\\([-a-zA-Z]+\\) version \\([0-9]+\\.[0-9]+\\.[0-9]+\\)"
            "v" (match-string 2) "-" (match-string 1)))

(eshell-starship-defmodule cmake
  :files '("CMakeLists.txt" "CMakeCache.txt")
  :icon "󰔶"
  :color "blue"
  :allow-remote nil
  :reload-on 'cwd
  :action (eshell-starship-find-version-function
            ("cmake" "--version")
            "cmake version \\([0-9]+\\.[0-9]+\\.[0-9]+\\)"
            "v" (match-string 1)))

(defun eshell-starship--cwd-clear-caches ()
  "Clear caches that should be cleared on cwd for eshell-starship."
  (dolist (module eshell-starship-modules)
    (if-let ((reload-on (nth 9 module))
             ((eq reload-on 'cwd)))
      (puthash (car module) nil eshell-starship--module-cache))))

(defun eshell-starship--exts-exist-p (&rest exts)
  "Test if any files with EXTS at the end of their name exist in default dir."
  (catch 'found
    (dolist (ext exts)
      (when (seq-filter #'(lambda (name)
                            (not (string-prefix-p "." name)))
                        (file-expand-wildcards (concat "*." ext)))
        (throw 'found t)))))

(defun eshell-starship--files-exist-p (&rest names)
  "Test if any of NAMES exists and are files in default directory."
  (catch 'found
    (dolist (name names)
      (when (file-exists-p name)
        (throw 'found t)))))

(defun eshell-starship--dirs-exist-p (&rest names)
  "Test if any of NAMES exists and are files in default directory."
  (catch 'found
    (dolist (name names)
      (when (file-directory-p name)
        (throw 'found t)))))

(defun eshell-starship--execute-modules ()
  "Execute all the modules in `eshell-starship-modules'."
  (cl-loop
   for (name pred files dirs exts icon color allow-remote action reload-on)
   in eshell-starship-modules
   when (and (or allow-remote (not (file-remote-p default-directory)))
             (or (and files (apply 'eshell-starship--files-exist-p files))
                 (and dirs (apply' eshell-starship--dirs-exist-p dirs))
                 (and exts (apply' eshell-starship--exts-exist-p exts))
                 (and pred (funcall pred))))
   concat (if-let (cache-val (gethash name eshell-starship--module-cache))
              cache-val
            (if-let ((result (funcall action))
                     (mod-string (concat " via " (propertize (concat icon " " result)
                                                             'face `(:foreground ,color)))))
                (unless (or (not reload-on) (eq reload-on 'always))
                  (puthash name mod-string eshell-starship--module-cache))
              result))))

(defun eshell-starship--replace-home-with-tilda (path)
  "If PATH beings with $HOME (the environment variable), replace it with ~."
  (let ((home (getenv "HOME")))
    (if (equal home path)
        "~"
      (setq home (file-name-as-directory home))
      (if (string-prefix-p home path)
          (concat "~/" (seq-subseq path (length home)))
        path))))

(defun eshell-starship--limit-path-parts (num path)
  "Cut PATH down to NUM components.
Example:
/this/is/a/path 3-> is/a/path"
  (let ((parts (string-split path "/" t nil)))
    (concat
     (when (and (file-name-absolute-p path)
                (not (equal "~" (car parts)))
                (<=  (length parts) num))
       "/")
     (string-join (last parts num) "/"))))

(defun eshell-starship--get-current-dir ()
  "Get dir for `eshell-starship--prompt-function'."
  (eshell-starship--limit-path-parts
   3 (if-let ((worktree (vc-root-dir))
              (parent (file-name-parent-directory worktree)))
         (file-relative-name default-directory parent)
       (eshell-starship--replace-home-with-tilda
        default-directory))))

(defun eshell-starship--git-parse-status-headers ()
  "Parse the status headers (read from the current buffer).
The headers are as described in the porcelain v2 section of the git-status(3)
man page.

The return value is a list of the form (oid head upstream ahead behind stash)"
  (let ((oid nil)
        (head nil)
        (upstream nil)
        (ahead nil)
        (behind nil)
        (stash nil))
    (while (and (char-after) (= (char-after) ?#))
      (forward-char 2)
      (cond
       ((looking-at "branch\\.oid ")
        (setq oid (buffer-substring-no-properties
                   (match-end 0)
                   (pos-eol))))
       ((looking-at "branch\\.head ")
        (setq head (buffer-substring-no-properties
                    (match-end 0)
                    (pos-eol))))
       ((looking-at "branch\\.upstream ")
        (setq upstream (buffer-substring-no-properties
                        (match-end 0)
                        (pos-eol))))
       ((looking-at "branch\\.ab ")
        (let ((ab-str (buffer-substring-no-properties
                       (match-end 0)
                       (pos-eol))))
          (when (string-match "\\(+[0-9]+\\) \\(-[0-9]+\\)$"
                              ab-str)
            (setq ahead (string-to-number (match-string 1 ab-str))
                  behind (string-to-number (match-string 2 ab-str))))))
       ((looking-at "stash ")
        (setq stash (string-to-number (buffer-substring-no-properties
                                       (match-end 0)
                                       (pos-eol))))))
      (forward-line))
    (list oid head upstream (or ahead 0) (or behind 0) stash)))

(defun eshell-starship--git-interpret-file-status (x y)
  "Return the prompt character for the status X and Y.
A description of X and Y can be found in the git-status(3) man page."
  (cond
   ((or (= x ?D) (= y ?D))
    ?)
   ((or (= x ?R) (= y ?R))
    ?»)
   ((= y ?M)
    ?!)
   ((or (= x ?A) (= x ?M))
    ?+)))

(defun eshell-starship--git-interpret-branch-status (ahead behind)
  "Get the status char for the current branch and its remote.
AHEAD should evaluate to t if the current branch is ahead of its remote, and
BEHIND should evaluate to t if the current branch is behind its remote."
  (cond
   ((and ahead behind) "󰹺")
   (ahead "󰜷")
   (behind "󰜮")))

(defun eshell-starship--git-file-status (stash ahead behind)
  "Get the file status string for the git prompt module.
STASH should be t if there is current stashed data stash.  AHEAD and BEHIND
should be as for `eshell-starship--git-interpret-branch-status'."
  (let ((merge-conflicts nil)
        (status-chars nil))
    (while (not (eobp))
      (cond
       ((= (char-after) ??)
        (push ?? status-chars))
       ((= (char-after) ?u)
        (setq merge-conflicts t))
       ((or (= (char-after) ?1)
            (= (char-after) ?2))
        (push (eshell-starship--git-interpret-file-status
               (char-after (+ 2 (point)))
               (char-after (+ 3 (point))))
              status-chars)))
      (forward-line))
    (concat (eshell-starship--git-interpret-branch-status (not (zerop ahead))
                                                          (not (zerop behind)))
            (when merge-conflicts "=")
            (when stash "$")
            (apply 'string (sort (seq-uniq status-chars) #'<)))))

(defun eshell-starship--git-current-operation ()
  "Return the current git operation.
For example, a revert.  If there is no current operation, return nil."
  (let ((git-dir (expand-file-name ".git" (vc-git-root default-directory))))
    (cond
     ((file-exists-p (expand-file-name "rebase-apply/applying" git-dir))
      "AM")
     ((file-exists-p (expand-file-name "rebase-apply/rebasing" git-dir))
      "REBASE")
     ((file-exists-p (expand-file-name "rebase-apply" git-dir))
      "AM/REBASE")
     ((file-exists-p (expand-file-name "rebase-merge" git-dir))
      "REBASING")
     ((file-exists-p (expand-file-name "CHERRY_PICK_HEAD" git-dir))
      "CHERRY-PICKING")
     ((file-exists-p (expand-file-name "MERGE_HEAD" git-dir))
      "MERGING")
     ((file-exists-p (expand-file-name "BISECT_LOG" git-dir))
      "BISECTING")
     ((file-exists-p (expand-file-name "REVERT_HEAD" git-dir))
      "REVERTING"))))

(defun eshell-starship--git-status ()
  "Return the text for the git module for `eshell-starship--prompt-function'."
  (with-temp-buffer
    (when (zerop (vc-git-command t nil nil "status" "--porcelain=v2"
                                 "--branch" "--show-stash"))
      (goto-char (point-min))
      (cl-destructuring-bind (oid head _upstream ahead behind stash)
          (eshell-starship--git-parse-status-headers)
        (let ((file-status (eshell-starship--git-file-status stash ahead
                                                             behind))
              (operation (eshell-starship--git-current-operation)))
          (concat
           (if (string= "(detached)" head)
               (propertize (concat " (" (substring oid 0 7) ")")
                           'face '(:foreground "lawn green"))
             (propertize (concat " 󰊢 " head)
                         'face '(:foreground "medium purple")))
           (unless (string-empty-p file-status)
             (propertize (concat " [" file-status "]")
                         'face '(:foreground "red")))
           (when operation
             (concat " (" (propertize
                           operation 'face
                           '(:inherit bold :foreground "yellow")) ")"))))))))

(defun eshell-starship--vc-status ()
  "Get vc status for `eshell-starship--prompt-function'."
  (if-let (backend (vc-responsible-backend default-directory t))
      (if (eq backend 'Git)
          (eshell-starship--git-status)
        (propertize
         (concat "  " (downcase (symbol-name backend)))
         'face '(:foreground "purple")))))

(defvar-local eshell-starship--last-start-time nil
  "Start time of last eshell command.")

(defun eshell-starship--timer-pre-cmd ()
  "Command run before each eshell program to record the time."
  (setq eshell-starship--last-start-time (current-time)))

(add-hook 'eshell-pre-command-hook #'eshell-starship--timer-pre-cmd)

(defun eshell-starship--prompt-format-span (span)
  "Format SPAN as \"XhXms\"."
  (let* ((hours (/ span 3600))
         (mins (% (/ span 60) 60))
         (secs (% span 60)))
    (concat (unless (= hours 0)
              (format "%dh" hours))
            (unless (= mins 0)
              (format "%dm" mins))
            (format "%ds" secs))))

(defun eshell-starship--last-command-time (end-time)
  "Return the prompt component for the time of the last command.
END-TIME is the time when the command finished executing."
  (if-let ((eshell-starship--last-start-time)
           (len (time-subtract end-time
                               eshell-starship--last-start-time))
           (float-len (float-time len))
           ((< 3 float-len))
           (int-len (round float-len)))
      (concat " time "
              (propertize (eshell-starship--prompt-format-span int-len)
                          'face '(:foreground "gold1")))))

(defun eshell-starship--prompt-function ()
  "Function for `eshell-prompt-function'."
  (let* ((end-time (current-time))
         (dir (eshell-starship--get-current-dir))
         (prompt (concat
                  "\n"
                  (if (file-remote-p default-directory)
                      (propertize "🌐 " 'face '(:foreground "light blue")))
                  (propertize dir 'face '(:foreground "dark turquoise"))
                  (unless (file-writable-p default-directory)
                    " ")
                  (eshell-starship--vc-status)
                  (eshell-starship--execute-modules)
                  (eshell-starship--last-command-time end-time)
                  (propertize "\n" 'read-only t 'rear-nonsticky t)
                  (propertize
                   "❯ " 'face `(:foreground
                                ,(if (= eshell-last-command-status 0)
                                     "lime green"
                                   "red"))
                   'rear-nonsticky t))))
    (setq eshell-starship--last-start-time nil)
    prompt))

(defvar-local ehsell-starship--restore-state nil
  "State of various variables set by `eshell-starship-prompt-mode'.")

;;;###autoload
(add-hook 'eshell-directory-change-hook #'eshell-starship--cwd-clear-caches)

;;;###autoload
(define-minor-mode eshell-starship-prompt-mode
  "Minor mode to make eshell prompts look like starship (https://starship.rs)."
  :global nil
  :init-value nil
  :interactive (eshell-mode)
  (if eshell-starship-prompt-mode
      (setq-local eshell-starship--restore-state
                  (buffer-local-set-state eshell-prompt-function
                                          'eshell-starship--prompt-function
                                          eshell-prompt-regexp "^❯ "
                                          eshell-highlight-prompt nil))
    (buffer-local-restore-state eshell-starship--restore-state)))

(provide 'eshell-starship)
;;; eshell-starship.el ends here