;;; my-yatex-review-highlight.el --- Highlight review macros in YaTeX -*- lexical-binding: t; -*-

(require 'cl-lib)

;; ============================================================
;; Faces
;;
;; \RA-\RD:
;;   第1引数: 文字色 + 背景色 + 太字
;;   第2引数: 文字色のみ、背景は変更しない
;;
;; \TO:
;;   第1引数: 赤太字、背景は変更しない
;; \XX:
;;   学生が修正した部分。第1引数を黄色太字で表示する。
;;
;; \CA-\CD:
;;   \CA は \RA と同じ色
;;   \CB は \RB と同じ色
;;   \CC は \RC と同じ色
;;   \CD は \RD と同じ色
;;
;;   {A1} などのラベル部分: label-face
;;   その後の本文: body-face
;; ============================================================

(defface my-latex-RA-label-face '((t ())) "Face for first argument of \\RA.")
(defface my-latex-RA-body-face  '((t ())) "Face for second argument of \\RA.")

(defface my-latex-RB-label-face '((t ())) "Face for first argument of \\RB.")
(defface my-latex-RB-body-face  '((t ())) "Face for second argument of \\RB.")

(defface my-latex-RC-label-face '((t ())) "Face for first argument of \\RC.")
(defface my-latex-RC-body-face  '((t ())) "Face for second argument of \\RC.")

(defface my-latex-RD-label-face '((t ())) "Face for first argument of \\RD.")
(defface my-latex-RD-body-face  '((t ())) "Face for second argument of \\RD.")

(defface my-latex-TO-body-face  '((t ())) "Face for argument of \\TO.")
(defface my-latex-XX-body-face  '((t ())) "Face for argument of \\XX.")


(defvar my-yatex-review-label-background "gray30"
  "Background color for review macro label faces.")


(defun my-yatex-review-set-faces ()
  "Set review macro faces explicitly."

  ;; RA / CA: red
  (set-face-attribute
   'my-latex-RA-label-face nil
   :foreground "tomato"
   :background my-yatex-review-label-background
   :weight 'bold)
  (set-face-attribute
   'my-latex-RA-body-face nil
   :foreground "tomato"
   :background 'unspecified
   :weight 'unspecified)

  ;; RB / CB: light blue
  (set-face-attribute
   'my-latex-RB-label-face nil
   :foreground "DeepSkyBlue"
   :background my-yatex-review-label-background
   :weight 'bold)
  (set-face-attribute
   'my-latex-RB-body-face nil
   :foreground "DeepSkyBlue"
   :background 'unspecified
   :weight 'unspecified)

  ;; RC / CC: green
  (set-face-attribute
   'my-latex-RC-label-face nil
   :foreground "SpringGreen2"
   :background my-yatex-review-label-background
   :weight 'bold)
  (set-face-attribute
   'my-latex-RC-body-face nil
   :foreground "SpringGreen2"
   :background 'unspecified
   :weight 'unspecified)

  ;; RD / CD: purple
  (set-face-attribute
   'my-latex-RD-label-face nil
   :foreground "MediumOrchid1"
   :background my-yatex-review-label-background
   :weight 'bold)
  (set-face-attribute
   'my-latex-RD-body-face nil
   :foreground "MediumOrchid1"
   :background 'unspecified
   :weight 'unspecified)

  ;; TO: red bold
  (set-face-attribute
   'my-latex-TO-body-face nil
   :foreground "tomato"
   :background 'unspecified
   :weight 'bold)

  ;; XX: student revision
  (set-face-attribute
   'my-latex-XX-body-face nil
   :foreground "DeepPink"
   :background 'unspecified
   :weight 'bold))


(my-yatex-review-set-faces)


;; ============================================================
;; Variables
;; ============================================================

(defvar-local my-yatex-review-overlays nil
  "Overlays used for review macro highlighting.")

(defvar-local my-yatex-review-highlight-enabled nil
  "Non-nil means review macro highlighting is enabled.")

(defvar-local my-yatex-review-refresh-timer nil
  "Idle timer for automatic review macro highlighting refresh.")

(defvar-local my-yatex-review-font-lock-removed-p nil
  "Non-nil means old font-lock review highlighting has already been removed.")

(defvar my-yatex-review-idle-delay 0.8
  "Idle delay seconds before automatic review macro highlighting refresh.")


;; ============================================================
;; Remove old font-lock highlighting
;; ============================================================

(defun my-yatex-review-remove-old-font-lock ()
  "Remove old font-lock based review macro highlighting."
  (interactive)

  (when (boundp 'my-yatex-review-macro-keywords)
    (font-lock-remove-keywords
     nil
     (symbol-value 'my-yatex-review-macro-keywords)))

  (font-lock-remove-keywords
   nil
   '(("\\\\RA{\\([^{}\n]*\\)}{\\([^{}\n]*\\)}"
      (1 'my-latex-RA-label-face t)
      (2 'my-latex-RA-body-face t))
     ("\\\\RB{\\([^{}\n]*\\)}{\\([^{}\n]*\\)}"
      (1 'my-latex-RB-label-face t)
      (2 'my-latex-RB-body-face t))
     ("\\\\RC{\\([^{}\n]*\\)}{\\([^{}\n]*\\)}"
      (1 'my-latex-RC-label-face t)
      (2 'my-latex-RC-body-face t))
     ("\\\\RD{\\([^{}\n]*\\)}{\\([^{}\n]*\\)}"
      (1 'my-latex-RD-label-face t)
      (2 'my-latex-RD-body-face t))
     ("\\\\TO{\\([^{}\n]*\\)}"
      (1 'my-latex-TO-body-face t))
     ("\\\\XX{\\([^{}\n]*\\)}"
      (1 'my-latex-XX-body-face t))))

  ;; 重いので refresh のたびには呼ばない
  (when font-lock-mode
    (font-lock-flush (point-min) (point-max))
    (font-lock-ensure (point-min) (point-max)))

  (when (called-interactively-p 'interactive)
    (message "Old font-lock review highlighting removed")))


;; ============================================================
;; Overlay utilities
;; ============================================================

(defun my-yatex-review--clear-overlays (&optional thorough)
  "Clear review macro overlays.

If THOROUGH is non-nil, also remove untracked overlays with
property `my-yatex-review'."
  (mapc #'delete-overlay my-yatex-review-overlays)
  (setq my-yatex-review-overlays nil)

  (when thorough
    (remove-overlays (point-min) (point-max) 'my-yatex-review t)))


(defun my-yatex-review--make-overlay (beg end face &optional priority)
  "Create overlay from BEG to END with FACE.
Optional PRIORITY specifies overlay priority."
  (when (< beg end)
    (let ((ov (make-overlay beg end)))
      (overlay-put ov 'face face)
      (overlay-put ov 'priority (or priority 1000))
      (overlay-put ov 'my-yatex-review t)
      (push ov my-yatex-review-overlays))))


;; ============================================================
;; Brace parser
;; ============================================================

(defun my-yatex-review--escaped-p (pos)
  "Return non-nil if character at POS is escaped by backslashes."
  (let ((n 0)
        (p (1- pos)))
    (while (and (>= p (point-min))
                (eq (char-after p) ?\\))
      (setq n (1+ n))
      (setq p (1- p)))
    (= (mod n 2) 1)))


(defun my-yatex-review--read-brace-arg ()
  "Read one LaTeX brace argument at point.

Point may be before spaces/newlines followed by `{`.
Return (INNER-BEG INNER-END OUTER-END), or nil."
  (skip-chars-forward " \t\r\n")
  (if (not (eq (char-after) ?{))
      nil
    (let ((open (point))
          (depth 0)
          close)
      (while (and (not close) (not (eobp)))
        (let ((ch (char-after)))
          (cond
           ((and (eq ch ?{)
                 (not (my-yatex-review--escaped-p (point))))
            (setq depth (1+ depth)))
           ((and (eq ch ?})
                 (not (my-yatex-review--escaped-p (point))))
            (setq depth (1- depth))
            (when (= depth 0)
              (setq close (point)))))
          (forward-char 1)))
      (when close
        ;; point is now just after closing brace
        (list (1+ open) close (point))))))


;; ============================================================
;; Macro faces
;; ============================================================

(defun my-yatex-review--faces-for-macro (macro)
  "Return label/body faces for \\RA, \\RB, \\RC, or \\RD."
  (cond
   ((string= macro "RA")
    '(my-latex-RA-label-face my-latex-RA-body-face))
   ((string= macro "RB")
    '(my-latex-RB-label-face my-latex-RB-body-face))
   ((string= macro "RC")
    '(my-latex-RC-label-face my-latex-RC-body-face))
   ((string= macro "RD")
    '(my-latex-RD-label-face my-latex-RD-body-face))
   (t nil)))


(defun my-yatex-review--faces-for-citem-macro (macro)
  "Return label/body faces for \\CA, \\CB, \\CC, or \\CD."
  (cond
   ((string= macro "CA")
    '(my-latex-RA-label-face my-latex-RA-body-face))
   ((string= macro "CB")
    '(my-latex-RB-label-face my-latex-RB-body-face))
   ((string= macro "CC")
    '(my-latex-RC-label-face my-latex-RC-body-face))
   ((string= macro "CD")
    '(my-latex-RD-label-face my-latex-RD-body-face))
   (t nil)))


;; ============================================================
;; \CA-\CD item-body highlighting
;; ============================================================

(defvar my-yatex-review-citem-stop-regexp
  "^[ \t]*\\(?:\\\\C[OABCD]\\_>\\|\\\\end[ \t\r\n]*{\\(?:enumerate\\|document\\)}\\)"
  "Regexp that stops highlighting body text after \\CA, \\CB, \\CC, or \\CD.")


(defun my-yatex-review--highlight-citem-macros ()
  "Highlight \\CA, \\CB, \\CC, and \\CD item bodies.

The body is highlighted until the next \\CO, \\CA, \\CB, \\CC,
\\CD, \\end{enumerate}, or \\end{document}.

Return number of highlighted answer macros."
  (let ((count 0))
    (goto-char (point-min))

    ;; 行頭またはインデント後の \CA, \CB, \CC, \CD を対象にする
    (while (re-search-forward "\\\\\\(R[ABCD]\\|TO\\|XX\\)\\_>" nil t)      
      (let* ((macro (match-string-no-properties 1))
             (faces (my-yatex-review--faces-for-citem-macro macro))
             (label-face (nth 0 faces))
             (body-face  (nth 1 faces)))
        (let ((arg1 (my-yatex-review--read-brace-arg)))
          (when arg1
            (let* ((body-beg (nth 2 arg1))
                   (body-end
                    (save-excursion
                      (goto-char body-beg)
                      (if (re-search-forward
                           my-yatex-review-citem-stop-regexp nil t)
                          (match-beginning 0)
                        (point-max)))))
              (setq count (1+ count))

              ;; \CA{A1} の A1 部分
              (my-yatex-review--make-overlay
               (nth 0 arg1) (nth 1 arg1) label-face 1000)

              ;; A1 に続く本文部分
              ;; \RA などの局所 overlay より少し低い priority にする
              (my-yatex-review--make-overlay
               body-beg body-end body-face 900)

              ;; 次の item まで飛ぶ
              (goto-char body-end))))))

    count))


;; ============================================================
;; Main refresh
;; ============================================================

(defun my-yatex-review-highlight-refresh ()
  "Refresh highlighting for \\RA-\\RD, \\TO, \\XX, and \\CA-\\CD."
  (interactive)

  ;; refresh 時は font-lock 削除や face 再設定をしない。
  ;; それらは enable 時に一度だけ行う。

  (my-yatex-review--clear-overlays)

  (let ((count 0))
    (save-excursion
      (save-restriction
        (widen)

        ;; ------------------------------------------------------------
        ;; \RA, \RB, \RC, \RD, \TO
        ;; ------------------------------------------------------------
        (goto-char (point-min))

        (while (re-search-forward "\\\\\\(R[ABCD]\\|TO\\|XX\\)\\_>" nil t)
          (let ((macro (match-string-no-properties 1)))
            (cond

             ;; \RA{label}{body}, \RB{label}{body}, ...
             ((member macro '("RA" "RB" "RC" "RD"))
              (let* ((faces (my-yatex-review--faces-for-macro macro))
                     (label-face (nth 0 faces))
                     (body-face  (nth 1 faces)))
                (let ((arg1 (my-yatex-review--read-brace-arg)))
                  (when arg1
                    (let ((arg2 (my-yatex-review--read-brace-arg)))
                      (when arg2
                        (setq count (1+ count))

                        ;; first argument
                        (my-yatex-review--make-overlay
                         (nth 0 arg1) (nth 1 arg1) label-face 1000)

                        ;; second argument
                        (my-yatex-review--make-overlay
                         (nth 0 arg2) (nth 1 arg2) body-face 1000)))))))
             ;; \TO{body}
             ((string= macro "TO")
              (let ((arg1 (my-yatex-review--read-brace-arg)))
                (when arg1
                  (setq count (1+ count))
                  (my-yatex-review--make-overlay
                   (nth 0 arg1) (nth 1 arg1)
                   'my-latex-TO-body-face 1000))))

             ;; \XX{body}: student revision
             ((string= macro "XX")
              (let ((arg1 (my-yatex-review--read-brace-arg)))
                (when arg1
                  (setq count (1+ count))
                  (my-yatex-review--make-overlay
                   (nth 0 arg1) (nth 1 arg1)
                   'my-latex-XX-body-face 1000)))))))
	
        ;; ------------------------------------------------------------
        ;; \CA, \CB, \CC, \CD
        ;; ------------------------------------------------------------
        (setq count
              (+ count
                 (my-yatex-review--highlight-citem-macros)))))

    (when (called-interactively-p 'interactive)
      (message "Review macro highlight refreshed: %d macros/items found" count))))


;; ============================================================
;; Automatic refresh after editing
;; ============================================================

(defun my-yatex-review-highlight-refresh-silent ()
  "Refresh highlighting silently."
  (cl-letf (((symbol-function 'message)
             (lambda (&rest _args) nil)))
    (my-yatex-review-highlight-refresh)))


(defun my-yatex-review-schedule-refresh (&rest _args)
  "Schedule automatic refresh after editing."
  (when my-yatex-review-highlight-enabled
    (when my-yatex-review-refresh-timer
      (cancel-timer my-yatex-review-refresh-timer))
    (setq my-yatex-review-refresh-timer
          (run-with-idle-timer
           my-yatex-review-idle-delay nil
           (lambda (buf)
             (when (buffer-live-p buf)
               (with-current-buffer buf
                 (when my-yatex-review-highlight-enabled
                   (my-yatex-review-highlight-refresh-silent)))))
           (current-buffer)))))


;; ============================================================
;; Enable / disable / toggle
;; ============================================================

(defun my-yatex-review-highlight-enable ()
  "Enable review macro highlighting."
  (interactive)

  (my-yatex-review-install-keybindings)
  (setq my-yatex-review-highlight-enabled t)

  ;; 古い font-lock 版の削除は一度だけ行う
  (unless my-yatex-review-font-lock-removed-p
    (my-yatex-review-remove-old-font-lock)
    (setq my-yatex-review-font-lock-removed-p t))

  ;; face 設定も enable 時だけで十分
  (my-yatex-review-set-faces)

  (add-hook 'after-change-functions
            #'my-yatex-review-schedule-refresh
            nil t)

  (my-yatex-review-highlight-refresh)
  (message "Review macro highlight: ON"))


(defun my-yatex-review-highlight-disable ()
  "Disable review macro highlighting."
  (interactive)

  (setq my-yatex-review-highlight-enabled nil)

  (remove-hook 'after-change-functions
               #'my-yatex-review-schedule-refresh
               t)

  (when my-yatex-review-refresh-timer
    (cancel-timer my-yatex-review-refresh-timer)
    (setq my-yatex-review-refresh-timer nil))

  ;; disable 時は念のため thorough に消す
  (my-yatex-review--clear-overlays t)

  ;; 古い font-lock 版が残っていれば削除
  (unless my-yatex-review-font-lock-removed-p
    (my-yatex-review-remove-old-font-lock)
    (setq my-yatex-review-font-lock-removed-p t))

  (message "Review macro highlight: OFF"))


(defun my-yatex-review-highlight-toggle ()
  "Toggle review macro highlighting."
  (interactive)
  (if (or my-yatex-review-highlight-enabled
          my-yatex-review-overlays)
      (my-yatex-review-highlight-disable)
    (my-yatex-review-highlight-enable)))


;; ============================================================
;; Key bindings
;; ============================================================

(defun my-yatex-review-install-keybindings ()
  "Install keybindings for review macro highlighting in current buffer."
  (interactive)

  ;; C-c h : toggle ON/OFF
  ;; C-c r : refresh
  (local-set-key (kbd "C-c h") #'my-yatex-review-highlight-toggle)
  (local-set-key (kbd "C-c r") #'my-yatex-review-highlight-refresh)

  (when (called-interactively-p 'interactive)
    (message "Review macro keybindings installed: C-c h toggle, C-c r refresh")))


(defun my-yatex-review-setup ()
  "Setup review macro highlighting for current YaTeX buffer."
  (interactive)
  (my-yatex-review-install-keybindings)
  (my-yatex-review-highlight-enable))


;; ============================================================
;; Hooks
;; ============================================================

(add-hook 'yatex-mode-hook #'my-yatex-review-setup)
(add-hook 'YaTeX-mode-hook #'my-yatex-review-setup)


(provide 'my-yatex-review-highlight)
;;; my-yatex-review-highlight.el ends here