pi-coding-agent-menu.el

;;; pi-coding-agent-menu.el --- Transient menu and session management -*- lexical-binding: t; -*-

;; Copyright (C) 2026 Daniel Nouri

;; Author: Daniel Nouri <daniel.nouri@gmail.com>
;; Maintainer: Daniel Nouri <daniel.nouri@gmail.com>
;; URL: https://github.com/dnouri/pi-coding-agent

;; SPDX-License-Identifier: GPL-3.0-or-later

;; This file is not part of GNU Emacs.

;; This program 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.
;;
;; This program 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 this program.  If not, see <https://www.gnu.org/licenses/>.

;;; Commentary:

;; Transient menu, session management, model selection, and command
;; infrastructure for pi-coding-agent.
;;
;; Key entry points:
;;   `pi-coding-agent-menu'            Transient menu (C-c C-p)
;;   `pi-coding-agent-new-session'     Start fresh session
;;   `pi-coding-agent-resume-session'  Resume previous session
;;   `pi-coding-agent-reload'          Restart pi process
;;   `pi-coding-agent-select-model'    Choose model interactively
;;   `pi-coding-agent-select-thinking' Choose thinking level interactively
;;   `pi-coding-agent-cycle-thinking'  Cycle thinking levels from header-line
;;   `pi-coding-agent-compact'         Compact conversation context
;;   `pi-coding-agent-fork'            Fork from previous message

;;; Code:

(require 'cl-lib)
(require 'pi-coding-agent-render)
(require 'transient)

(defconst pi-coding-agent--minimum-transient-version "0.9.0"
  "Minimum supported transient version.")

(defun pi-coding-agent--normalize-version (version)
  "Return the numeric prefix of VERSION, or nil when none is present."
  (when (and (stringp version)
             (string-match "[0-9]+\\(?:\\.[0-9]+\\)*" version))
    (match-string 0 version)))

(defun pi-coding-agent--version-at-least-p (version minimum)
  "Return non-nil when VERSION satisfies MINIMUM.
VERSION may include a leading prefix like `v' or extra suffix text."
  (let ((normalized (pi-coding-agent--normalize-version version)))
    (and normalized
         (not (version< normalized minimum)))))

