;;; citre-backend.el --- Live and persistent tags through Citre -*- lexical-binding: t; -*-

(require 'citre)
(require 'citre-tags)
(require 'citre-index)
(require 'subr-x)

(defcustom my/citre-completion-deduplication 'name
  "How Citre completion candidates are deduplicated.
The value `name' keeps one ranked candidate for each insertion text.  Nil keeps every tag returned by Citre."
  :type '(choice (const :tag "One candidate per name" name)
                 (const :tag "Keep every tag" nil))
  :group 'my/citre)

(defcustom my/citre-completion-context-filter 'safe
  "How receiver context affects Citre completions.
The value `safe' filters by an explicitly declared receiver type only when at least one returned tag matches it.  Nil only uses context for ranking."
  :type '(choice (const :tag "Filter only verified receiver types" safe)
                 (const :tag "Only rank by context" nil))
  :group 'my/citre)

(defcustom my/citre-completion-source-order
  '(current-buffer live-buffer project dependency tags-file)
  "Preferred source order for Citre completion candidates."
  :type '(repeat (choice (const current-buffer) (const live-buffer)
                         (const project) (const dependency) (const tags-file)))
  :group 'my/citre)

(defcustom my/citre-completion-annotation '(scope source count)
  "Components shown after a Citre completion candidate."
  :type '(set (const scope) (const source) (const count))
  :group 'my/citre)

(defcustom my/citre-reference-program "rg"
  "Program used to search references.
It must accept ripgrep's options and honour the project's ignore files."
  :type 'string
  :group 'my/citre)

(defcustom my/citre-completion-type-resolvers
  '(my/citre-backend--resolve-go-type
    my/citre-backend--resolve-rust-or-python-type
    my/citre-backend--resolve-c-type)
  "Functions used in order to resolve an explicitly declared receiver type.
Each function receives a receiver name and the beginning of its member expression and returns a type name or nil."
  :type 'hook
  :group 'my/citre)

(defvar company--capf-cache)
(defvar-local my/citre-backend--completion-cache nil)

;; Retire the popup restart timer when reloading an earlier version of this module.
(when (boundp 'my/citre-backend--company-timer)
  (my/citre-index--cancel-timer (symbol-value 'my/citre-backend--company-timer))
  (set 'my/citre-backend--company-timer nil))

(defun my/citre-backend--sources ()
  "Return tag source plists in priority order."
  (if-let* ((state my/citre-index--state))
      (let* ((entries (my/citre-index-live-entries state))
             (current (my/citre-index--path))
             (live (sort entries (lambda (a b) (and (equal (car a) current) (not (equal (car b) current)))))))
        (append
         (mapcar (lambda (entry)
                   (list :file (plist-get (cdr entry) :file)
                         :live-source (car entry)
                         :source (if (equal (car entry) current) 'current-buffer 'live-buffer)))
                 live)
         (let ((project (expand-file-name "project.tags" (my/citre-index-state-cache state))))
           (if (file-readable-p project)
               (list (list :file project :source 'project))
             (when-let* ((legacy (citre-tags-file-path)))
               (list (list :file legacy :source 'tags-file)))))
         (let ((dependencies (expand-file-name "dependencies.tags" (my/citre-index-state-cache state))))
           (when (file-readable-p dependencies)
             (list (list :file dependencies :source 'dependency))))))
    (when-let* ((file (citre-tags-file-path)))
      (list (list :file file :source 'tags-file)))))

(defun my/citre-backend--source-rank (source)
  "Return the configured rank of SOURCE."
  (or (seq-position my/citre-completion-source-order source) most-positive-fixnum))

(defun my/citre-backend--query (operation &optional identifier)
  "Merge tags for OPERATION, optionally using IDENTIFIER, while masking superseded files."
  (let* ((sources (my/citre-backend--sources))
         (masked (delq nil (mapcar (lambda (source) (plist-get source :live-source)) sources)))
         (seen (make-hash-table :test #'equal))
         result)
    (dolist (source sources)
      (let* ((file (plist-get source :file))
             (symbol (or (and identifier
                              (citre-put-property (copy-sequence identifier)
                                                  'file-path buffer-file-name 'tags-file file))
                         (and (not (memq operation '(identifiers imenu)))
                              (citre-tags-get-symbol file))))
             (filter
              (pcase operation
                ('completion (when symbol (or (citre-tags--get-value-in-language-alist :completion-filter symbol)
                                              (citre-tags-completion-default-filter symbol))))
                ('definition (when symbol (or (citre-tags--get-value-in-language-alist :definition-filter symbol)
                                              (citre-tags-definition-default-filter symbol))))
                ('imenu (list 'and (citre-readtags-filter-input (my/citre-index--path) file)
                              (list 'not citre-tags-filter-file-tags)))
                (_ (list 'not citre-tags-filter-file-tags))))
             (sorter
              (pcase operation
                ('completion (when symbol (or (citre-tags--get-value-in-language-alist :completion-sorter symbol)
                                              citre-tags-completion-default-sorter)))
                ('imenu (citre-readtags-sorter 'line))
                (_ citre-tags-definition-default-sorter)))
             (tags (when (or symbol (memq operation '(identifiers imenu)))
                     (citre-tags-get-tags
                      file symbol (if (eq operation 'completion) (if citre-tags-substr-completion 'substr 'prefix) 'exact)
                      :filter filter :sorter sorter
                      :require '(name ext-abspath)
                      :optional '(pattern line end ext-kind-full signature scope typeref extras language)))))
        (dolist (tag tags)
          (let* ((path (citre-get-tag-field 'ext-abspath tag))
                 (key (mapcar (lambda (field) (citre-get-tag-field field tag))
                              '(name ext-abspath line scope ext-kind-full))))
            (unless (or (and (not (plist-get source :live-source)) (member path masked))
                        (gethash key seen))
              (puthash key t seen)
              (when (eq operation 'completion)
                (setq tag (copy-hash-table tag))
                (citre-set-tag-field 'my-source (plist-get source :source) tag)
                (citre-set-tag-field 'my-order (length result) tag))
              (push tag result))))))
    (nreverse result)))

(defun my/citre-backend--member-context (beginning)
  "Return the member-access context immediately before BEGINNING."
  (save-excursion
    (goto-char beginning)
    (skip-chars-backward " \t")
    (let ((end (point)))
      (when (and (re-search-backward
                  "\\_<\\([[:alpha:]_][[:alnum:]_]*\\)\\_>[ \t]*\\(\\.\\|->\\|::\\)[ \t]*"
                  (line-beginning-position) t)
                 (= (match-end 0) end))
        (let* ((receiver (match-string-no-properties 1))
               (operator (match-string-no-properties 2))
               (type (run-hook-with-args-until-success
                      'my/citre-completion-type-resolvers receiver beginning)))
          (list :receiver receiver :operator operator :type type))))))

(defun my/citre-backend--defun-start ()
  "Return the beginning of the enclosing defun without moving point."
  (save-excursion
    (condition-case nil
        (progn (beginning-of-defun) (point))
      (error (point-min)))))

(defun my/citre-backend--unique-match (regexp text)
  "Return the sole normalized type captured by REGEXP in TEXT."
  (let (matches)
    (with-temp-buffer
      (insert text)
      (goto-char (point-min))
      (while (re-search-forward regexp nil t)
        (push (my/citre-backend--normalize-type (match-string-no-properties 1)) matches)))
    (car (when (= (length (delete-dups (delq nil matches))) 1) matches))))

(defun my/citre-backend--normalize-type (type)
  "Return a comparable nominal name for TYPE."
  (when type
    (let ((value (string-trim type)))
      (setq value (replace-regexp-in-string "\\_<\\(?:const\\|volatile\\|mut\\|struct\\|class\\|enum\\|dyn\\)\\_>" "" value)
            value (replace-regexp-in-string "[&*[:space:]]" "" value)
            value (replace-regexp-in-string "<.*>\\'" "" value)
            value (replace-regexp-in-string "::" "." value))
      (unless (string-empty-p value) value))))

(defun my/citre-backend--context-text (beginning)
  "Return text from the enclosing defun through BEGINNING."
  (buffer-substring-no-properties (my/citre-backend--defun-start) beginning))

(defun my/citre-backend--resolve-go-type (receiver beginning)
  "Resolve RECEIVER from an explicit Go declaration before BEGINNING."
  (when (derived-mode-p 'go-mode 'go-ts-mode)
    (let* ((text (my/citre-backend--context-text beginning))
           (name (regexp-quote receiver))
           (regexp (format "\\(?:\\_<var[ \t]+%s\\_>\\|[,( \t\n]%s\\_>\\)[ \t]+\\(\\*?[[:alpha:]_][[:alnum:]_./]*\\)" name name)))
      (my/citre-backend--unique-match regexp text))))

(defun my/citre-backend--resolve-rust-or-python-type (receiver beginning)
  "Resolve an explicitly annotated Rust or Python RECEIVER before BEGINNING."
  (when (or (derived-mode-p 'rust-mode 'rust-ts-mode)
            (derived-mode-p 'python-mode 'python-ts-mode))
    (let* ((text (my/citre-backend--context-text beginning))
           (regexp (format "\\_<%s\\_>[ \t]*:[ \t]*\\(?:&[ \t]*\\)?\\(?:mut[ \t]+\\)?\\([[:alpha:]_][[:alnum:]_.:]*\\(?:<[^>\n]+>\\)?\\)"
                           (regexp-quote receiver))))
      (my/citre-backend--unique-match regexp text))))

(defun my/citre-backend--resolve-c-type (receiver beginning)
  "Resolve RECEIVER from a simple explicit C or C++ declaration."
  (when (derived-mode-p 'c-mode 'c-ts-mode 'c++-mode 'c++-ts-mode)
    (let* ((text (my/citre-backend--context-text beginning))
           (regexp (format "\\(?:[(,;{]\\|^\\)[ \t\n]*\\(?:const[ \t]+\\)?\\(?:struct[ \t]+\\)?\\([[:alpha:]_][[:alnum:]_:]*\\(?:[ \t]*<[^>;\n]+>\\)?\\)[ \t*&]+%s\\_>"
                           (regexp-quote receiver))))
      (my/citre-backend--unique-match regexp text))))

(defun my/citre-backend--tag-scope (tag)
  "Return TAG's compact normalized scope."
  (my/citre-backend--normalize-type (citre-get-tag-field 'scope tag t)))

(defun my/citre-backend--scope-match-p (type tag)
  "Whether the nominal TYPE matches TAG's scope."
  (when-let* ((type (my/citre-backend--normalize-type type))
              (scope (my/citre-backend--tag-scope tag)))
    (or (equal scope type)
        (string-suffix-p (concat "." type) scope)
        (and (not (string-match-p "\\." type))
             (equal type (car (last (split-string scope "\\." t))))))))

(defun my/citre-backend--same-directory-p (tag)
  "Whether TAG belongs to the current file's directory."
  (when-let* ((current buffer-file-name)
              (path (citre-get-tag-field 'ext-abspath tag)))
    (equal (file-name-directory (file-truename current))
           (file-name-directory (file-truename path)))))

(defun my/citre-backend--current-scope ()
  "Return the current defun name when Emacs can determine it cheaply."
  (when (fboundp 'add-log-current-defun)
    (ignore-errors (add-log-current-defun))))

(defun my/citre-backend--tag-rank (tag context current-scope)
  "Return a lexicographically comparable rank for TAG in CONTEXT."
  (let* ((path (citre-get-tag-field 'ext-abspath tag))
         (scope (my/citre-backend--tag-scope tag))
         (type (plist-get context :type)))
    (list (if (my/citre-backend--scope-match-p type tag) 0 1)
          (if (and path buffer-file-name (equal (file-truename path) (file-truename buffer-file-name))) 0 1)
          (if (and current-scope scope (string-match-p (regexp-quote current-scope) scope)) 0 1)
          (if (my/citre-backend--same-directory-p tag) 0 1)
          (my/citre-backend--source-rank (citre-get-tag-field 'my-source tag))
          (or (citre-get-tag-field 'my-order tag) most-positive-fixnum))))

(defun my/citre-backend--rank-less-p (left right context current-scope)
  "Whether LEFT should sort before RIGHT in CONTEXT."
  (let ((a (my/citre-backend--tag-rank left context current-scope))
        (b (my/citre-backend--tag-rank right context current-scope)))
    (catch 'result
      (while a
        (unless (= (car a) (car b))
          (throw 'result (< (car a) (car b))))
        (setq a (cdr a) b (cdr b)))
      nil)))

(defun my/citre-backend--prepare-completions (tags context)
  "Context-filter, rank and deduplicate completion TAGS."
  (let* ((type (plist-get context :type))
         (matching (and type (seq-filter (lambda (tag) (my/citre-backend--scope-match-p type tag)) tags)))
         (visible (if (and (eq my/citre-completion-context-filter 'safe) matching) matching tags))
         (current-scope (my/citre-backend--current-scope)))
    (setq visible (sort visible (lambda (a b) (my/citre-backend--rank-less-p a b context current-scope))))
    (if (not (eq my/citre-completion-deduplication 'name))
        visible
      (let ((groups (make-hash-table :test #'equal))
            (seen (make-hash-table :test #'equal)) result)
        (dolist (tag tags)
          (push tag (gethash (citre-get-tag-field 'name tag) groups)))
        (dolist (tag visible (nreverse result))
          (let ((name (citre-get-tag-field 'name tag)))
            (unless (gethash name seen)
              (puthash name t seen)
              (citre-set-tag-field 'my-tags (nreverse (gethash name groups)) tag)
              (citre-set-tag-field 'my-count (length (citre-get-tag-field 'my-tags tag)) tag)
              (push tag result))))))))

(defun my/citre-backend--source-label (source)
  "Return a compact display label for SOURCE."
  (alist-get source '((current-buffer . "buf") (live-buffer . "live")
                      (project . "proj") (dependency . "dep") (tags-file . "tags"))))

(defun my/citre-backend--annotation (tag)
  "Return the compact completion annotation for TAG."
  (let ((count (or (citre-get-tag-field 'my-count tag) 1)) parts)
    (when (memq 'scope my/citre-completion-annotation)
      (when-let* ((scope (my/citre-backend--tag-scope tag))) (push scope parts)))
    (when (memq 'source my/citre-completion-annotation)
      (when-let* ((source (my/citre-backend--source-label (citre-get-tag-field 'my-source tag))))
        (push source parts)))
    (when (and (memq 'count my/citre-completion-annotation) (> count 1))
      (push (format "+%d" (1- count)) parts))
    (when parts (concat " " (string-join (nreverse parts) " · ")))))

(defun my/citre-backend--make-candidate (tag)
  "Make one Company/CAPF candidate from TAG."
  (citre-put-property
   (citre-make-tag-str tag nil '(name))
   'annotation (my/citre-backend--annotation tag)
   'signature (or (citre-get-tag-field 'signature tag)
                  (when-let* ((str (citre-get-tag-field 'extra-matched-str tag))) (string-trim str)))
   'kind (when-let* ((kind (citre-get-tag-field 'ext-kind-full tag))) (intern kind))
   'tags (citre-get-tag-field 'my-tags tag)))

(defun my/citre-backend--make-collection (tags start end)
  "Return a completion-at-point value for TAGS between START and END."
  (when-let* ((candidates (mapcar #'my/citre-backend--make-candidate tags)))
    (let ((collection (lambda (string predicate action)
                        (if (eq action 'metadata)
                            '(metadata (category . citre-completion)
                                       (cycle-sort-function . identity)
                                       (display-sort-function . identity))
                          (complete-with-action action candidates string predicate)))))
      (list start end collection
            :annotation-function (lambda (candidate) (citre-get-property 'annotation candidate))
            :company-docsig (lambda (candidate) (citre-get-property 'signature candidate))
            :company-kind (lambda (candidate) (citre-get-property 'kind candidate))
            :exclusive 'no))))

(defun my/citre-backend-completions ()
  "Return current Citre completion bounds and tags, reusing a narrowing prefix."
  (when-let* ((source (car (my/citre-backend--sources)))
              (symbol (citre-tags-get-symbol (plist-get source :file)))
              (bounds (citre-get-property 'bounds symbol)))
    (let* ((revision (and my/citre-index--state (my/citre-index-state-revision my/citre-index--state)))
           (member-context (my/citre-backend--member-context (car bounds)))
           (context (list revision (car bounds) major-mode citre-tags-substr-completion
                          (buffer-substring-no-properties (line-beginning-position) (car bounds))
                          member-context))
           (cache my/citre-backend--completion-cache)
           (tags (if (and revision (equal context (plist-get cache :context))
                          (string-prefix-p (plist-get cache :prefix) symbol))
                     (plist-get cache :tags)
                   (let ((tags (my/citre-backend--prepare-completions
                                (my/citre-backend--query 'completion) member-context)))
                     (setq my/citre-backend--completion-cache
                           (list :context context :prefix (substring-no-properties symbol)
                                 :tags tags))
                     tags))))
      (when tags (list (car bounds) (cdr bounds) tags)))))

(defun my/citre-completion-at-point ()
  "Complete with ranked, deduplicated Citre candidates."
  (when-let* ((result (if citre-capf-optimize-for-popup
                          (pcase (while-no-input (citre-get-completions)) ('t nil) (value value))
                        (citre-get-completions)))
              (beginning (nth 0 result))
              (end (nth 1 result))
              (tags (nth 2 result)))
    (if citre-capf-optimize-for-popup
        (pcase (while-no-input (my/citre-backend--make-collection tags beginning end))
          ('t nil) (value value))
      (my/citre-backend--make-collection tags beginning end))))

(defun my/citre-backend-definitions ()
  "Find definitions in the current live and persistent indexes."
  (my/citre-backend--query 'definition))

(defun my/citre-backend-definitions-of-id (identifier)
  "Find definitions of IDENTIFIER in the current project."
  (my/citre-backend--query 'definition identifier))

(defun my/citre-backend-identifiers ()
  "Return all current identifiers for Citre's query UI."
  (delete-dups (mapcar (lambda (tag) (citre-get-tag-field 'name tag))
                      (my/citre-backend--query 'identifiers))))

(defun my/citre-backend-imenu ()
  "Return the current file's symbols for Imenu, including unsaved changes."
  (my/citre-backend--query 'imenu))

(defun my/citre-backend--reference-root ()
  "Return the directory searched for references."
  (or (and my/citre-index--state (my/citre-index-state-root my/citre-index--state))
      (when-let* ((project (project-current nil))) (project-root project))
      default-directory))

(defun my/citre-backend-references-of-id (identifier)
  "Return textual references to IDENTIFIER below the project root.
Universal Ctags records definitions only, so references come from a
whole-word search that honours the project's ignore files."
  (when-let* ((program (executable-find my/citre-reference-program))
              (name (substring-no-properties identifier))
              (root (my/citre-backend--reference-root)))
    (unless (string-empty-p name)
      (let ((default-directory root) tags)
        (with-temp-buffer
          (process-file program nil t nil
                        "--no-heading" "--line-number" "--color=never"
                        "--word-regexp" "--fixed-strings" "--" name ".")
          (goto-char (point-min))
          (while (not (eobp))
            ;; Citre renders the matched line from the file itself, so only the
            ;; location is recorded here.
            (when (looking-at "\\([^:\n]+\\):\\([0-9]+\\):")
              (push (citre-make-tag
                     'name name
                     'ext-abspath (expand-file-name (match-string-no-properties 1) root)
                     'line (match-string-no-properties 2)
                     'extras "reference")
                    tags))
            (forward-line 1)))
        (nreverse tags)))))

(defun my/citre-backend-references ()
  "Return textual references to the symbol at point."
  (when-let* ((symbol (citre-tags-symbol-at-point)))
    (my/citre-backend-references-of-id symbol)))

(defun my/citre-backend-usable-p ()
  "Whether the current buffer has an automatic or conventional tags index."
  (or my/citre-index--state (citre-tags-usable-p)))

(defun my/citre-backend--updated (state)
  "Invalidate data for STATE without restarting or rearranging an open Company menu."
  (dolist (buffer (my/citre-index--buffers state))
    (with-current-buffer buffer
      (setq my/citre-backend--completion-cache nil)
      (when (boundp 'imenu--index-alist) (setq imenu--index-alist nil))
      ;; Fresh data is consumed on the next edit or completion request. Navigation keeps its current candidates and selection.
      (when (and (boundp 'company--capf-cache)
                 (eq (car company--capf-cache) buffer))
        (setq company--capf-cache nil)))))

(citre-register-backend
 'my-index
 (citre-backend-create :usable-probe #'my/citre-backend-usable-p
                      :symbol-at-point-fn #'citre-tags-symbol-at-point
                      :completions-fn #'my/citre-backend-completions
                      :defs-fn #'my/citre-backend-definitions
                      :defs-of-id-fn #'my/citre-backend-definitions-of-id
                      :refs-fn #'my/citre-backend-references
                      :refs-of-id-fn #'my/citre-backend-references-of-id
                      :id-list-fn #'my/citre-backend-identifiers
                      :tags-in-buffer-fn #'my/citre-backend-imenu
                      :after-jump-fn #'my/citre-index-attach))

(add-hook 'my/citre-index-updated-hook #'my/citre-backend--updated)

(provide 'citre-backend)
;;; citre-backend.el ends here