;;; jabber-mam.el --- XEP-0313 Message Archive Management  -*- lexical-binding: t; -*-

;; Copyright (C) 2026 Thanos Apollo

;; Author: Thanos Apollo <public@thanosapollo.org>
;; Maintainer: Thanos Apollo <public@thanosapollo.org>

;; This file is part of emacs-jabber.

;; emacs-jabber is free software: you can redistribute it and/or modify
;; it under the terms of the GNU General Public License as published by
;; the Free Software Foundation, either version 3 of the License, or
;; (at your option) any later version.

;; emacs-jabber is distributed in the hope that it will be useful,
;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
;; GNU General Public License for more details.

;; You should have received a copy of the GNU General Public License
;; along with emacs-jabber.  If not, see <https://www.gnu.org/licenses/>.

;;; Commentary:

;; XEP-0313 (Message Archive Management) support.
;; Queries the server's message archive on connect to sync missed
;; messages across devices.  Results are stored in the local database
;; and displayed in open chat buffers.
;;
;; Pagination uses XEP-0059 (Result Set Management).
;; Deduplication uses XEP-0359 stanza-id / server-id.

;;; Code:

(require 'jabber-xml)
(require 'jabber-util)
(require 'jabber-db)
(require 'jabber-chat)
(require 'jabber-disco)
(require 'jabber-iq)
(require 'jabber-message-correct)
(require 'jabber-muc-state)

(eval-when-compile (require 'cl-lib))

(defvar jabber-buffer-connection)       ; jabber-chatbuffer.el
(defvar jabber-chatting-with)           ; jabber-chat.el
(defvar jabber-group)                   ; jabber-muc.el

;;; Constants

(defconst jabber-mam-xmlns "urn:xmpp:mam:2"
  "Namespace for XEP-0313 MAM.")

(defconst jabber-mam-rsm-xmlns "http://jabber.org/protocol/rsm"
  "Namespace for XEP-0059 Result Set Management.")

(defconst jabber-mam-forward-xmlns "urn:xmpp:forward:0"
  "Namespace for XEP-0297 Stanza Forwarding.")

(defconst jabber-mam-delay-xmlns "urn:xmpp:delay"
  "Namespace for XEP-0203 Delayed Delivery.")

;;; Customization

(defcustom jabber-mam-enable t
  "Whether to sync messages via MAM on connect."
  :type 'boolean
  :group 'jabber)

(defcustom jabber-mam-page-size 50
  "Number of messages to request per MAM page."
  :type 'integer
  :group 'jabber)

(defcustom jabber-mam-catch-up-days 3
  "Limit initial MAM catch-up to this many days back.
Only used when no previous sync point exists (first sync).
Set to nil to fetch the entire archive."
  :type '(choice integer (const :tag "Fetch all" nil))
  :group 'jabber)

;;; Hooks

(defvar jabber-mam-peer-syncing-functions nil
  "Hook run when a peer's MAM sync state changes.
Each function receives three arguments: PEER (bare JID), TYPE
\(\"groupchat\" or \"chat\"), and SYNCING-P (non-nil when sync
starts, nil when it ends).")

(defvar jabber-mam-sync-complete-functions nil
  "Hook run after MAM stores messages for one or more peers.
Each function receives one argument: a list of (ACCOUNT PEER TYPE) entries.")

;;; Internal state

(defvar jabber-mam--syncing nil
  "Non-nil while a MAM sync is in progress.
Alist of (JC . QUERYID) for active queries.")

(defvar jabber-mam--dirty-peers nil
  "Peers that received MAM messages during sync.
Each entry is (ACCOUNT PEER TYPE), where ACCOUNT is the local bare JID.
Accumulated during sync, drained after COMMIT.")

(defvar jabber-mam--tx-depth 0
  "Reference count for the shared MAM transaction.
BEGIN when 0->1, COMMIT when 1->0.  Allows concurrent MAM queries
to share one SQLite transaction.")

(defvar jabber-mam--completion-callbacks nil
  "Alist of (QUERYID . CALLBACK) for per-query completion hooks.
CALLBACK is called with no arguments when the query finishes.")

(defvar jabber-mam--peer-syncing nil
  "Alist of active automatic peer catch-ups.
Each entry has the shape ((JC PEER) . TOKEN).")

(defvar jabber-mam--sync-received nil
  "Alist of (QUERYID . PLIST) for sync-buffer reconciliation.
PLIST keys: :ids (hash-table), :min-ts, :max-ts, :account, :peer.
Populated during sync; consumed by `jabber-mam--reconcile-sync'.")

(defvar jabber-mam--query-targets nil
  "Alist of (QUERYID . TARGET) for active MAM queries.
TARGET is a room JID for MUC MAM, or nil for 1:1 MAM.")

;;; Public predicates

(defun jabber-mam-syncing-p ()
  "Return non-nil if any MAM sync is in progress."
  (not (null jabber-mam--syncing)))

;;; Query building

(defvar jabber-mam--queryid-counter 0
  "Monotonic counter for unique MAM query IDs.")

(defun jabber-mam--make-queryid ()
  "Generate a unique query ID for MAM."
  (format "mam-%d-%d"
          (cl-incf jabber-mam--queryid-counter)
          (floor (float-time))))

(defun jabber-mam--begin-peer-sync (jc peer)
  "Record an automatic MAM catch-up for JC and PEER.
Return its unique token, or nil when one was already active."
  (let ((key (list jc peer)))
    (unless (assoc key jabber-mam--peer-syncing #'equal)
      (let ((token (make-symbol "jabber-mam-peer-sync-")))
        (push (cons key token) jabber-mam--peer-syncing)
        token))))

(defun jabber-mam--peer-sync-active-p (jc peer token)
  "Return non-nil when TOKEN owns JC and PEER's catch-up."
  (eq token
      (cdr (assoc (list jc peer) jabber-mam--peer-syncing #'equal))))

(defun jabber-mam--end-peer-sync (jc peer token)
  "Forget TOKEN's automatic MAM catch-up for JC and PEER."
  (when-let* ((entry (assoc (list jc peer)
                            jabber-mam--peer-syncing #'equal))
              ((eq token (cdr entry))))
    (setq jabber-mam--peer-syncing
          (delq entry jabber-mam--peer-syncing))
    t))

(defun jabber-mam--finish-peer-sync (jc peer token)
  "Finish TOKEN's peer catch-up for JC and PEER."
  (when (jabber-mam--end-peer-sync jc peer token)
    (run-hook-with-args 'jabber-mam-peer-syncing-functions
                        peer "chat" nil)))

(defun jabber-mam--complete-query (queryid)
  "Remove and run the completion callback for QUERYID."
  (when-let* ((cb (assoc queryid jabber-mam--completion-callbacks
                         #'string=)))
    (setq jabber-mam--completion-callbacks
          (delq cb jabber-mam--completion-callbacks))
    (condition-case err
        (funcall (cdr cb))
      (error
       (message "MAM: completion callback error: %s"
                (error-message-string err))))
    t))

(defun jabber-mam--build-query (queryid &optional with start after-id max
                                        before-id)
  "Build a MAM <query> sexp.
QUERYID is echoed in results for correlation.
WITH filters by JID, START is an XEP-0082 datetime string.
AFTER-ID is an RSM cursor for forward pagination.
MAX is the page size.
BEFORE-ID is an RSM cursor for backward pagination; when t, emit
an empty <before/> element (meaning \"last page\")."
  (let ((form-fields nil)
        (rsm-children nil))
    ;; Data form fields
    (when (or with start)
      (push `(field ((var . "FORM_TYPE") (type . "hidden"))
                    (value () ,jabber-mam-xmlns))
            form-fields)
      (when with
        (push `(field ((var . "with"))
                      (value () ,with))
              form-fields))
      (when start
        (push `(field ((var . "start"))
                      (value () ,start))
              form-fields))
      (setq form-fields (nreverse form-fields)))
    ;; RSM
    (when max
      (push `(max () ,(number-to-string max)) rsm-children))
    (when after-id
      (push `(after () ,after-id) rsm-children))
    (when before-id
      (if (eq before-id t)
          (push '(before ()) rsm-children)
        (push `(before () ,before-id) rsm-children)))
    (setq rsm-children (nreverse rsm-children))
    ;; Build query
    `(query ((xmlns . ,jabber-mam-xmlns)
             (queryid . ,queryid))
            ,@(when form-fields
                (list `(x ((xmlns . "jabber:x:data")
                           (type . "submit"))
                          ,@form-fields)))
            ,@(when rsm-children
                (list `(set ((xmlns . ,jabber-mam-rsm-xmlns))
                            ,@rsm-children))))))

;;; Result parsing

(defun jabber-mam--parse-result (xml-data)
  "Extract MAM result from the <message> stanza XML-DATA.
Return (ARCHIVE-ID DELAY-STAMP INNER-MESSAGE) or nil."
  (when-let* ((result-el (jabber-xml-child-with-xmlns
                          xml-data jabber-mam-xmlns)))
    (let* ((archive-id (jabber-xml-get-attribute result-el 'id))
           (fwd-el (car (jabber-xml-get-children result-el 'forwarded)))
           (delay-el (and fwd-el
                          (car (jabber-xml-get-children fwd-el 'delay))))
           (stamp (and delay-el
                       (jabber-xml-get-attribute delay-el 'stamp)))
           (inner-msg (and fwd-el
                           (car (jabber-xml-get-children fwd-el 'message)))))
      (when inner-msg
        (list archive-id stamp inner-msg)))))

(defun jabber-mam--parse-fin (xml-data)
  "Parse a MAM <fin> IQ result XML-DATA.
Return plist (:complete BOOL :first ID :last ID)."
  (let* ((fin-el (jabber-xml-child-with-xmlns xml-data jabber-mam-xmlns))
         (complete (string= (or (jabber-xml-get-attribute fin-el 'complete) "")
                            "true"))
         (set-el (and fin-el
                      (car (jabber-xml-get-children fin-el 'set))))
         (first-el (and set-el
                        (car (jabber-xml-get-children set-el 'first))))
         (last-el (and set-el
                       (car (jabber-xml-get-children set-el 'last))))
         (first-id (and first-el
                        (car (jabber-xml-node-children first-el))))
         (last-id (and last-el
                       (car (jabber-xml-node-children last-el)))))
    (list :complete complete :first first-id :last last-id)))

;;; Message chain handler

(defun jabber-mam--unwrap-into (outer inner &optional archive-id)
  "Replace OUTER stanza's attributes and children with INNER's.
Marks the stanza as MAM-origin so downstream handlers can suppress
outgoing receipts.  ARCHIVE-ID is preserved for bodyless archive
stanzas whose downstream handlers need the MAM result id.  Mutates
OUTER in place."
  (setcar (cdr outer) (append (jabber-xml-node-attributes inner)
                              `((jabber-mam--origin . "t")
                                ,@(and archive-id
                                       `((jabber-mam--archive-id
                                          . ,archive-id))))))
  (setcdr (cdr outer) (cddr inner)))

(defun jabber-mam--active-query-p (queryid)
  "Return non-nil if QUERYID is an active MAM query."
  (cl-find queryid jabber-mam--syncing :key #'cdr :test #'string=))

(defun jabber-mam--query-target (queryid)
  "Return the archive target for QUERYID, excluding control markers."
  (cdr
   (cl-find-if
    (lambda (entry)
      (and (string= queryid (car entry))
           (not (eq (cdr entry) 'one-shot))))
    jabber-mam--query-targets)))

(defun jabber-mam--valid-sender-p (jc from queryid)
  "Return non-nil if FROM is a valid MAM result sender for JC.
Valid senders are the entity recorded for QUERYID: our own bare JID
for user archives, or the queried room JID for room archives.  A
nil FROM is accepted only for own-archive queries because some
servers omit the attribute when the message originates from the
user's own archive.
QUERYID identifies the active query for target lookup."
  (let ((target (jabber-mam--query-target queryid)))
    (if (null from)
        (null target)
      (let ((bare (jabber-jid-user from))
            (our-jid (jabber-connection-bare-jid jc)))
        (if target
            (string= bare target)
          (string= bare our-jid))))))

(defun jabber-mam--classify-direction (jc from to type)
  "Classify message direction and peer from MAM result fields.
JC is the connection, FROM/TO are stanza JIDs, TYPE is message type.
Return (DIRECTION . PEER) where DIRECTION is \"in\" or \"out\"."
  (let* ((our-jid (jabber-connection-bare-jid jc))
         (groupchat-p (string= type "groupchat"))
         (direction (if groupchat-p
			(let ((nick (jabber-jid-resource from))
                              (room (jabber-jid-user from)))
                          (if (and nick
                                   (jabber-mam--our-muc-nick-p
                                    room nick jc))
                              "out" "in"))
                      (if (string= (jabber-jid-user from) our-jid)
                          "out" "in")))
         (peer (if groupchat-p
                   (jabber-jid-user from)
                 (jabber-jid-user
                  (if (string= direction "out") to from)))))
    (cons direction peer)))

(defun jabber-mam--extract-fields (jc inner-msg stamp)
  "Extract message fields from INNER-MSG for storage.
JC is the connection.  STAMP is the MAM delay timestamp string.
Returns a plist with :from :to :type :body :stanza-id :our-jid
:direction :peer :timestamp :oob-entries, or nil if direction
cannot be determined."
  (let* ((from (jabber-xml-get-attribute inner-msg 'from))
         (to (jabber-xml-get-attribute inner-msg 'to))
         (type (or (jabber-xml-get-attribute inner-msg 'type) "chat"))
         (body-el (car (jabber-xml-get-children inner-msg 'body)))
         (body (and body-el (car (jabber-xml-node-children body-el))))
         (stanza-id (jabber-xml-get-attribute inner-msg 'id))
         (our-jid (jabber-connection-bare-jid jc))
         (dir-peer (jabber-mam--classify-direction jc from to type))
         (direction (car dir-peer))
         (peer (cdr dir-peer))
         (timestamp (and stamp (jabber-parse-time stamp)))
         (oob-entries (jabber-db--extract-oob-entries inner-msg)))
    (list :from from :to to :type type :body body
          :stanza-id stanza-id :our-jid our-jid
          :direction direction :peer peer
          :timestamp timestamp :oob-entries oob-entries)))

(defun jabber-mam--track-sync-ids (qid archive-id stanza-id ts)
  "Update sync-received tracking for query QID.
ARCHIVE-ID and STANZA-ID are recorded as seen.  TS updates the
min/max timestamp range when non-nil."
  (when-let* ((sync-data (cdr (assoc qid jabber-mam--sync-received
                                     #'string=)))
              (ids (plist-get sync-data :ids)))
    (when archive-id (puthash archive-id t ids))
    (when stanza-id (puthash stanza-id t ids))
    (when ts
      (when (or (null (plist-get sync-data :min-ts))
                (< ts (plist-get sync-data :min-ts)))
        (plist-put sync-data :min-ts ts))
      (when (or (null (plist-get sync-data :max-ts))
                (> ts (plist-get sync-data :max-ts)))
        (plist-put sync-data :max-ts ts)))))

(defun jabber-mam--store-new-message-p (jc inner-msg)
  "Return non-nil when INNER-MSG should be stored as a new message for JC."
  (not (run-hook-with-args-until-success
        'jabber-history-inhibit-received-message-functions
        jc inner-msg)))

(defun jabber-mam--message-action (jc inner-msg fields)
  "Return how MAM should handle INNER-MSG with FIELDS for JC."
  (cond
   ((or (not (plist-get fields :peer))
        (not (plist-get fields :body)))
    :unwrap)
   ((jabber-message-correct--replace-id inner-msg) :correct)
   ((jabber-mam--store-new-message-p jc inner-msg) :store)
   (t :unwrap)))

(defun jabber-mam--process-message (jc xml-data)
  "Handle a MAM result <message> from the message chain.
JC is the Jabber connection.  XML-DATA is the stanza."
  (when-let* ((result-el (jabber-xml-child-with-xmlns
                          xml-data jabber-mam-xmlns))
              (qid (jabber-xml-get-attribute result-el 'queryid))
              ((jabber-mam--active-query-p qid))
              (parsed (jabber-mam--parse-result xml-data))
              ((jabber-mam--valid-sender-p
                jc (jabber-xml-get-attribute xml-data 'from) qid)))
    (let* ((archive-id (nth 0 parsed))
           (stamp (nth 1 parsed))
           (inner-msg (nth 2 parsed))
           (encrypted (jabber-xml-encrypted-p inner-msg))
           (inner-msg
            (let ((inhibit-message t)
                  (message-log-max nil))
              (jabber-chat--decrypt-if-needed jc inner-msg)))
           (fields (jabber-mam--extract-fields jc inner-msg stamp))
           (peer (plist-get fields :peer))
           (body (plist-get fields :body))
           (timestamp (and-let* ((time (plist-get fields :timestamp)))
                        (floor (float-time time))))
           (action (jabber-mam--message-action jc inner-msg fields)))
      (pcase action
        ((or :correct :store)
         (pcase action
           (:correct
            (unless (jabber--decrypt-failure-body-p body)
              (jabber-message-correct--apply
               (jabber-message-correct--replace-id inner-msg)
               body (plist-get fields :from)
               (string= (plist-get fields :type) "groupchat") nil
               (jabber-db--extract-occupant-id inner-msg)
               (plist-get fields :our-jid) peer nil)))
           (:store
            (let ((jabber-db-message-thread-stored-functions nil))
              (jabber-db-store-message
               (plist-get fields :our-jid) peer
               (plist-get fields :direction) (plist-get fields :type)
               body timestamp (jabber-jid-resource (plist-get fields :from))
               (plist-get fields :stanza-id) archive-id
               (jabber-db--extract-occupant-id inner-msg)
               (plist-get fields :oob-entries) encrypted
               (jabber-db--extract-reply-fields inner-msg)
               (jabber-db--extract-thread-fields inner-msg)))))
         (jabber-mam--track-sync-ids qid archive-id
                                     (plist-get fields :stanza-id) timestamp)
         (jabber-mam--mark-dirty jc peer (plist-get fields :type))
         (setcdr (cdr xml-data) nil))
        (:unwrap
         (jabber-mam--unwrap-into xml-data inner-msg archive-id))))))

(defun jabber-mam--our-muc-nick-p (room nick jc)
  "Return non-nil if NICK in ROOM is us on connection JC.
Checks the current room nickname first, then falls back to
comparing with the account username to handle nick changes."
  (require 'jabber-muc)
  (or (and-let* ((my-nick (jabber-muc-nickname room jc)))
        (string= nick my-nick))
      (string= nick (plist-get (fsm-get-state-data jc) :username))))


(defun jabber-mam--mark-dirty (jc peer type)
  "Record that PEER's buffer needs redisplay after sync.
JC identifies the local account.  TYPE is the message type."
  (let ((entry (list (jabber-connection-bare-jid jc) peer type)))
    (cl-pushnew entry jabber-mam--dirty-peers :test #'equal)))


(defun jabber-mam--redraw-dirty ()
  "Signal that accumulated dirty peers need display refresh.
Drains `jabber-mam--dirty-peers' and runs
`jabber-mam-sync-complete-functions'."
  (let ((peers (prog1 jabber-mam--dirty-peers
                 (setq jabber-mam--dirty-peers nil))))
    (when peers
      (run-hook-with-args 'jabber-mam-sync-complete-functions peers))))

;;; Shared transaction management

(defun jabber-mam--tx-begin ()
  "Increment the MAM transaction ref count.
BEGIN a SQLite transaction when transitioning from 0 to 1."
  (when (zerop jabber-mam--tx-depth)
    (when-let* ((db (jabber-db-ensure-open)))
      (sqlite-execute db "BEGIN")))
  (cl-incf jabber-mam--tx-depth))

(defun jabber-mam--tx-end ()
  "Decrement the MAM transaction ref count.
COMMIT the SQLite transaction when transitioning from 1 to 0."
  (when (> jabber-mam--tx-depth 0)
    (cl-decf jabber-mam--tx-depth)
    (when (zerop jabber-mam--tx-depth)
      (when-let* ((db (jabber-db-ensure-open)))
        (sqlite-execute db "COMMIT")))))

;;; Query and pagination

(defun jabber-mam--query (jc &optional after-id queryid with start to
                             before-id max)
  "Send a MAM query via JC, paginating from AFTER-ID.
QUERYID correlates results; generated if nil.
WITH and START are optional filters.
TO is the query target; nil for user archive, a room JID for MUC MAM.
BEFORE-ID and MAX support backward pagination (last-page queries).
When BEFORE-ID is non-nil, the query is one-shot (no forward pagination)."
  (let ((queryid (or queryid (jabber-mam--make-queryid)))
        (page-size (or max jabber-mam-page-size)))
    (push (cons jc queryid) jabber-mam--syncing)
    (when to
      (push (cons queryid to) jabber-mam--query-targets))
    ;; Mark one-shot queries so handle-fin skips forward pagination.
    (when before-id
      (push (cons queryid 'one-shot) jabber-mam--query-targets))
    ;; Open a shared transaction for concurrent MAM queries.
    ;; COMMIT happens when the last active query finishes.
    (jabber-mam--tx-begin)
    (condition-case err
        (jabber-send-iq
         jc to "set"
         (jabber-mam--build-query queryid with start after-id
                                  page-size before-id)
         #'jabber-mam--handle-fin
         (list queryid with start to)
         #'jabber-mam--handle-error
         (list queryid to))
      (error
       (jabber-mam--tx-end)
       (setq jabber-mam--syncing
             (cl-remove queryid jabber-mam--syncing
                        :key #'cdr :test #'string=))
       (setq jabber-mam--query-targets
             (cl-remove queryid jabber-mam--query-targets
                        :key #'car :test #'string=))
       (jabber-mam--complete-query queryid)
       (message "MAM: query failed to send: %s"
                (error-message-string err))))))

(defun jabber-mam--handle-fin (jc xml-data closure)
  "Handle the <fin> IQ result for a MAM query.
JC is the connection.  XML-DATA is the IQ response.
CLOSURE is (QUERYID WITH START TO)."
  (let* ((queryid (nth 0 closure))
         (with (nth 1 closure))
         (start (nth 2 closure))
         (to (nth 3 closure))
         (fin (jabber-mam--parse-fin xml-data))
         (complete (plist-get fin :complete))
         (last-id (plist-get fin :last)))
    (jabber-mam--tx-end)
    ;; Remove from syncing list
    (setq jabber-mam--syncing
          (cl-remove queryid jabber-mam--syncing
                     :key #'cdr :test #'string=))
    (let ((one-shot-p (assoc queryid jabber-mam--query-targets
                             #'string=)))
      ;; One-shot queries (before-id based) never paginate forward.
      (setq one-shot-p (and one-shot-p
                            (eq (cdr one-shot-p) 'one-shot)))
      (if (or complete (null last-id) one-shot-p)
          (progn
            ;; Clean up query target tracking.
            (setq jabber-mam--query-targets
                  (cl-remove queryid jabber-mam--query-targets
                             :key #'car :test #'string=))
            (let ((inhibit-message t))
              (message "MAM: sync complete%s"
                       (if to (format " for %s" to)
                         (if with (format " for %s" with) ""))))
            (jabber-mam--complete-query queryid)
            ;; Redraw affected buffers from DB.
            (jabber-mam--redraw-dirty))
        ;; More pages: yield to the event loop for redisplay and input,
        ;; then continue pagination.
        (run-with-timer 0.1 nil #'jabber-mam--query
                        jc last-id queryid with start to)))))

(defun jabber-mam--handle-error (jc xml-data closure)
  "Handle a MAM query error.
JC is the connection.  XML-DATA is the IQ error.
CLOSURE is (QUERYID TO).
On item-not-found (stale sync point), falls back to time-based query."
  (let ((queryid (car closure))
        (to (cadr closure)))
    (jabber-mam--tx-end)
    (setq jabber-mam--syncing
          (cl-remove queryid jabber-mam--syncing
                     :key #'cdr :test #'string=))
    (setq jabber-mam--query-targets
          (cl-remove queryid jabber-mam--query-targets
                     :key #'car :test #'string=))
    (let ((error-el (car (jabber-xml-get-children xml-data 'error))))
      (if (and error-el
               (car (jabber-xml-get-children error-el 'item-not-found)))
          (progn
            (let ((inhibit-message t))
              (message "MAM: sync point expired%s, falling back to time-based query"
                       (if to (format " for %s" to) "")))
            ;; Transfer completion callback to the fallback query.
            (let* ((old-cb (assoc queryid jabber-mam--completion-callbacks
                                  #'string=))
                   (new-queryid (jabber-mam--make-queryid))
                   (start (when jabber-mam-catch-up-days
                            (format-time-string
                             "%Y-%m-%dT%H:%M:%SZ"
                             (time-subtract (current-time)
                                            (* jabber-mam-catch-up-days 86400))
                             t))))
              (when old-cb
                (setq jabber-mam--completion-callbacks
                      (delq old-cb jabber-mam--completion-callbacks))
                (push (cons new-queryid (cdr old-cb))
                      jabber-mam--completion-callbacks))
              (jabber-mam--query jc nil new-queryid nil start to)))
        ;; Permanent error: fire completion callback so callers aren't stuck.
        (jabber-mam--complete-query queryid)
        (message "MAM: query failed: %s"
                 (jabber-sexp2xml xml-data))))))

;;; Post-connect catch-up

(defun jabber-mam--catch-up (jc)
  "Sync missed messages for JC via MAM."
  (let* ((account (jabber-connection-bare-jid jc))
         (last-id (jabber-db-last-server-id account)))
    (if last-id
        ;; Resume from last known server-id
        (jabber-mam--query jc last-id)
      ;; First sync: limit to N days back
      (let ((start (when jabber-mam-catch-up-days
                     (format-time-string
                      "%Y-%m-%dT%H:%M:%SZ"
                      (time-subtract (current-time)
                                     (* jabber-mam-catch-up-days 86400))
                      t))))
        (jabber-mam--query jc nil nil nil start)))))

(defun jabber-mam-maybe-catchup (jc)
  "Post-connect hook on JC: sync messages via MAM if enabled.
Added to `jabber-post-connect-hooks'."
  (when jabber-mam-enable
    (jabber-disco-get-info
     jc (jabber-connection-bare-jid jc) nil
     (lambda (jc _closure-data result)
       (when (and (listp result)
                  (not (eq (car result) 'error))
                  (member jabber-mam-xmlns (cadr result)))
         (jabber-mam--catch-up jc)))
     nil)))

;;; 1:1 chat MAM catch-up

(defun jabber-mam--chat-catch-up (jc peer token)
  "Sync missed messages for PEER via MAM.
JC is the Jabber connection.  PEER is the bare JID.  TOKEN identifies
the automatic peer sync attempt.
Registers a completion callback to clear the syncing indicator."
  (let* ((account (jabber-connection-bare-jid jc))
         (last-id (jabber-db-last-server-id account peer))
         (queryid (jabber-mam--make-queryid)))
    (push (cons queryid
                (lambda ()
                  (jabber-mam--finish-peer-sync jc peer token)))
          jabber-mam--completion-callbacks)
    (if last-id
        (jabber-mam--query jc last-id queryid peer nil nil)
      (let ((start (when jabber-mam-catch-up-days
                     (format-time-string
                      "%Y-%m-%dT%H:%M:%SZ"
                      (time-subtract (current-time)
                                     (* jabber-mam-catch-up-days 86400))
                      t))))
        (jabber-mam--query jc nil queryid peer start nil)))))

(defun jabber-mam--handle-chat-disco (jc closure-data result)
  "Start a peer catch-up after disco completes.
JC is the connection.  CLOSURE-DATA is (PEER TOKEN).  RESULT is
the disco response."
  (pcase-let ((`(,peer ,token) closure-data))
    (when (jabber-mam--peer-sync-active-p jc peer token)
      (if (and (listp result)
               (not (eq (car result) 'error))
               (member jabber-mam-xmlns (cadr result)))
          (condition-case err
              (jabber-mam--chat-catch-up jc peer token)
            (error
             (jabber-mam--finish-peer-sync jc peer token)
             (message "MAM: peer catch-up failed to start: %s"
                      (error-message-string err))))
        (jabber-mam--finish-peer-sync jc peer token)))))

(defun jabber-mam-chat-opened (jc peer)
  "Trigger 1:1 MAM catch-up when opening a chat with PEER.
JC is the Jabber connection.  Called from `jabber-chat-with'.
Sets the syncing indicator immediately; clears it when the catch-up
query completes (or when disco reveals MAM is not supported)."
  (when-let* ((token (and jabber-mam-enable
                          (jabber-mam--begin-peer-sync jc peer))))
    (run-hook-with-args 'jabber-mam-peer-syncing-functions peer "chat" t)
    (condition-case err
        (jabber-disco-get-info
         jc (jabber-connection-bare-jid jc) nil
         #'jabber-mam--handle-chat-disco (list peer token))
      (error
       (jabber-mam--finish-peer-sync jc peer token)
       (signal (car err) (cdr err))))))

;;; MUC MAM catch-up

(defun jabber-mam--muc-catch-up (jc group)
  "Sync missed messages for GROUP via MUC MAM.
JC is the Jabber connection.  GROUP is the room bare JID.
Registers a completion callback to clear the syncing indicator."
  (let* ((account (jabber-connection-bare-jid jc))
         (last-id (jabber-db-last-server-id account group))
         (queryid (jabber-mam--make-queryid)))
    (push (cons queryid
                (lambda ()
                  (run-hook-with-args 'jabber-mam-peer-syncing-functions
                                      group "groupchat" nil)))
          jabber-mam--completion-callbacks)
    (if last-id
        (jabber-mam--query jc last-id queryid nil nil group)
      (let ((start (when jabber-mam-catch-up-days
                     (format-time-string
                      "%Y-%m-%dT%H:%M:%SZ"
                      (time-subtract (current-time)
                                     (* jabber-mam-catch-up-days 86400))
                      t))))
        (jabber-mam--query jc nil queryid nil start group)))))

(defun jabber-mam-muc-joined (jc group)
  "Trigger MUC MAM catch-up after joining GROUP.
JC is the Jabber connection.  Called from MUC self-presence handler.
Sets the syncing indicator immediately; clears it when the catch-up
query completes (or when disco reveals MAM is not supported)."
  (when jabber-mam-enable
    (run-hook-with-args 'jabber-mam-peer-syncing-functions
                        group "groupchat" t)
    (jabber-disco-get-info
     jc group nil
     (lambda (jc closure-data result)
       (let ((group (car closure-data)))
         (if (and (listp result)
                  (not (eq (car result) 'error))
                  (member jabber-mam-xmlns (cadr result)))
             (jabber-mam--muc-catch-up jc group)
           (run-hook-with-args 'jabber-mam-peer-syncing-functions
                               group "groupchat" nil))))
     (list group))))

(defun jabber-mam--reconcile-sync (queryid)
  "Delete local messages not found in the remote archive for QUERYID.
Uses the IDs and timestamp range accumulated during sync to find
local messages that the server no longer has."
  (when-let* ((entry (assoc queryid jabber-mam--sync-received #'string=)))
    (let* ((data (cdr entry))
           (ids (plist-get data :ids))
           (min-ts (plist-get data :min-ts))
           (max-ts (plist-get data :max-ts))
           (account (plist-get data :account))
           (peer (plist-get data :peer)))
      (when (and min-ts max-ts (> (hash-table-count ids) 0))
        (when-let* ((db (jabber-db-ensure-open)))
          (let ((local-rows
                 (sqlite-select db "\
SELECT id, stanza_id, server_id FROM message \
WHERE account = ? AND peer = ? AND timestamp BETWEEN ? AND ? \
AND retracted_by IS NULL"
                                (list account peer min-ts max-ts)))
                (deleted 0))
            (dolist (row local-rows)
              (let ((row-id (nth 0 row))
                    (sid (nth 1 row))
                    (svid (nth 2 row)))
                ;; Only consider messages that have a server-side ID.
                ;; Messages without IDs can't be compared.
                (when (and (or sid svid)
                           (not (and svid (gethash svid ids)))
                           (not (and sid (gethash sid ids))))
                  (sqlite-execute db "DELETE FROM message WHERE id = ?"
                                  (list row-id))
                  (cl-incf deleted))))
            (when (> deleted 0)
              (jabber-db-prune-empty-message-threads account peer)
              (message "MAM: removed %d messages not found on server"
                       deleted))))))
    (setq jabber-mam--sync-received
          (cl-remove queryid jabber-mam--sync-received
                     :key #'car :test #'string=))))


(defun jabber-mam-sync-buffer ()
  "Sync messages from the server archive for this buffer.
Uses `jabber-chat-buffer-msg-count' for the number of messages.
Fetches recent messages using RSM backward pagination.  New messages
are decrypted and stored; existing messages are preserved via dedup.
Failed-decrypt placeholders are replaced if decryption now succeeds.
Local messages in the synced time range whose IDs are not found on
the server are deleted.  The buffer is refreshed in place after sync."
  (interactive)
  (unless (memq jabber-buffer-connection jabber-connections)
    (user-error "Not connected"))
  (let* ((jc jabber-buffer-connection)
         (count (jabber-chat-buffer-msg-count))
         (group (bound-and-true-p jabber-group))
         (peer (or group
                   (jabber-jid-user (bound-and-true-p jabber-chatting-with))))
         (account (jabber-connection-bare-jid jc))
         (muc-p (not (null group)))
         (queryid (jabber-mam--make-queryid)))
    ;; Register ID tracking for post-sync reconciliation.
    (push (cons queryid (list :ids (make-hash-table :test #'equal)
                              :min-ts nil :max-ts nil
                              :account account :peer peer))
          jabber-mam--sync-received)
    (let ((type (if group "groupchat" "chat")))
      (push (cons queryid
                  (lambda ()
                    (jabber-mam--reconcile-sync queryid)
                    (run-hook-with-args 'jabber-mam-peer-syncing-functions
                                        peer type nil)))
            jabber-mam--completion-callbacks)
      (run-hook-with-args 'jabber-mam-peer-syncing-functions peer type t))
    (jabber-mam--mark-dirty jc peer (if group "groupchat" "chat"))
    (message "MAM: syncing last %d messages for %s..." count peer)
    (if muc-p
        (jabber-mam--query jc nil queryid nil nil peer t count)
      (jabber-mam--query jc nil queryid peer nil nil t count))))

;;; Disconnect cleanup

(defun jabber-mam--cleanup-connection (jc)
  "Clean up MAM state for connection JC.
Called from `jabber-lost-connection-hooks' on involuntary disconnect."
  (dolist (entry (copy-sequence jabber-mam--peer-syncing))
    (when (eq (caar entry) jc)
      (jabber-mam--finish-peer-sync
       jc (cadar entry) (cdr entry))))
  (let ((jc-queries (cl-remove-if-not
                     (lambda (entry) (eq (car entry) jc))
                     jabber-mam--syncing)))
    (when jc-queries
      (setq jabber-mam--syncing
            (cl-set-difference jabber-mam--syncing jc-queries))
      (condition-case nil
          (dotimes (_ (length jc-queries))
            (jabber-mam--tx-end))
        (error nil))
      ;; Fire and remove leaked completion callbacks and query targets.
      (dolist (entry jc-queries)
        (let ((qid (cdr entry)))
          (jabber-mam--complete-query qid)
          (setq jabber-mam--query-targets
                (cl-remove qid jabber-mam--query-targets
                           :key #'car :test #'string=))))
      ;; Redraw affected buffers.
      (jabber-mam--redraw-dirty))))

(defun jabber-mam--cleanup-all ()
  "Clean up all MAM state on voluntary disconnect.
Called from `jabber-pre-disconnect-hook'."
  (condition-case nil
      (dotimes (_ jabber-mam--tx-depth)
        (jabber-mam--tx-end))
    (error nil))
  ;; Fire remaining completion callbacks to clear syncing flags.
  (dolist (cb (copy-sequence jabber-mam--completion-callbacks))
    (jabber-mam--complete-query (car cb)))
  (dolist (entry (copy-sequence jabber-mam--peer-syncing))
    (jabber-mam--finish-peer-sync
     (caar entry) (cadar entry) (cdr entry)))
  (setq jabber-mam--syncing nil
        jabber-mam--tx-depth 0
        jabber-mam--completion-callbacks nil
        jabber-mam--peer-syncing nil
        jabber-mam--query-targets nil
        jabber-mam--sync-received nil)
  (jabber-mam--redraw-dirty))

;;; MUC query cancellation

(defun jabber-mam--cancel-muc-query (room)
  "Cancel any active MUC MAM query for ROOM.
Removes the query from syncing state and decrements the transaction
depth.  Called when leaving a room to stop wasting bandwidth."
  (when-let* ((target-entry (cl-find room jabber-mam--query-targets
                                     :key #'cdr :test #'string=)))
    (let ((qid (car target-entry)))
      (setq jabber-mam--syncing
            (cl-remove qid jabber-mam--syncing
                       :key #'cdr :test #'string=))
      (setq jabber-mam--query-targets
            (delq target-entry jabber-mam--query-targets))
      (jabber-mam--complete-query qid)
      (condition-case nil
          (jabber-mam--tx-end)
        (error nil))
      (when (zerop jabber-mam--tx-depth)
        (jabber-mam--redraw-dirty)))))

;;; Registration

(jabber-disco-advertise-feature jabber-mam-xmlns)

(jabber-chain-add 'jabber-message-chain #'jabber-mam--process-message -10)

(add-hook 'jabber-post-connect-hooks #'jabber-mam-maybe-catchup)
(add-hook 'jabber-pre-disconnect-hook #'jabber-mam--cleanup-all)
(add-hook 'jabber-lost-connection-hooks #'jabber-mam--cleanup-connection)

(provide 'jabber-mam)
;;; jabber-mam.el ends here