(when (and (not (bound-and-true-p byte-compile-current-file))
           (or (not (boundp 'transient-version))
               (not (pi-coding-agent--version-at-least-p
                     transient-version
                     pi-coding-agent--minimum-transient-version))))
  (display-warning 'pi-coding-agent
                   (format "pi-coding-agent requires transient >= %s, \
but %s is loaded.
  Fix: upgrade transient from MELPA.  If Emacs is using an older built-in
  copy, set `package-install-upgrade-built-in' to t before running
  M-x package-install RET transient RET, then restart Emacs."
                           pi-coding-agent--minimum-transient-version
                           (if (boundp 'transient-version)
                               transient-version
                             "unknown"))
                   :error))

;;;; Slash Commands via RPC

(defun pi-coding-agent--normalize-command (cmd &optional anchor)
  "Normalize a command plist from the RPC wire format.
Lift `sourceInfo.scope' to `:location' and `sourceInfo.path' to
`:path' when present, mapping Pi's temporary scope to the menu's path bucket,
then drop the raw `:sourceInfo' key.  Path values from
Pi are normalized to Emacs paths using ANCHOR or `default-directory'.  Unsafe
passive backend path metadata is ignored rather than stored as navigable state.
Returns CMD (modified in place)."
  (when-let* ((info (plist-get cmd :sourceInfo)))
    (when-let* ((scope (plist-get info :scope)))
      (plist-put cmd :location
                 (if (equal scope "temporary") "path" scope)))
    (when-let* ((path (plist-get info :path)))
      (plist-put cmd :path path))
    (cl-remf cmd :sourceInfo))
  (when-let* ((path (plist-get cmd :path)))
    (if-let* ((emacs-path (pi-coding-agent--passive-emacs-path path anchor)))
        (plist-put cmd :path emacs-path)
      (cl-remf cmd :path)))
  cmd)

(defun pi-coding-agent--fetch-commands (proc callback &optional anchor)
  "Fetch available commands via RPC, call CALLBACK with result.
PROC is the pi process.  CALLBACK receives the command list on success.
ANCHOR is the session directory used to normalize command source paths."
  (pi-coding-agent--rpc-async proc '(:type "get_commands")
    (lambda (response)
      (when (eq (plist-get response :success) t)
        (let* ((data (plist-get response :data))
               (commands-vec (plist-get data :commands))
               (commands (mapcar (lambda (cmd)
                                   (pi-coding-agent--normalize-command
                                    cmd anchor))
                                 (append commands-vec nil))))
          (funcall callback commands))))))

(defun pi-coding-agent--refresh-commands-ignoring-errors
    (proc chat-buf generation anchor)
  "Refresh CHAT-BUF commands from PROC without owning transition completion.
Synchronous command-fetch errors are reported but do not finish GENERATION;
state/history latches remain responsible for unlocking session switches once
those refreshes have been scheduled."
  (condition-case err
      (pi-coding-agent--fetch-commands
       proc
       (lambda (commands)
         (when (pi-coding-agent--session-transition-current-p
                chat-buf proc generation)
           (with-current-buffer chat-buf
             (pi-coding-agent--set-commands commands)
             (pi-coding-agent--rebuild-commands-menu))))
       anchor)
    (error
     (message "Pi: Failed to refresh commands - %s"
              (error-message-string err)))))

;;;; Session Management

(defun pi-coding-agent--menu-state ()
  "Return session state from the chat buffer.
State is buffer-local in the chat buffer; this accessor works
from either chat or input buffer."
  (let ((chat-buf (pi-coding-agent--get-chat-buffer)))
    (and chat-buf (buffer-local-value 'pi-coding-agent--state chat-buf))))

(defun pi-coding-agent--menu-model-description ()
  "Return model description for transient menu."
  (let* ((state (pi-coding-agent--menu-state))
         (model (plist-get (plist-get state :model) :name))
         (short (and model (pi-coding-agent--shorten-model-name model))))
    (format "Model: %s" (or short "unknown"))))

(defun pi-coding-agent--menu-thinking-description ()
  "Return thinking level description for transient menu."
  (let* ((state (pi-coding-agent--menu-state))
         (level (plist-get state :thinking-level)))
    (format "Thinking: %s" (or level "off"))))

(defun pi-coding-agent--menu-description ()
  "Return the transient menu summary line."
  (concat (pi-coding-agent--menu-model-description) " • "
          (pi-coding-agent--menu-thinking-description)))

(defun pi-coding-agent--menu-default-thinking-display-mode ()
  "Return the completed-thinking display mode used for new chat buffers."
  pi-coding-agent-thinking-display)

(defun pi-coding-agent--menu-current-thinking-display-mode ()
  "Return the completed-thinking display mode for the linked chat buffer."
  (let ((chat-buf (pi-coding-agent--get-chat-buffer)))
    (if (and chat-buf (buffer-live-p chat-buf))
        (with-current-buffer chat-buf
          (pi-coding-agent--thinking-display-mode))
      (pi-coding-agent--menu-default-thinking-display-mode))))

(defun pi-coding-agent--next-thinking-display-mode (mode)
  "Return the thinking-display mode after MODE in the visible/hidden cycle."
  (if (eq mode 'hidden) 'visible 'hidden))

(defclass pi-coding-agent--thinking-display-setting (transient-variable)
  ((getter :initarg :getter)
   (setter :initarg :setter))
  "Transient row that shows and changes a thinking-display mode.")

(cl-defmethod transient-init-value ((obj pi-coding-agent--thinking-display-setting))
  "Initialize OBJ from its current thinking-display getter."
  (oset obj value (funcall (oref obj getter))))

(cl-defmethod transient-infix-read ((obj pi-coding-agent--thinking-display-setting))
  "Return the next visible/hidden thinking-display value for OBJ."
  (pi-coding-agent--next-thinking-display-mode (oref obj value)))

(cl-defmethod transient-infix-set ((obj pi-coding-agent--thinking-display-setting) value)
  "Set OBJ to VALUE using its configured thinking-display setter."
  (funcall (oref obj setter) value)
  (oset obj value value))

(cl-defmethod transient-format-value ((obj pi-coding-agent--thinking-display-setting))
  "Format OBJ's current thinking-display value for the transient menu."
  (propertize (symbol-name (oref obj value)) 'face 'transient-value))

;;;###autoload
(defun pi-coding-agent-new-session ()
  "Start a new pi session (reset)."
  (interactive)
  (when-let* ((proc (pi-coding-agent--get-process))
             (chat-buf (pi-coding-agent--get-chat-buffer)))
    (pi-coding-agent--rpc-async proc '(:type "new_session")
                   (lambda (response)
                     (let* ((data (plist-get response :data))
                            (cancelled (plist-get data :cancelled)))
                       (if (and (eq (plist-get response :success) t)
                                (pi-coding-agent--json-false-p cancelled))
                           (when (buffer-live-p chat-buf)
                             (with-current-buffer chat-buf
                               (pi-coding-agent--clear-chat-buffer)
                               (pi-coding-agent--refresh-header))
                             (pi-coding-agent--refresh-session-state proc chat-buf)
                             (message "Pi: New session started"))
                         (message "Pi: New session cancelled")))))))

(defun pi-coding-agent--session-list-directory (&optional chat-buf)
  "Return the directory containing CHAT-BUF's current JSONL session file.
Return nil when the current state has no usable session file.  Relative
session file names are resolved from the chat buffer's stable session
directory."
  (let ((chat-buf (or chat-buf (pi-coding-agent--get-chat-buffer))))
    (when (and chat-buf (buffer-live-p chat-buf))
      (with-current-buffer chat-buf
        (when-let* ((session-file (plist-get pi-coding-agent--state
                                             :session-file))
                    ((stringp session-file))
                    ((not (string-empty-p session-file))))
          (when-let* ((emacs-session-file
                       (pi-coding-agent--emacs-path
                        session-file
                        (pi-coding-agent--chat-session-directory chat-buf))))
            (pi-coding-agent--route-preserving-file-name-directory
             emacs-session-file)))))))

(defun pi-coding-agent--with-session-list-directory (proc chat-buf callback)
  "Call CALLBACK with the session list directory for CHAT-BUF.
When cached state has no session file, fetch fresh state from PROC first."
  (if-let* ((session-dir (pi-coding-agent--session-list-directory chat-buf)))
      (funcall callback session-dir)
    (pi-coding-agent--rpc-async proc '(:type "get_state")
      (lambda (response)
        (when (and (eq (plist-get response :success) t)
                   (buffer-live-p chat-buf))
          (pi-coding-agent--apply-state-response chat-buf response))
        (when (buffer-live-p chat-buf)
          (funcall callback
                   (pi-coding-agent--session-list-directory chat-buf)))))))

(defconst pi-coding-agent--session-line-type-re
  "[ \t]*{[ \t]*\"type\"[ \t]*:[ \t]*\"%s\""
  "Format string matching a JSONL line whose top-level type appears first.")

(defun pi-coding-agent--session-line-type-p (type)
  "Return non-nil when the current line has top-level session TYPE.
Pi writes session JSONL with `type' as the first key.  Matching that cheap
prefix lets resume count ordinary message lines without parsing their full
payloads."
  (looking-at-p (format pi-coding-agent--session-line-type-re
                        (regexp-quote type))))

(defconst pi-coding-agent--session-metadata-fallback-max-line-length 8192
  "Maximum line length parsed when the cheap session type prefix misses.")

(defun pi-coding-agent--session-current-line-data ()
  "Parse the current JSONL line as a plist."
  (json-parse-string (buffer-substring-no-properties
                      (point) (line-end-position))
                     :object-type 'plist))

(defun pi-coding-agent--session-small-current-line-data ()
  "Parse the current line when it is small enough for metadata fallback."
  (let ((end (line-end-position)))
    (when (and (<= (- end (point))
                   pi-coding-agent--session-metadata-fallback-max-line-length)
               (save-excursion
                 (skip-chars-forward " \t" end)
                 (< (point) end)))
      (pi-coding-agent--session-current-line-data))))

(defun pi-coding-agent--session-first-message-text (data)
  "Return first visible message text from parsed session DATA, or nil."
  (let* ((message (plist-get data :message))
         (content (plist-get message :content)))
    (cond
     ((stringp content)
      (unless (string-empty-p content) content))
     ((and (vectorp content) (> (length content) 0))
      (plist-get (aref content 0) :text)))))

(defun pi-coding-agent--session-metadata (path)
  "Extract metadata from session file PATH.
Returns plist with :modified-time, :first-message, :message-count,
:session-name, and :cwd, or nil on error.  Session name comes from the
most recent session_info entry if present."
  (condition-case nil
      (let* ((attrs (file-attributes path))
             (modified-time (file-attribute-modification-time attrs)))
        (with-temp-buffer
          (insert-file-contents path)
          (let ((first-message nil)
                (message-count 0)
                (session-name nil)
                (session-cwd nil)
                (has-session-header nil))
            (cl-labels
                ((apply-entry
                  (data &optional message-counted)
                  (pcase (plist-get data :type)
                    ("message"
                     (unless message-counted
                       (setq message-count (1+ message-count)))
                     (unless first-message
                       (setq first-message
                             (pi-coding-agent--session-first-message-text
                              data))))
                    ("session"
                     (setq has-session-header t
                           session-cwd
                           (pi-coding-agent--normalize-string-or-null
                            (plist-get data :cwd))))
                    ("session_info"
                     (setq session-name
                           (pi-coding-agent--normalize-string-or-null
                            (plist-get data :name)))))))
              (goto-char (point-min))
              ;; Count ordinary message lines by their cheap prefix.  Only parse
              ;; session headers, session_info lines, and the first message whose
              ;; text is needed for the resume picker.
              (while (not (eobp))
                (cond
                 ((pi-coding-agent--session-line-type-p "message")
                  (setq message-count (1+ message-count))
                  (unless first-message
                    (apply-entry (pi-coding-agent--session-current-line-data)
                                 t)))
                 ((or (pi-coding-agent--session-line-type-p "session")
                      (pi-coding-agent--session-line-type-p "session_info"))
                  (apply-entry (pi-coding-agent--session-current-line-data)))
                 (t
                  (when-let* ((data (pi-coding-agent--session-small-current-line-data)))
                    (apply-entry data))))
                (forward-line 1)))
            ;; Only return metadata if we found a valid session header.
            (when has-session-header
              (list :modified-time modified-time
                    :first-message first-message
                    :message-count message-count
                    :session-name session-name
                    :cwd session-cwd)))))
    (error nil)))

(defun pi-coding-agent--session-file-cwd-or-error (path)
  "Return the recorded cwd from session file PATH, or signal `user-error'.
The returned directory is an Emacs path with a trailing slash.  For remote
session files, the recorded process-local cwd is anchored to PATH's TRAMP
prefix.  PATH must be a readable pi session file whose session header contains
a non-empty absolute cwd that names an existing directory."
  (let ((session-file (pi-coding-agent--route-preserving-expand-file-name path)))
    (unless (file-readable-p session-file)
      (user-error "Session file is not readable: %s" session-file))
    (let ((metadata (pi-coding-agent--session-metadata session-file)))
      (unless metadata
        (user-error "Not a pi session file: %s" session-file))
      (let ((cwd (plist-get metadata :cwd)))
        (unless (and (stringp cwd) (not (string-empty-p cwd)))
          (user-error "Session file has no usable cwd: %s" session-file))
        (when (file-remote-p cwd)
          (user-error "Session file cwd must be process-local, not remote: %s\nSession file: %s"
                      cwd session-file))
        (when (pi-coding-agent--remote-home-path-p cwd)
          (user-error "Session file cwd must be absolute, not home-relative: %s\nSession file: %s"
                      cwd session-file))
        (unless (file-name-absolute-p cwd)
          (user-error "Session file cwd is not absolute: %s\nSession file: %s"
                      cwd session-file))
        (let ((expanded-cwd (pi-coding-agent--emacs-directory cwd session-file)))
          (unless (file-directory-p expanded-cwd)
            (user-error "Stored session cwd is not an existing directory: %s\nSession file: %s"
                        expanded-cwd session-file))
          expanded-cwd)))))

(defun pi-coding-agent--update-session-name-from-file (session-file)
  "Update `pi-coding-agent--session-name' from SESSION-FILE metadata.
Call this from the chat buffer after switching or loading a session.
Return the parsed metadata, or nil when SESSION-FILE was not a pi session."
  (when session-file
    (let ((metadata (pi-coding-agent--session-metadata session-file)))
      (setq pi-coding-agent--session-name (plist-get metadata :session-name))
      metadata)))

(defun pi-coding-agent--session-entry (path)
  "Return a resume-list entry for session file PATH, or nil.
The entry carries PATH and its parsed metadata so the resume picker does not
read the same JSONL file again while formatting completion choices."
  (when-let* ((metadata (pi-coding-agent--session-metadata path)))
    (list :path path :metadata metadata)))

(defun pi-coding-agent--list-session-entries (session-dir)
  "List valid session entries in SESSION-DIR.
Entries are sorted by modification time with most recently used first."
  (when (and session-dir (file-directory-p session-dir))
    (let ((entries (delq nil
                         (mapcar #'pi-coding-agent--session-entry
                                 (directory-files session-dir t
                                                  "\\.jsonl\\'")))))
      (sort entries
            (lambda (a b)
              (time-less-p (plist-get (plist-get b :metadata) :modified-time)
                           (plist-get (plist-get a :metadata) :modified-time)))))))

(defun pi-coding-agent--list-sessions (session-dir)
  "List valid session files in SESSION-DIR.
Returns absolute paths to JSONL pi sessions, sorted by modification time with
most recently used first."
  (mapcar (lambda (entry) (plist-get entry :path))
          (pi-coding-agent--list-session-entries session-dir)))

(defun pi-coding-agent--format-session-choice (path &optional metadata)
  "Format session PATH for display in selector.
Returns (display-string . path) for `completing-read', or nil when PATH is not
a valid pi session.  Prefers session name over first message when available.
Optional METADATA reuses data already collected by the session-list step."
  (when-let* ((metadata (or metadata (pi-coding-agent--session-metadata path))))
    (let* ((modified-time (plist-get metadata :modified-time))
           (session-name (plist-get metadata :session-name))
           (first-msg (plist-get metadata :first-message))
           (msg-count (plist-get metadata :message-count))
           (relative-time (pi-coding-agent--format-relative-time modified-time))
           (label (cond
                   (session-name (pi-coding-agent--truncate-string session-name 50))
                   (first-msg (pi-coding-agent--truncate-string first-msg 50))
                   (t nil)))
           (display (if label
                        (format "%s · %s (%d msgs)"
                                label relative-time msg-count)
                      (format "[empty session] · %s" relative-time))))
      (cons display path))))

(defun pi-coding-agent--format-session-entry-choice (entry)
  "Format resume-list ENTRY for display in the session selector."
  (pi-coding-agent--format-session-choice
   (plist-get entry :path)
   (plist-get entry :metadata)))

(defun pi-coding-agent--uniquify-session-choices (choices)
  "Return CHOICES with duplicate display strings made unique.
CHOICES is a list of (DISPLAY . PATH) pairs.  Entries whose DISPLAY occurs
once are returned unchanged.  Entries whose DISPLAY collides get their session
file base name appended."
  (let ((counts (make-hash-table :test 'equal)))
    (dolist (choice choices)
      (puthash (car choice) (1+ (gethash (car choice) counts 0)) counts))
    (mapcar
     (lambda (choice)
       (let ((display (car choice))
             (path (cdr choice)))
         (if (> (gethash display counts 0) 1)
             (cons (format "%s · %s"
                           display
                           (file-name-base (directory-file-name path)))
                   path)
           choice)))
     choices)))

(defun pi-coding-agent--reset-session-state ()
  "Reset all session-specific state for a new session.
Call this when starting a new session to ensure no stale state persists."
  (dolist (marker (list pi-coding-agent--message-start-marker
                        pi-coding-agent--streaming-marker
                        pi-coding-agent--thinking-marker
                        pi-coding-agent--thinking-start-marker))
    (when (markerp marker)
      (set-marker marker nil)))
  (setq pi-coding-agent--session-name nil
        pi-coding-agent--cached-stats nil
        pi-coding-agent--assistant-header-shown nil
        pi-coding-agent--local-user-message nil
        pi-coding-agent--extension-status nil
        pi-coding-agent--working-message nil
        pi-coding-agent--pre-compaction-status nil
        pi-coding-agent--in-code-block nil
        pi-coding-agent--in-thinking-block nil
        pi-coding-agent--thinking-marker nil
        pi-coding-agent--thinking-start-marker nil
        pi-coding-agent--thinking-raw nil
        pi-coding-agent--line-parse-state 'line-start
        pi-coding-agent--pending-tool-overlay nil
        pi-coding-agent--tool-block-order-counter 0
        pi-coding-agent--thinking-block-order-counter 0)
  (pi-coding-agent--set-activity-phase "idle" 'reset t)
  (pi-coding-agent--clear-unsupported-extension-ui-warnings)
  (pi-coding-agent--invalidate-history-loads)
  (pi-coding-agent--finish-session-transition
   pi-coding-agent--session-transition-generation)
  ;; Use accessors for cross-module state
  (pi-coding-agent--cancel-followup-drain-timer)
  (pi-coding-agent--invalidate-prompt-start-wait)
  (pi-coding-agent--clear-followup-queue)
  (pi-coding-agent--set-aborted nil)
  (pi-coding-agent--set-canonical-messages nil)
  (pi-coding-agent--set-message-start-marker nil)
  (pi-coding-agent--set-streaming-marker nil)
  (when pi-coding-agent--tool-args-cache
    (clrhash pi-coding-agent--tool-args-cache))
  (when pi-coding-agent--live-tool-blocks
    (clrhash pi-coding-agent--live-tool-blocks)))

(defun pi-coding-agent--clear-chat-buffer ()
  "Clear the chat buffer and display fresh startup header.
Used when starting a new session."
  (when-let* ((chat-buf (pi-coding-agent--get-chat-buffer)))
    (with-current-buffer chat-buf
      (let ((inhibit-read-only t))
        (pi-coding-agent--clear-render-artifacts)
        (erase-buffer)
        (insert (pi-coding-agent--format-startup-header))
        (insert "\n")
        (pi-coding-agent--reset-session-state)
        (goto-char (point-max))))))

(defun pi-coding-agent--load-session-history
    (proc callback &optional chat-buf completion-callback)
  "Load and display session history from PROC.
Calls CALLBACK with message count when history is applied successfully.
CHAT-BUF is the target buffer; if nil, uses `pi-coding-agent--get-chat-buffer'.
Optional COMPLETION-CALLBACK is called after a current RPC response is handled,
even when the response failed or was not safe to render.  Note: When called
from async callbacks, pass CHAT-BUF explicitly."
  (let ((chat-buf (or chat-buf (pi-coding-agent--get-chat-buffer))))
    (when (and chat-buf (buffer-live-p chat-buf))
      (with-current-buffer chat-buf
        (let ((generation (pi-coding-agent--invalidate-history-loads)))
          (pi-coding-agent--rpc-async proc '(:type "get_messages")
                         (lambda (response)
                           (unwind-protect
                               (when (and (eq (plist-get response :success) t)
                                          (buffer-live-p chat-buf))
                                 (with-current-buffer chat-buf
                                   (when (and (eq pi-coding-agent--process proc)
                                              (= generation
                                                 pi-coding-agent--history-load-generation)
                                              (pi-coding-agent--canonical-rerender-safe-p))
                                     (let* ((messages (plist-get (plist-get response :data)
                                                                 :messages))
                                            (count (if (vectorp messages)
                                                       (length messages)
                                                     0)))
                                       (pi-coding-agent--display-session-history
                                        messages chat-buf)
                                       ;; Refresh header after loading history (resume/fork).
                                       (pi-coding-agent--refresh-header)
                                       (when callback
                                         (funcall callback count))))))
                             (when completion-callback
                               (funcall completion-callback response))))))))))

(defun pi-coding-agent--session-transition-ready-p (chat-buf action)
  "Return non-nil when CHAT-BUF may ACTION another session.
ACTION should be a short verb such as resume or fork for user messages."
  (with-current-buffer chat-buf
    (cond
     ((not (eq pi-coding-agent--status 'idle))
      (message "Pi: Cannot %s while streaming" action)
      nil)
     ((pi-coding-agent--session-busy-p)
      (message "Pi: Cannot %s while Pi is busy" action)
      nil)
     (pi-coding-agent--local-user-message
      (message "Pi: Wait for pi to echo your prompt before you %s" action)
      nil)
     (t t))))

(defun pi-coding-agent--make-session-transition-latch
    (chat-buf proc generation count)
  "Return a callback to finish transition GENERATION after COUNT invocations.
Only completions that still belong to CHAT-BUF, PROC, and GENERATION count, so
stale callbacks from older session switches cannot unlock newer transitions."
  (let ((remaining count)
        (finished nil))
    (lambda (&rest _)
      (when (and (not finished)
                 (pi-coding-agent--session-transition-current-p
                  chat-buf proc generation))
        (setq remaining (1- remaining))
        (when (<= remaining 0)
          (setq finished t)
          (when (buffer-live-p chat-buf)
            (with-current-buffer chat-buf
              (pi-coding-agent--finish-session-transition generation))))))))

(defun pi-coding-agent--refresh-transition-state-and-history
    (proc chat-buf generation &optional session-file history-callback)
  "Refresh state and history for CHAT-BUF through PROC in parallel.
The transition remains active until both the state and history RPC callbacks
for GENERATION have completed or failed.  Synchronous scheduling errors finish
the transition before being re-signaled so the UI cannot stay wedged."
  (let ((latch (pi-coding-agent--make-session-transition-latch
                chat-buf proc generation 2)))
    (condition-case err
        (progn
          (pi-coding-agent--refresh-session-state
           proc chat-buf session-file generation latch)
          (pi-coding-agent--load-session-history
           proc history-callback chat-buf latch))
      (error
       (when (pi-coding-agent--session-transition-current-p
              chat-buf proc generation)
         (with-current-buffer chat-buf
           (pi-coding-agent--finish-session-transition generation)))
       (signal (car err) (cdr err))))))

(defun pi-coding-agent--refresh-session-state
    (proc chat-buf &optional session-file generation completion-callback)
  "Refresh session state for CHAT-BUF from PROC.
SESSION-FILE seeds the session-name cache when the switching action already
knows the selected file.  Optional GENERATION ties the refresh to an existing
session transition; otherwise this function starts and finishes its own guard.
Optional COMPLETION-CALLBACK runs after a current get_state response is
handled, even when the response failed."
  (when (buffer-live-p chat-buf)
    (with-current-buffer chat-buf
      (pi-coding-agent--set-canonical-messages nil)
      (when session-file
        (pi-coding-agent--update-session-name-from-file session-file))
      (let* ((own-generation (null generation))
             (generation (or generation
                             (pi-coding-agent--begin-session-transition))))
        (pi-coding-agent--rpc-async proc '(:type "get_state")
          (lambda (response)
            (when (pi-coding-agent--session-transition-current-p
                   chat-buf proc generation)
              (unwind-protect
                  (when (eq (plist-get response :success) t)
                    (pi-coding-agent--apply-state-response chat-buf response)
                    (when (buffer-live-p chat-buf)
                      (with-current-buffer chat-buf
                        (unless session-file
                          (when-let* ((current-session-file
                                       (plist-get pi-coding-agent--state
                                                  :session-file)))
                            (pi-coding-agent--update-session-name-from-file
                             current-session-file)))
                        (force-mode-line-update t))))
                (when completion-callback
                  (funcall completion-callback response))
                (when (and own-generation (buffer-live-p chat-buf))
                  (with-current-buffer chat-buf
                    (pi-coding-agent--finish-session-transition
                     generation)))))))))))

;;;###autoload
(defun pi-coding-agent-reload ()
  "Reload the current session by restarting the pi process.
Useful for reloading extensions, skills, prompts, and themes after
editing them, or when the pi process has died or become unresponsive.
Starts a fresh process, switches it back to the cached session file, then
replaces the old process, refreshes state and commands, and rebuilds the chat
buffer from session history."
  (interactive)
  (let* ((chat-buf (pi-coding-agent--get-chat-buffer))
         (session-file (and chat-buf
                            (buffer-local-value 'pi-coding-agent--state chat-buf)
                            (plist-get (buffer-local-value 'pi-coding-agent--state chat-buf)
                                       :session-file))))
    (cond
     ((not chat-buf)
      (message "Pi: No session to reload"))
     ((not session-file)
      (message "Pi: No session file available - cannot reload"))
     (t
      (message "Pi: Reloading...")
      (with-current-buffer chat-buf
        (let ((dir (pi-coding-agent--session-directory)))
          (unless (and (stringp dir)
                       (not (string-empty-p dir))
                       (file-name-absolute-p dir))
            (user-error "Pi: Cannot reload from invalid session directory: %s"
                        dir))
          (let ((session-path (pi-coding-agent--process-local-path
                               session-file dir))
                (old-proc pi-coding-agent--process))
            (setq pi-coding-agent--status 'idle)
            (let* ((new-proc (pi-coding-agent--start-process dir))
                   (generation (pi-coding-agent--begin-session-transition
                                new-proc)))
              (when (processp new-proc)
                (set-process-buffer new-proc chat-buf)
                (process-put new-proc 'pi-coding-agent-chat-buffer chat-buf)
                (pi-coding-agent--register-display-handler new-proc)
                (pi-coding-agent--rpc-async
                 new-proc
                 (list :type "switch_session" :sessionPath session-path)
                 (lambda (response)
                   (when (pi-coding-agent--session-transition-current-p
                          chat-buf new-proc generation)
                     (let* ((data (plist-get response :data))
                            (cancelled (plist-get data :cancelled)))
                       (if (and (eq (plist-get response :success) t)
                                (pi-coding-agent--json-false-p cancelled))
                           (let ((refresh-scheduled nil))
                             (condition-case err
                                 (progn
                                   (when (and (processp old-proc)
                                              (not (eq old-proc new-proc)))
                                     (pi-coding-agent--skip-process-kill-confirmation
                                      old-proc)
                                     (pi-coding-agent--unregister-display-handler
                                      old-proc)
                                     (when (process-live-p old-proc)
                                       (delete-process old-proc)))
                                   (when (buffer-live-p chat-buf)
                                     (with-current-buffer chat-buf
                                       (pi-coding-agent--set-process new-proc)))
                                   (pi-coding-agent--refresh-transition-state-and-history
                                    new-proc chat-buf generation session-file
                                    (lambda (_count)
                                      (message "Pi: Session reloaded")))
                                   (setq refresh-scheduled t))
                               (error
                                (when (pi-coding-agent--session-transition-current-p
                                       chat-buf new-proc generation)
                                  (with-current-buffer chat-buf
                                    (pi-coding-agent--finish-session-transition
                                     generation)))
                                (message "Pi: Failed to reload - %s"
                                         (error-message-string err))))
                             (when refresh-scheduled
                               (pi-coding-agent--refresh-commands-ignoring-errors
                                new-proc chat-buf generation dir)))
                         (pi-coding-agent--unregister-display-handler new-proc)
                         (when (process-live-p new-proc)
                           (delete-process new-proc))
                         (if (and cancelled
                                  (not (pi-coding-agent--json-false-p cancelled)))
                             (message "Pi: Reload cancelled")
                           (message "Pi: Failed to reload - %s"
                                    (or (plist-get response :error)
                                        "unknown error")))
                         (when (buffer-live-p chat-buf)
                           (with-current-buffer chat-buf
                             (pi-coding-agent--finish-session-transition
                              generation)))))))))))))))))

(defun pi-coding-agent--resume-selected-session (proc chat-buf selected-path)
  "Resume SELECTED-PATH using PROC and rebuild CHAT-BUF from its history."
  (let* ((target-dir (pi-coding-agent--session-file-cwd-or-error selected-path))
         (session (and (buffer-live-p chat-buf)
                       (pi-coding-agent--chat-session-name chat-buf)))
         (existing-target (pi-coding-agent--find-session target-dir session)))
    (when (and existing-target (not (eq existing-target chat-buf)))
      (user-error "Pi session already open for: %s" target-dir))
    (let* ((session-path (if (buffer-live-p chat-buf)
                             (with-current-buffer chat-buf
                               (pi-coding-agent--process-local-path
                                selected-path
                                (pi-coding-agent--chat-session-directory chat-buf)))
                           (pi-coding-agent--process-local-path selected-path)))
           (generation (when (buffer-live-p chat-buf)
                         (with-current-buffer chat-buf
                           (pi-coding-agent--begin-session-transition proc)))))
      (pi-coding-agent--rpc-async
       proc
       (list :type "switch_session" :sessionPath session-path)
       (lambda (response)
         (when (pi-coding-agent--session-transition-current-p
                chat-buf proc generation)
           (let* ((data (plist-get response :data))
                  (cancelled (plist-get data :cancelled)))
             (if (and (eq (plist-get response :success) t)
                      (pi-coding-agent--json-false-p cancelled))
                 (let ((refresh-scheduled nil))
                   (condition-case err
                       (progn
                         (when (buffer-live-p chat-buf)
                           (with-current-buffer chat-buf
                             (pi-coding-agent--retarget-session-buffers
                              target-dir)))
                         (pi-coding-agent--refresh-transition-state-and-history
                          proc chat-buf generation selected-path
                          (lambda (count)
                            (message "Pi: Resumed session (%d messages)" count)))
                         (setq refresh-scheduled t))
                     (error
                      (when (pi-coding-agent--session-transition-current-p
                             chat-buf proc generation)
                        (with-current-buffer chat-buf
                          (pi-coding-agent--finish-session-transition generation)))
                      (message "Pi: Failed to resume session - %s"
                               (error-message-string err))))
                   (when refresh-scheduled
                     (pi-coding-agent--refresh-commands-ignoring-errors
                      proc chat-buf generation target-dir)))
               (message "Pi: Failed to resume session")
               (when (buffer-live-p chat-buf)
                 (with-current-buffer chat-buf
                   (pi-coding-agent--finish-session-transition generation)))))))))))

(defun pi-coding-agent--resume-session-from-directory (proc chat-buf session-dir)
  "Prompt for a session from SESSION-DIR, then resume it using PROC.
CHAT-BUF is rebuilt from the selected session history."
  (let ((entries (pi-coding-agent--list-session-entries session-dir)))
    (if (null entries)
        (message "Pi: No previous sessions found")
      (let* ((choices
              (pi-coding-agent--uniquify-session-choices
               (delq nil
                     (mapcar #'pi-coding-agent--format-session-entry-choice
                             entries))))
             (choice-strings (mapcar #'car choices)))
        (if (null choices)
            (message "Pi: No previous sessions found")
          (let* ((choice (completing-read
                          "Resume session: "
                          (lambda (string pred action)
                            (if (eq action 'metadata)
                                '(metadata (display-sort-function . identity))
                              (complete-with-action action choice-strings
                                                    string pred)))
                          nil t))
                 (selected-path (cdr (assoc choice choices))))
            (when selected-path
              (pi-coding-agent--resume-selected-session
               proc chat-buf selected-path))))))))

(defun pi-coding-agent-resume-session ()
  "Resume a previous pi session stored beside the current session file."
  (interactive)
  (when-let* ((proc (pi-coding-agent--get-process))
              (chat-buf (pi-coding-agent--get-chat-buffer)))
    (when (pi-coding-agent--session-transition-ready-p chat-buf "resume")
      (pi-coding-agent--with-session-list-directory
       proc chat-buf
       (lambda (session-dir)
         (cond
          ((not session-dir)
           (message "Pi: Session file not available"))
          ((pi-coding-agent--session-transition-ready-p chat-buf "resume")
           (pi-coding-agent--resume-session-from-directory
            proc chat-buf session-dir))))))))

;;;; Model and Thinking

(defun pi-coding-agent-set-session-name (name)
  "Set the session NAME for the current session.
The name is displayed in the resume picker and header-line."
  (interactive
   (let ((chat-buf (pi-coding-agent--get-chat-buffer)))
     (list (read-string "Session name: "
                        (or (and chat-buf
                                 (buffer-local-value 'pi-coding-agent--session-name chat-buf))
                            "")))))
  (let* ((trimmed-name (string-trim name))
         (chat-buf (pi-coding-agent--get-chat-buffer)))
    (if (string-empty-p trimmed-name)
        ;; Consistent with TUI /name behavior
        (let ((current-name (and chat-buf
                                 (buffer-local-value 'pi-coding-agent--session-name chat-buf))))
          (if current-name
              (message "Pi: Session name: %s" current-name)
            (message "Pi: No session name set")))
      (let ((proc (pi-coding-agent--get-process)))
        (unless proc
          (user-error "No pi process running"))
        (pi-coding-agent--rpc-async proc
            (list :type "set_session_name" :name trimmed-name)
            (lambda (response)
              (if (eq (plist-get response :success) t)
                  (progn
                    (when (buffer-live-p chat-buf)
                      (with-current-buffer chat-buf
                        (setq pi-coding-agent--session-name trimmed-name)
                        (force-mode-line-update t)))
                    (message "Pi: Session name set to \"%s\"" trimmed-name))
                (message "Pi: Failed to set session name: %s"
                         (or (plist-get response :error) "unknown error")))))))))

(defun pi-coding-agent-select-model (&optional initial-input)
  "Select a model interactively.
Optional INITIAL-INPUT pre-fills the completion prompt for filtering."
  (interactive)
  (let ((proc (pi-coding-agent--get-process))
        (chat-buf (pi-coding-agent--get-chat-buffer)))
    (unless proc
      (user-error "No pi process running"))
    (let* ((state (pi-coding-agent--menu-state))
           (response (pi-coding-agent--rpc-sync proc '(:type "get_available_models") 5))
           (data (plist-get response :data))
           (models (plist-get data :models))
           (current-name (plist-get (plist-get state :model) :name))
           (current-provider (plist-get (plist-get state :model) :provider))
           (current-short (and current-name
                               (pi-coding-agent--shorten-model-name current-name)))
           (current-display (and current-short current-provider
                                (format "%s [%s]" current-short current-provider)))
           ;; Build alist of (display-string . model-plist) for selection
           ;; Display includes provider for clarity
           (model-alist (mapcar (lambda (m)
                                  (let ((short (pi-coding-agent--shorten-model-name
                                               (plist-get m :name)))
                                        (prov (plist-get m :provider)))
                                    (cons (format "%s [%s]" short (or prov "?"))
                                          m)))
                                models))
           (names (mapcar #'car model-alist))
           (choice (let ((completion-ignore-case t)
                         (completion-styles '(basic flex)))
                     (if initial-input
                         ;; Try auto-selecting on unique match
                         (let ((matches (completion-all-completions
                                         initial-input names nil
                                         (length initial-input))))
                           (when (consp matches)
                             (setcdr (last matches) nil))
                           (cond
                            ((= (length matches) 1) (car matches))
                            ((null matches)
                             (message "Pi: No model matching \"%s\"" initial-input)
                             nil)
                            (t (completing-read
                                (format "Model (current: %s): "
                                        (or current-display "unknown"))
                                names nil t initial-input))))
                       (completing-read
                        (format "Model (current: %s): "
                                (or current-display "unknown"))
                        names nil t)))))
      (when (and choice (not (equal choice current-display)))
        (let* ((selected-model (cdr (assoc choice model-alist)))
               (model-id (plist-get selected-model :id))
               (provider (plist-get selected-model :provider)))
          (pi-coding-agent--rpc-async proc (list :type "set_model"
                                    :provider provider
                                    :modelId model-id)
                         (lambda (resp)
                           (when (and (eq (plist-get resp :success) t)
                                      (buffer-live-p chat-buf))
                             (with-current-buffer chat-buf
                               (pi-coding-agent--update-state-from-response resp)
                               (force-mode-line-update))
                             (message "Pi: Model set to %s" choice)))))))))

(defconst pi-coding-agent--thinking-levels '("off" "minimal" "low" "medium" "high" "xhigh")
  "Thinking levels accepted by `set_thinking_level' RPC.

Unsupported levels are clamped to the current model's capabilities.")

(defun pi-coding-agent-cycle-thinking ()
  "Cycle through thinking levels."
  (interactive)
  (when-let* ((proc (pi-coding-agent--get-process))
             (chat-buf (pi-coding-agent--get-chat-buffer)))
    (pi-coding-agent--rpc-async proc '(:type "cycle_thinking_level")
                   (lambda (response)
                     (when (and (eq (plist-get response :success) t)
                                (buffer-live-p chat-buf))
                       (with-current-buffer chat-buf
                         (pi-coding-agent--update-state-from-response response)
                         (force-mode-line-update)
                         (message "Pi: Thinking level: %s"
                                  (plist-get pi-coding-agent--state :thinking-level))))))))

(defun pi-coding-agent--refresh-thinking-level-state (proc chat-buf)
  "Refresh CHAT-BUF state from PROC after a thinking-level change.
Uses `get_state' so the UI reflects the server's actual level,
including any model-specific clamping."
  (pi-coding-agent--rpc-async
   proc '(:type "get_state")
   (lambda (response)
     (if (eq (plist-get response :success) t)
         (let* ((data (plist-get response :data))
                (level (or (plist-get data :thinkingLevel) "off")))
           (pi-coding-agent--apply-state-response chat-buf response)
           (message "Pi: Thinking level: %s" level))
       (message "Pi: Thinking level updated, but failed to refresh state%s"
                (if-let* ((error-text (plist-get response :error)))
                    (format ": %s" error-text)
                  ""))))))

(defun pi-coding-agent-select-thinking ()
  "Select a thinking level from the minibuffer."
  (interactive)
  (let ((proc (pi-coding-agent--get-process))
        (chat-buf (pi-coding-agent--get-chat-buffer)))
    (unless proc
      (user-error "No pi process running"))
    (unless chat-buf
      (user-error "No pi session buffer"))
    (let* ((state (pi-coding-agent--menu-state))
           (current (or (plist-get state :thinking-level) "off"))
           (choice (completing-read
                    (format "Thinking level (current: %s): " current)
                    pi-coding-agent--thinking-levels
                    nil t)))
      (unless (equal choice current)
        (pi-coding-agent--rpc-async
         proc (list :type "set_thinking_level" :level choice)
         (lambda (response)
           (if (eq (plist-get response :success) t)
               (pi-coding-agent--refresh-thinking-level-state proc chat-buf)
             (message "Pi: Failed to set thinking level: %s"
                      (or (plist-get response :error) "unknown error")))))))))

(defun pi-coding-agent-toggle-thinking-display ()
  "Toggle completed-thinking display for the current chat buffer."
  (interactive)
  (pi-coding-agent--set-chat-thinking-display
   (pi-coding-agent--next-thinking-display-mode
    (pi-coding-agent--menu-current-thinking-display-mode))))

(defun pi-coding-agent--set-default-thinking-display (mode)
  "Set MODE as the default completed-thinking display for new chat buffers."
  (setq pi-coding-agent-thinking-display mode)
  (message "Pi: New chat buffers will %s completed thinking by default"
           (if (eq mode 'hidden) "hide" "show")))

(defun pi-coding-agent-toggle-default-thinking-display ()
  "Toggle the completed-thinking display default for new chat buffers.
This changes the live default for future chat buffers in the current Emacs
session.  Persist it with Customize or your init file if you want it to stick
across restarts."
  (interactive)
  (pi-coding-agent--set-default-thinking-display
   (pi-coding-agent--next-thinking-display-mode
    (pi-coding-agent--menu-default-thinking-display-mode))))

(transient-define-infix pi-coding-agent-menu-chat-thinking-display ()
  :class 'pi-coding-agent--thinking-display-setting
  :key "h"
  :description "This chat"
  :getter #'pi-coding-agent--menu-current-thinking-display-mode
  :setter #'pi-coding-agent--set-chat-thinking-display)

(transient-define-infix pi-coding-agent-menu-default-thinking-display ()
  :class 'pi-coding-agent--thinking-display-setting
  :key "H"
  :description "New chat default"
  :getter #'pi-coding-agent--menu-default-thinking-display-mode
  :setter #'pi-coding-agent--set-default-thinking-display)

;;;; Session Info and Actions

(defun pi-coding-agent--format-session-stats (stats)
  "Format STATS plist as human-readable string."
  (let* ((tokens (plist-get stats :tokens))
         (input (or (plist-get tokens :input) 0))
         (output (or (plist-get tokens :output) 0))
         (total (or (plist-get tokens :total) 0))
         (cache-read (or (plist-get tokens :cacheRead) 0))
         (cache-write (or (plist-get tokens :cacheWrite) 0))
         (cost (or (plist-get stats :cost) 0))
         (messages (or (plist-get stats :userMessages) 0))
         (tools (or (plist-get stats :toolCalls) 0)))
    (format "Tokens: %s in / %s out (%s total) | Cache: R%s / W%s | Cost: $%.2f | Messages: %d | Tools: %d"
            (pi-coding-agent--format-number input)
            (pi-coding-agent--format-number output)
            (pi-coding-agent--format-number total)
            (pi-coding-agent--format-number cache-read)
            (pi-coding-agent--format-number cache-write)
            cost messages tools)))

(defun pi-coding-agent-session-stats ()
  "Display session statistics in the echo area."
  (interactive)
  (when-let* ((proc (pi-coding-agent--get-process)))
    (pi-coding-agent--rpc-async proc '(:type "get_session_stats")
                   (lambda (response)
                     (if (eq (plist-get response :success) t)
                         (let ((data (plist-get response :data)))
                           (message "Pi: %s" (pi-coding-agent--format-session-stats data)))
                       (message "Pi: Failed to get session stats"))))))

(defun pi-coding-agent-process-info ()
  "Display process information for debugging.
Shows PID, status, and session file."
  (interactive)
  (let* ((chat-buf (pi-coding-agent--get-chat-buffer))
         (proc (and chat-buf (buffer-local-value 'pi-coding-agent--process chat-buf)))
         (state (and chat-buf (buffer-local-value 'pi-coding-agent--state chat-buf)))
         (status (and chat-buf (buffer-local-value 'pi-coding-agent--status chat-buf)))
         (session-file (and state (plist-get state :session-file))))
    (cond
     ((not chat-buf)
      (message "Pi: No session"))
     ((not proc)
      (message "Pi: No process (status: %s, session: %s)"
               status
               (or session-file "none")))
     (t
      (message "Pi: PID %s, %s (status: %s, session: %s)"
               (process-id proc)
               (if (process-live-p proc) "alive" "dead")
               status
               (or (and session-file (file-name-nondirectory session-file)) "none"))))))

(defun pi-coding-agent--handle-manual-compaction-response (chat-buf response)
  "Handle manual compact command RESPONSE for CHAT-BUF.
Canonical compaction events render success, failure, and queue effects.
This callback reports only command-level failure seen before any end event."
  (when (buffer-live-p chat-buf)
    (with-current-buffer chat-buf
      (unless (eq (plist-get response :success) t)
        (when (eq pi-coding-agent--status 'compacting)
          (setq pi-coding-agent--status 'idle)
          (pi-coding-agent--set-activity-phase "idle")
          (pi-coding-agent--restore-followup-queue-to-input)
          (message "Pi: Compact failed%s"
                   (if-let* ((error-text (plist-get response :error)))
                       (format ": %s" error-text)
                     "")))))))

(defun pi-coding-agent-compact (&optional custom-instructions)
  "Compact conversation context to reduce token usage.
Optional CUSTOM-INSTRUCTIONS provide guidance for the compaction summary."
  (interactive)
  (when-let* ((chat-buf (pi-coding-agent--get-chat-buffer)))
    (let ((proc (pi-coding-agent--get-process)))
      (cond
       ((null proc)
        (message "Pi: No process available - try M-x pi-coding-agent-reload or C-c C-p R"))
       ((not (process-live-p proc))
        (message "Pi: Process died - try M-x pi-coding-agent-reload or C-c C-p R"))
       (t
        (message "Pi: Compacting...")
        (with-current-buffer chat-buf
          (setq pi-coding-agent--status 'compacting)
          (pi-coding-agent--set-activity-phase "compact"))
        (pi-coding-agent--rpc-async
         proc
         (if custom-instructions
             (list :type "compact" :customInstructions custom-instructions)
           '(:type "compact"))
         (lambda (response)
           (pi-coding-agent--handle-manual-compaction-response chat-buf response))))))))

(defun pi-coding-agent-export-html (&optional output-path)
  "Export session to HTML file.
Optional OUTPUT-PATH specifies where to save; nil uses pi's default."
  (interactive
   (list (let ((path (read-string
                      "Export path on session host (RET for default): ")))
           (and (not (string-empty-p path)) path))))
  (when-let* ((proc (pi-coding-agent--get-process)))
    (let* ((chat-buf (pi-coding-agent--get-chat-buffer))
           (anchor (when (buffer-live-p chat-buf)
                     (with-current-buffer chat-buf
                       (pi-coding-agent--chat-session-directory chat-buf))))
           (process-output-path
            (pi-coding-agent--process-local-path output-path anchor)))
      (pi-coding-agent--rpc-async
       proc
       (if process-output-path
           (list :type "export_html" :outputPath process-output-path)
         '(:type "export_html"))
       (lambda (response)
         (if (eq (plist-get response :success) t)
             (let* ((data (plist-get response :data))
                    (path (plist-get data :path))
                    (emacs-path (pi-coding-agent--passive-emacs-path
                                 path anchor)))
               (if emacs-path
                   (message "Pi: Exported to %s" emacs-path)
                 (message "Pi: Exported, but Pi did not return a usable path")))
           (message "Pi: Export failed")))))))

(defun pi-coding-agent-copy-last-message ()
  "Copy last assistant message to kill ring."
  (interactive)
  (when-let* ((proc (pi-coding-agent--get-process)))
    (pi-coding-agent--rpc-async proc '(:type "get_last_assistant_text")
                   (lambda (response)
                     (if (eq (plist-get response :success) t)
                         (let* ((data (plist-get response :data))
                                (text (plist-get data :text)))
                           (if text
                               (progn
                                 (kill-new text)
                                 (message "Pi: Copied to kill ring"))
                             (message "Pi: No assistant message to copy")))
                       (message "Pi: Failed to get message"))))))

;;;; Fork

(defun pi-coding-agent--flatten-tree (nodes)
  "Flatten tree NODES into a hash table mapping id to node plist.
NODES is a vector of tree node plists, each with `:children' vector.
Returns a hash table for O(1) lookup by id.

Uses iterative traversal to avoid `max-lisp-eval-depth' errors on deep
session trees."
  (let ((index (make-hash-table :test 'equal))
        (stack nil))
    ;; Push roots in reverse so popping preserves original order.
    (let ((i (1- (length nodes))))
      (while (>= i 0)
        (push (aref nodes i) stack)
        (setq i (1- i))))
    (while stack
      (let* ((node (pop stack))
             (children (plist-get node :children)))
        (puthash (plist-get node :id) node index)
        (let ((i (1- (length children))))
          (while (>= i 0)
            (push (aref children i) stack)
            (setq i (1- i))))))
    index))

(defun pi-coding-agent--active-branch-user-ids (index leaf-id)
  "Return chronological list of user message IDs on the active branch.
INDEX is a hash table from `pi-coding-agent--flatten-tree'.
LEAF-ID is the current leaf node ID.  Walk from leaf to root via
`:parentId', collecting IDs of nodes with type \"message\" and role
\"user\".  Returns list in root-to-leaf (chronological) order."
  (when leaf-id
    (let ((user-ids nil)
          (current-id leaf-id))
      (while current-id
        (let ((node (gethash current-id index)))
          (when (and node
                     (equal (plist-get node :type) "message")
                     (equal (plist-get node :role) "user"))
            (push (plist-get node :id) user-ids))
          (setq current-id (and node (plist-get node :parentId)))))
      user-ids)))

(defun pi-coding-agent--format-fork-message (msg &optional index)
  "Format MSG for display in fork selector.
MSG is a plist with :entryId and :text.
INDEX is the display index (1-based) for the message."
  (let* ((text (or (plist-get msg :text) ""))
         (preview (truncate-string-to-width text 60 nil nil "...")))
    (if index
        (format "%d: %s" index preview)
      preview)))

(defun pi-coding-agent-fork ()
  "Fork conversation from a previous user message.
Shows a selector of user messages and creates a fork from the selected one."
  (interactive)
  (when-let* ((proc (pi-coding-agent--get-process))
              (chat-buf (pi-coding-agent--get-chat-buffer)))
    (when (pi-coding-agent--session-transition-ready-p chat-buf "fork")
      (pi-coding-agent--rpc-async proc '(:type "get_fork_messages")
                     (lambda (response)
                       (if (eq (plist-get response :success) t)
                           (let* ((data (plist-get response :data))
                                  (messages (plist-get data :messages)))
                             ;; Note: messages is a vector from JSON, use seq-empty-p not null
                             (if (seq-empty-p messages)
                                 (message "Pi: No messages to fork from")
                               (pi-coding-agent--show-fork-selector proc messages)))
                         (message "Pi: Failed to get fork messages")))))))

(defun pi-coding-agent--resolve-fork-entry (response ordinal heading-count)
  "Resolve a fork entry ID from get_fork_messages RESPONSE.
ORDINAL is the 0-based user turn index.  HEADING-COUNT is the number
of visible You headings in the buffer.  Returns (ENTRY-ID . PREVIEW)
or nil if the ordinal could not be mapped."
  (when (eq (plist-get response :success) t)
    (let* ((data (plist-get response :data))
           (messages (append (plist-get data :messages) nil))
           ;; Use last N messages to align with visible headings in
           ;; compacted sessions.
           (visible-messages (last messages heading-count))
           (selected (nth ordinal visible-messages))
           (entry-id (plist-get selected :entryId)))
      (when entry-id
        (cons entry-id (pi-coding-agent--format-fork-message selected))))))

(defun pi-coding-agent-fork-at-point ()
  "Fork conversation from the user turn at point.
Determines which user message point is in (or after), confirms with
a preview, then forks.  Only works when the session is idle."
  (interactive)
  (let ((chat-buf (pi-coding-agent--get-chat-buffer)))
    (unless chat-buf
      (user-error "Pi: No chat buffer"))
    (with-current-buffer chat-buf
      (let* ((headings (pi-coding-agent--collect-you-headings))
             (ordinal (pi-coding-agent--user-turn-index-at-point headings)))
        (cond
         ((not (pi-coding-agent--session-transition-ready-p chat-buf "fork")))
         ((not ordinal)
          (message "Pi: No user message at point"))
         (t
          (let ((heading-count (length headings))
                (proc (pi-coding-agent--get-process)))
            (unless proc
              (user-error "Pi: No active process"))
            (pi-coding-agent--rpc-async proc '(:type "get_fork_messages")
              (lambda (response)
                (if (not (eq (plist-get response :success) t))
                    (if-let* ((error-text (plist-get response :error)))
                        (message "Pi: Failed to get fork messages: %s" error-text)
                      (message "Pi: Failed to get fork messages"))
                  (let ((result (pi-coding-agent--resolve-fork-entry
                                 response ordinal heading-count)))
                    (cond
                     ((not result)
                      (message "Pi: Could not map turn to entry ID"))
                     ((with-current-buffer chat-buf
                        (y-or-n-p (format "Fork from: %s? " (or (cdr result) "?"))))
                      (with-current-buffer chat-buf
                        (pi-coding-agent--execute-fork proc (car result))))))))))))))))

(defun pi-coding-agent--execute-fork (proc entry-id)
  "Execute fork to ENTRY-ID via PROC.
Sends the fork RPC, then on success: refreshes state, reloads history,
and pre-fills the input buffer with the forked message text.
Captures chat and input buffers at call time (before the async RPC)."
  (let* ((chat-buf (pi-coding-agent--get-chat-buffer))
         (input-buf (pi-coding-agent--get-input-buffer))
         (generation (when (buffer-live-p chat-buf)
                       (with-current-buffer chat-buf
                         (pi-coding-agent--begin-session-transition proc)))))
    (pi-coding-agent--rpc-async proc (list :type "fork" :entryId entry-id)
      (lambda (response)
        (when (pi-coding-agent--session-transition-current-p
               chat-buf proc generation)
          (if (eq (plist-get response :success) t)
              (let ((refresh-scheduled nil)
                    text)
                (condition-case err
                    (progn
                      (setq text (plist-get (plist-get response :data) :text))
                      (pi-coding-agent--refresh-transition-state-and-history
                       proc chat-buf generation nil
                       (lambda (count)
                         (message "Pi: Branched to new session (%d messages)" count)))
                      (setq refresh-scheduled t))
                  (error
                   (when (pi-coding-agent--session-transition-current-p
                          chat-buf proc generation)
                     (with-current-buffer chat-buf
                       (pi-coding-agent--finish-session-transition generation)))
                   (message "Pi: Branch failed - %s"
                            (error-message-string err))))
                ;; Pre-fill input with the forked message text.  Sending stays
                ;; blocked by the transition until state and history settle.
                (when refresh-scheduled
                  (condition-case err
                      (when (buffer-live-p input-buf)
                        (with-current-buffer input-buf
                          (erase-buffer)
                          (when text (insert text))))
                    (error
                     (message "Pi: Failed to prefill fork prompt - %s"
                              (error-message-string err))))))
            (message "Pi: Branch failed")
            (when (buffer-live-p chat-buf)
              (with-current-buffer chat-buf
                (pi-coding-agent--finish-session-transition generation)))))))))

(defun pi-coding-agent--show-fork-selector (proc messages)
  "Show selector for MESSAGES and fork on selection.
PROC is the pi process.
MESSAGES is a vector of plists from get_fork_messages."
  (let* ((index 0)
         ;; Reverse so most recent messages appear first (upstream sends chronological order)
         (reversed-messages (reverse (append messages nil)))
         (formatted (mapcar (lambda (msg)
                              (setq index (1+ index))
                              (cons (pi-coding-agent--format-fork-message msg index) msg))
                            reversed-messages))
         (choice-strings (mapcar #'car formatted))
         ;; Use completion table with metadata to preserve our sort order
         ;; (completing-read normally re-sorts alphabetically)
         (choice (completing-read "Branch from: "
                                  (lambda (string pred action)
                                    (if (eq action 'metadata)
                                        '(metadata (display-sort-function . identity))
                                      (complete-with-action action choice-strings string pred)))
                                  nil t))
         (selected (cdr (assoc choice formatted))))
    (when selected
      (pi-coding-agent--execute-fork proc (plist-get selected :entryId)))))

;;;; Custom Commands

(defun pi-coding-agent--command-chat-buffer-or-error ()
  "Return the current pi chat buffer for running a slash command.
Signal `user-error' when no live chat buffer is linked to the current buffer."
  (let ((chat-buf (pi-coding-agent--get-chat-buffer)))
    (unless (and chat-buf (buffer-live-p chat-buf))
      (user-error "No pi session in current buffer"))
    chat-buf))

(defun pi-coding-agent-run-command (name &optional args)
  "Run pi slash command NAME with optional ARGS in the current session.
NAME is the command name without the leading slash.  ARGS, when
non-nil and non-empty, is appended after one space.

This command sends through the pi session associated with the current
pi chat buffer or its linked input buffer.  Signal `user-error' when the
current buffer is not part of a pi session."
  (interactive
   (progn
     (pi-coding-agent--command-chat-buffer-or-error)
     (list (completing-read "Pi command: "
                            (mapcar (lambda (cmd) (plist-get cmd :name))
                                    pi-coding-agent--commands)
                            nil t)
           (read-string "Args: "))))
  (let ((chat-buf (pi-coding-agent--command-chat-buffer-or-error)))
    (let ((full-command (if (or (null args) (string-empty-p args))
                            (format "/%s" name)
                          (format "/%s %s" name args))))
      (with-current-buffer chat-buf
        (pi-coding-agent--prepare-and-send full-command)))))

(defun pi-coding-agent--run-custom-command (cmd)
  "Execute custom command CMD.
Always prompts for arguments - user can press Enter if none needed.
Sends the literal /command text to pi, which handles expansion."
  (let* ((name (plist-get cmd :name))
         (args-string (read-string (format "/%s: " name))))
    (pi-coding-agent-run-command name args-string)))

(defun pi-coding-agent-run-custom-command ()
  "Select and run a custom command.
Uses commands from pi's `get_commands' RPC."
  (interactive)
  (if (null pi-coding-agent--commands)
      (message "Pi: No commands available")
    (let* ((choices (mapcar (lambda (cmd)
                              (cons (format "%s - %s"
                                            (plist-get cmd :name)
                                            (or (plist-get cmd :description) ""))
                                    cmd))
                            pi-coding-agent--commands))
           (choice (completing-read "Command: " choices nil t))
           (cmd (cdr (assoc choice choices))))
      (when cmd
        (pi-coding-agent--run-custom-command cmd)))))

;;;; Transient Menu

(transient-define-prefix pi-coding-agent-menu ()
  "Pi coding agent menu."
  [:description #'pi-coding-agent--menu-description
   :class transient-row]
  [["Session"
    ("n" "new" pi-coding-agent-new-session)
    ("r" "resume" pi-coding-agent-resume-session)
    ("R" "reload" pi-coding-agent-reload)
    ("N" "name" pi-coding-agent-set-session-name)
    ("e" "export" pi-coding-agent-export-html)
    ("Q" "quit" pi-coding-agent-quit)]
   ["Context"
    ("c" "compact" pi-coding-agent-compact)
    ("f" "fork" pi-coding-agent-fork)]
   ["Actions"
    ("RET" "send" pi-coding-agent-send)
    ("s" "steer" pi-coding-agent-queue-steering)
    ("k" "abort" pi-coding-agent-abort)]]
  [["Model"
    ("m" "select" pi-coding-agent-select-model)
    ("t" "thinking" pi-coding-agent-select-thinking)]
   ["Completed thinking"
    (pi-coding-agent-menu-chat-thinking-display)
    (pi-coding-agent-menu-default-thinking-display)]
   ["Info"
    ("i" "stats" pi-coding-agent-session-stats)
    ("y" "copy last" pi-coding-agent-copy-last-message)]])

(defun pi-coding-agent-refresh-commands ()
  "Refresh commands from pi via RPC."
  (interactive)
  (if-let* ((proc (pi-coding-agent--get-process)))
      (let* ((chat-buf (pi-coding-agent--get-chat-buffer))
             (anchor (when (buffer-live-p chat-buf)
                       (with-current-buffer chat-buf
                         (pi-coding-agent--chat-session-directory chat-buf)))))
        (pi-coding-agent--fetch-commands proc
          (lambda (commands)
            (pi-coding-agent--set-commands commands)
            (pi-coding-agent--rebuild-commands-menu)
            (message "Pi: Refreshed %d commands" (length commands)))
          anchor))
    (message "Pi: No active process")))

;;;; Command Submenus (Templates, Extensions, Skills)

(defun pi-coding-agent--commands-by-source (source)
  "Return commands filtered by SOURCE, sorted alphabetically."
  (sort (seq-filter (lambda (c) (equal (plist-get c :source) source))
                    pi-coding-agent--commands)
        (lambda (a b)
          (string< (plist-get a :name) (plist-get b :name)))))

(defun pi-coding-agent--commands-by-source-and-location (source location)
  "Return commands filtered by SOURCE and LOCATION, sorted alphabetically."
  (sort (seq-filter (lambda (c)
                      (and (equal (plist-get c :source) source)
                           (equal (plist-get c :location) location)))
                    pi-coding-agent--commands)
        (lambda (a b)
          (string< (plist-get a :name) (plist-get b :name)))))

(defun pi-coding-agent--submenu-commands-ordered (source)
  "Return commands for SOURCE ordered by location then name.
Location order: path, project, user, then commands without location.
Within each location group, commands are sorted alphabetically by name.
This ordering is shared by run keys (a-z) and edit keys (A-Z)."
  (let ((path-cmds (pi-coding-agent--commands-by-source-and-location source "path"))
        (project-cmds (pi-coding-agent--commands-by-source-and-location source "project"))
        (user-cmds (pi-coding-agent--commands-by-source-and-location source "user"))
        (no-location-cmds (seq-filter (lambda (c)
                                        (and (equal (plist-get c :source) source)
                                             (null (plist-get c :location))))
                                      pi-coding-agent--commands)))
    (append path-cmds project-cmds user-cmds no-location-cmds)))

(defun pi-coding-agent--make-submenu-children (source)
  "Build transient children for commands with SOURCE.
Returns a list suitable for `transient-parse-suffixes'.
Commands are grouped by location (path, project, user).
Descriptions are truncated to fit the current frame width."
  (let* ((path-cmds (pi-coding-agent--commands-by-source-and-location source "path"))
         (project-cmds (pi-coding-agent--commands-by-source-and-location source "project"))
         (user-cmds (pi-coding-agent--commands-by-source-and-location source "user"))
         ;; Extensions don't have location, get them separately
         (no-location-cmds (seq-filter (lambda (c)
                                          (and (equal (plist-get c :source) source)
                                               (null (plist-get c :location))))
                                        pi-coding-agent--commands))
         (key 0)
         ;; Calculate available width for descriptions
         (available-width (max 20 (- (frame-width) 28)))
         (children '()))
    ;; Build location groups in order: path, project, user (then no-location for extensions)
    (dolist (group `(("Path" . ,path-cmds)
                     ("Project" . ,project-cmds)
                     ("User" . ,user-cmds)
                     (nil . ,no-location-cmds)))
      (let ((label (car group))
            (cmds (cdr group)))
        (when cmds
          ;; Add section header if there's a label
          (when label
            (push label children))
          ;; Add commands
          (dolist (cmd cmds)
            (when (< key 26)
              (let* ((name (plist-get cmd :name))
                     (desc (or (plist-get cmd :description) "")))
                ;; Run command with letter key (a-z)
                (push (list (format "%c" (+ ?a key))
                            (format "%-20s  %s"
                                    (truncate-string-to-width name 20)
                                    (truncate-string-to-width desc available-width))
                            `(lambda ()
                               (interactive)
                               (pi-coding-agent--run-custom-command ',cmd)))
                      children)
                (cl-incf key)))))))
    (nreverse children)))

(defun pi-coding-agent--edit-command-source (path)
  "Visit command source PATH, normalizing Pi paths for the current session."
  (let* ((chat-buf (pi-coding-agent--command-chat-buffer-or-error))
         (anchor (with-current-buffer chat-buf
                   (pi-coding-agent--chat-session-directory chat-buf)))
         (emacs-path (or (pi-coding-agent--emacs-path path anchor)
                         path)))
    (find-file-other-window emacs-path)))

(defun pi-coding-agent--make-submenu-edit-children (source)
  "Build edit suffixes for commands with SOURCE.
Returns a list suitable for `transient-parse-suffixes'.
Edit keys use uppercase letters (A-Z), matching the run keys (a-z).
Keys are assigned from the full command list so that `a' and `A'
always refer to the same command.  Commands without a :path are
skipped but still consume a key position."
  (let* ((all-cmds (pi-coding-agent--submenu-commands-ordered source))
         (key 0)
         (children '()))
    (dolist (cmd all-cmds)
      (when (< key 26)
        (let ((path (plist-get cmd :path)))
          (when path
            (let ((name (plist-get cmd :name)))
              (push (list (format "%c" (+ ?A key))
                          (truncate-string-to-width name 12)
                          `(lambda ()
                             (interactive)
                             (pi-coding-agent--edit-command-source ,path)))
                    children)))
          (cl-incf key))))
    (nreverse children)))

(defun pi-coding-agent--make-edit-columns (prefix source)
  "Build edit section as columns for SOURCE.
PREFIX is the transient command symbol.
Returns children for `:setup-children' as column group vectors."
  (let* ((items (pi-coding-agent--make-submenu-edit-children source))
         (len (length items)))
    (when (> len 0)
      (let* ((num-cols (min 3 len))
             (per-col (ceiling len (float num-cols)))
             (columns '()))
        (dotimes (i num-cols)
          (let* ((start (* i per-col))
                 (col-items (seq-subseq items start (min (+ start per-col) len))))
            (when col-items
              (push (vector 'transient-column
                            nil
                            (transient-parse-suffixes prefix col-items))
                    columns))))
        (nreverse columns)))))

(transient-define-prefix pi-coding-agent-templates-menu ()
  "All prompt templates.
Press letter to run, Shift+letter to edit source file."
  [:class transient-column
   :setup-children
   (lambda (_)
     (when-let* ((items (pi-coding-agent--make-submenu-children "prompt")))
       (transient-parse-suffixes 'pi-coding-agent-templates-menu items)))]
  [:class transient-columns
   :description "Edit"
   :setup-children
   (lambda (_)
     (pi-coding-agent--make-edit-columns
      'pi-coding-agent-templates-menu "prompt"))])

(transient-define-prefix pi-coding-agent-extensions-menu ()
  "All extension commands.
Press letter to run, Shift+letter to edit source file."
  [:class transient-column
   :setup-children
   (lambda (_)
     (when-let* ((items (pi-coding-agent--make-submenu-children "extension")))
       (transient-parse-suffixes 'pi-coding-agent-extensions-menu items)))]
  [:class transient-columns
   :description "Edit"
   :setup-children
   (lambda (_)
     (pi-coding-agent--make-edit-columns
      'pi-coding-agent-extensions-menu "extension"))])

(transient-define-prefix pi-coding-agent-skills-menu ()
  "All available skills.
Press letter to run, Shift+letter to edit source file."
  [:class transient-column
   :setup-children
   (lambda (_)
     (when-let* ((items (pi-coding-agent--make-submenu-children "skill")))
       (transient-parse-suffixes 'pi-coding-agent-skills-menu items)))]
  [:class transient-columns
   :description "Edit"
   :setup-children
   (lambda (_)
     (pi-coding-agent--make-edit-columns
      'pi-coding-agent-skills-menu "skill"))])

;;;; Main Menu Command Sections

(defun pi-coding-agent--rebuild-commands-menu ()
  "Rebuild command entries in transient menu.
Groups commands by source (extension, skill, template) with up to 3
quick-access commands per category and links to full submenus.
Sections are displayed side-by-side to use horizontal space."
  (let* ((extensions (pi-coding-agent--commands-by-source "extension"))
         (skills (pi-coding-agent--commands-by-source "skill"))
         (templates (pi-coding-agent--commands-by-source "prompt"))
         (columns '())
         (key 1))
    ;; Remove existing command group (index 3 if it exists)
    (ignore-errors (transient-remove-suffix 'pi-coding-agent-menu '(3)))
    ;; Build columns in display order: extensions, skills, templates
    ;; Keys are assigned sequentially across all categories
    (when extensions
      (push (pi-coding-agent--build-command-section
             "Extensions" extensions key 3 "E" 'pi-coding-agent-extensions-menu)
            columns)
      (setq key (+ key (min 3 (length extensions)))))
    (when skills
      (push (pi-coding-agent--build-command-section
             "Skills" skills key 3 "S" 'pi-coding-agent-skills-menu)
            columns)
      (setq key (+ key (min 3 (length skills)))))
    (when templates
      (push (pi-coding-agent--build-command-section
             "Templates" templates key 3 "T" 'pi-coding-agent-templates-menu)
            columns)
      (setq key (+ key (min 3 (length templates)))))
    ;; Add all columns as a single transient-columns group after the static rows.
    (when columns
      (transient-append-suffix 'pi-coding-agent-menu '(2)
        (apply #'vector (nreverse columns))))))

(defun pi-coding-agent--build-command-section (title commands start-key max-shown more-key more-menu)
  "Build a transient section for TITLE with COMMANDS.
Shows up to MAX-SHOWN commands starting at START-KEY.
MORE-KEY and MORE-MENU provide access to the full list (shown first)."
  (let ((shown (seq-take commands max-shown))
        (suffixes (list title))
        (key start-key))
    ;; Add "all..." link first for discovery
    (push (list more-key "all..." more-menu) suffixes)
    ;; Add quick-access commands
    (dolist (cmd shown)
      (let ((name (plist-get cmd :name)))
        (push (list (number-to-string key)
                    (truncate-string-to-width name 18)
                    `(lambda () (interactive) (pi-coding-agent--run-custom-command ',cmd)))
              suffixes)
        (setq key (1+ key))))
    (apply #'vector (nreverse suffixes))))

(provide 'pi-coding-agent-menu)
;;; pi-coding-agent-menu.el ends here