;;; org-word-count.el --- org-word-count in the modeline -*- lexical-binding: t; -*-

;;; Commentary:

;;; Code:

(require 'org)
(require 'cl-lib)

(defgroup org-word-count nil
  "Extra fast word-counting in `org-mode'."
  :group 'org)

(defvar-local org-word-count-word-count nil
  "Running total of words in this buffer.")

(defvar-local org-word-count-string nil
  "String for the modeline.")

(defcustom org-word-count-format "%sw "
  "Format for org word count in modeline."
  :type 'string)

(defcustom org-word-count-huge-string "huge"
  "String to display with a huge buffer."
  :type 'string)

(defcustom org-word-count-update-after-funcs '(org-narrow-to-subtree
                                               org-narrow-to-block
                                               org-narrow-to-element
                                               org-capture-narrow)
  "Functions after which to update the word count."
  :type '(repeat function))

(defcustom org-word-count-deletion-idle-timer 0.25
  "Length of time, in seconds, to wait before updating word-count."
  :type 'number)

(defcustom org-word-count-huge-change 5000
  "Number of characters that constitute a \"huge\" insertion."
  :type 'number)

(defcustom org-word-count-huge-buffer 10000
  "Number of words past which we're not going to try to count."
  :type 'number)

(defvar org-word-count-correction -5
  "Number to add to `org-word-count-word-count', for some reason?
`org-word-count-word-count' seems to consistently be off by 5.  Thus
this correction.  (At some point I should correct the underlying
code... probably).")

(defvar-local org-word-count-update-timer nil)

;;; Variables from org-wc

(defun org-word-count-list-of-strings-p (arg)
  (cl-every #'stringp arg))

(defun org-word-count--downcase-list-of-strings-set-default (var val)
  (set-default var (mapcar #'downcase val)))

(defcustom org-word-count-ignored-tags '("nowc" "noexport" "ARCHIVE")
  "List of tags for which subtrees will be ignored in word counts"
  :type '(repeat string)
  :safe #'org-word-count-list-of-strings-p)

(defcustom org-word-count-ignore-commented-trees t
  "Ignore trees with COMMENT-prefix if non-nil."
  :type 'boolean
  :safe #'booleanp)

(defcustom org-word-count-default-link-count 'description-or-path
  "Default way of counting words in links.
This is applied to any link type not specified in any of
‘org-word-count-ignored-link-types’,‘org-word-count-one-word-link-types’, or
‘org-word-count-only-description-link-types’ "
  :type '(choice
          (const :tag "Count words in description or else path part of links" description-or-path)
          (const :tag "Count words only in description part of links" description)
          (const :tag "Count links as 0 words" ignore)
          (const :tag "Count links as 1 word" oneword)
          (const :tag "Count words only in path part of links" path))
  :safe 'symbolp)

(defcustom org-word-count-ignored-link-types nil
  "Link types which won't be counted as a word"
  :type '(repeat string)
  :safe #'org-word-count-list-of-strings-p)

(defcustom org-word-count-one-word-link-types '("zotero")
  "Link types which will be counted as one word"
  :type '(repeat string)
  :safe #'org-word-count-list-of-strings-p)

(defcustom org-word-count-description-or-path-link-types '()
  "Link types for which the description or the path should be counted"
  :type '(repeat string)
  :safe #'org-word-count-list-of-strings-p)

(defcustom org-word-count-only-description-link-types '("note")
  "Link types for which only the description should be counted"
  :type '(repeat string)
  :safe #'org-word-count-list-of-strings-p)

(defcustom org-word-count-only-path-link-types '()
  "Link types for which only the path should be counted"
  :type '(repeat string)
  :safe #'org-word-count-list-of-strings-p)

(defcustom org-word-count-blocks-to-count '("quote" "verse")
  "List of blocks which should be included in word count.

Use lower case block names"
  :type '(repeat string)
  :safe #'org-word-count-list-of-strings-p
  :set #'org-word-count--downcase-list-of-strings-set-default)

(defun org-word-count-delayed-update (&rest _)
  (if org-word-count-update-timer
      (setq org-word-count-update-timer nil)
    (setq org-word-count-update-timer
          (run-with-idle-timer org-word-count-deletion-idle-timer nil
                               #'org-word-count-update))))

(defun org-word-count-force-update ()
  (interactive)
  (message "Counting words...")
  (when (timerp org-word-count-update-timer)
    (cancel-timer org-word-count-update-timer))
  (org-word-count-update)
  (message "Counting words...done"))

(defun org-word-count-update (&rest _)         ; Needs variadic parameters, since it's advice
  (dlet ((org-word-count-counting t))
    (org-word-count-buffer)
    (org-word-count-modeline)
    (setq org-word-count-update-timer nil)))

(defun org-word-count-changed (start end length)
  (org-word-count-delayed-update))

(defun org-word-count-buffer ()
  "Count the words in the buffer."
  (when (and (derived-mode-p 'org-mode)
             (not (eq org-word-count-word-count 'huge)))
    (setq org-word-count-word-count
          (cond
           ((> (count-words (point-min) (point-max))
               org-word-count-huge-buffer)
            'huge)
           (t (org-word-count-aux (point-min) (point-max)))))))

;;; From org-wc.el:
;; https://github.com/tesujimath/org-wc/
(defun org-word-count-aux (beg end)
  "Return the number of words between BEG and END."
  (let ((wc 0)
        subtreecount
        (latex-macro-regexp "\\\\[A-Za-z]+\\(\\[[^]]*\\]\\|\\){\\([^}]*\\)}"))
    (save-excursion
      (goto-char beg)
      ;; Handle the case where we start in a drawer
      (when (org-at-drawer-p)
        (org-end-of-meta-data t))
      (while (< (point) end)
        (cond
         ;; Handle headlines and subtrees
         ((org-at-heading-p)
          (cond
           ;; Ignore commented and org-wc-ignored-tags trees
           ((or (and org-word-count-ignore-commented-trees (org-in-commented-heading-p))
                (cl-intersection org-word-count-ignored-tags (org-get-tags) :test #'string=))
            (org-end-of-subtree t t))
           ;; Re-use count for subtrees already counted
           ((setq subtreecount (get-text-property (point) :org-wc))
            (cl-incf wc subtreecount)
            (org-end-of-subtree t t))
           ;; Skip counting words in headline
           (t (org-word-count--goto-char (point-at-eol) end))))
         ;; Ignore most blocks.
         ((when (save-excursion
                  (beginning-of-line 1)
                  (looking-at org-block-regexp))
            (if (member (downcase (match-string 1)) org-word-count-blocks-to-count)
                (progn ;; go inside block and subtract count of end line
                  (org-word-count--goto-char (match-beginning 4) end)
                  (cl-decf wc))
              (org-word-count--goto-char (match-end 0) end))))
         ;; Ignore comments.
         ((org-at-comment-p)
          (org-word-count--goto-char (point-at-eol) end))
         ;; Ignore drawers.
         ((org-at-drawer-p)
          (org-end-of-meta-data t))
         ;; Ignore all other #+ lines
         ((looking-at "#+")
          (org-word-count--goto-char (point-at-eol) end))
         ;; Handle links
         ((save-excursion
            (when (< (1+ (point-min)) (point)) (backward-char 2))
            (looking-at org-link-bracket-re))
          (let* ((type (car (save-match-data (split-string (match-string 1) ":"))))
                 (pathstart (+ 1 (length type) (match-beginning 1))))
            (cl-case (cond ((member type org-word-count-ignored-link-types) 'ignore)
                           ((member type org-word-count-one-word-link-types) 'oneword)
                           ((member type org-word-count-only-description-link-types)
                            'description)
                           ((member type org-word-count-only-path-link-types) 'path)
                           ((member type org-word-count-description-or-path-link-types)
                            'description-or-path)
                           (t org-word-count-default-link-count))
              (ignore (org-word-count--goto-char (match-end 0) end))
              (oneword (org-word-count--goto-char (match-end 0) end)
                       (cl-incf wc))
              (description (if (match-beginning 2)
                               (goto-char (match-beginning 2))
                             (org-word-count--goto-char
                              (match-end 0) end)))
              (path (cl-incf wc (count-words-region pathstart
                                                    (match-end 1)))
                    (org-word-count--goto-char (match-end 0) end))
              (description-or-path
               (if (match-beginning 2)
                   (goto-char (match-beginning 2))
                 (cl-incf wc (count-words-region pathstart
                                                 (match-end 1)))
                 (org-word-count--goto-char (match-end 0) end)))
              (t (user-error "Error in org-word-count link configuration")))))
         ;; Count latex macros as 1 word, ignoring their arguments.
         ((save-excursion
            (when (< (point-min) (point)) (backward-char))
            (looking-at latex-macro-regexp))
          (org-word-count--goto-char (match-end 0) end)
          (cl-incf wc))
         (t
          (and (re-search-forward "\\w+\\W*" end 'skip)
               (cl-incf wc))))))
    wc))

(defun org-word-count--goto-char (char end)
  "Moves point to CHAR and from there passes 0+ non-word characters.
Searchers to end as a maximum.

This ensures that we are in an expected state (at the first word
character after some non-word characters) after moving beyond
headlines, links etc."
  (goto-char char)
  (re-search-forward "\\W*" end 'skip))

(defvar org-word-count-counting nil
  "Are we currently counting?")

(defun org-word-count-recount-widen (&rest _)
  (when (and (not org-word-count-counting))
    (org-word-count-update)))

(defun org-word-count-modeline ()
  (setq org-word-count-string
        (cond
         ((eq org-word-count-word-count 'huge)
          org-word-count-huge-string)
         (org-word-count-word-count
          (format org-word-count-format
                  (max 0 (+ org-word-count-word-count
                            org-word-count-correction))))))
  (force-mode-line-update))

(define-minor-mode org-word-count-mode
  "Count words in `org-mode' buffers in the mode-line."
  :lighter ""
  :keymap (let ((map (make-sparse-keymap)))
            (define-key map (kbd "C-c C-.") #'org-word-count-force-update)
            map)
  (cond (org-word-count-mode
         (org-word-count-buffer)
         (add-hook 'after-change-functions
                   #'org-word-count-delayed-update nil t)
         (unless (member '(org-word-count-mode org-word-count-string)
                         mode-line-misc-info)
           (add-to-list 'mode-line-misc-info
                        '(org-word-count-mode org-word-count-string)
                        nil
                        #'equal))
         (dolist (fn org-word-count-update-after-funcs)
           (advice-add fn :after #'org-word-count-update)))
        (:else
         (remove-hook 'after-change-functions
                      #'org-word-count-delayed-update t)
         (setf mode-line-misc-info
               (delete '(org-word-count-mode org-word-count-string)
                       mode-line-misc-info))
         (dolist (fn org-word-count-update-after-funcs)
           (advice-remove fn #'org-word-count-update)))))

(provide 'org-word-count)
;;; org-word-count.el ends here