;;; codex.el --- Simple Codex chat buffer -*- lexical-binding: t; -*-

;; Copyright (C) 2025

;; Author: Grant Horner
;; Version: 0.1
;; Package-Requires: ((emacs "27.1"))
;; Keywords: tools, convenience

;;; Commentary:
;; Basic chat buffer for Codex using `codex exec --json`.

;; {"type":"thread.started","thread_id":"019ba390-af9c-7872-8617-c8038b61d559"}
;; {"type":"turn.started"}
;; {"type":"item.completed","item":{"id":"item_0","type":"reasoning","text":"**Responding with greeting**"}}
;; {"type":"item.completed","item":{"id":"item_1","type":"agent_message","text":"Hi! How can I help?"}}
;; {"type":"turn.completed","usage":{"input_tokens":4050,"cached_input_tokens":3712,"output_tokens":13}}

;;; Code:

(require 'json)
(require 'cl-lib)
(require 'subr-x)

(defgroup codex nil
  "Chat with Codex."
  :group 'tools)

(defface codex-prompt-delimiter-face
  '((t :inherit shadow))
  "Face for prompt delimiters."
  :group 'codex)

(defface codex-user-prompt-face
  '((t :inherit font-lock-keyword-face))
  "Face for the user prompt label."
  :group 'codex)

(defface codex-assistant-prompt-face
  '((t :inherit font-lock-function-name-face))
  "Face for the Codex prompt label."
  :group 'codex)

(defcustom codex-command "codex"
  "Codex executable to invoke."
  :type 'string
  :group 'codex)

(defcustom codex-provider nil
  "Custom codex provider to use.
Note on Copilot:
In order for copilot integration to work, you need a section in your
~/.codex/config.toml like this:
[model_providers.github-copilot]
name = \"Github Copilot\"
base_url = \"https://api.githubcopilot.com\"
env_key = \"GITHUB_COPILOT_TOKEN\"
wire_api = \"chat\"
http_headers = {
    Copilot-Integration-Id = \"vscode-chat\"
}"
  :type 'string
  :group 'codex)

(defcustom codex-model nil
  "Model for Codex to use."
  :type 'string
  :group 'codex)

(defcustom codex-skip-git-repo-check t
  "Whether to pass --skip-git-repo-check to Codex."
  :type 'boolean
  :group 'codex)

(defcustom codex-buffer-name "*codex*"
  "What the default name should be fore the codex chat buffer."
  :type 'string
  :group 'codex)

(defcustom codex-wrap-lines t
  "Whether to truncate long lines in `codex-mode`."
  :type 'boolean
  :group 'codex)

(defcustom codex-sessions-file (concat (file-name-parent-directory user-init-file) "codex-sessions")
  "What the default name should be fore the codex chat buffer."
  :type 'string
  :group 'codex)

(defcustom codex-sessions-history-dir
  (concat (file-name-parent-directory user-init-file) "codex-sessions-history/")
  "Directory where Codex session histories are stored."
  :type 'directory
  :group 'codex)

(defvar codex--session-id nil)

(defvar codex-mode-map
  (let ((map (make-sparse-keymap)))
    (define-key map (kbd "C-c C-c") #'codex-send)
    (define-key map (kbd "C-c C-s") #'codex-switch-sessions)
    map)
  "Keymap for `codex-mode'.")

(defconst codex--font-lock-keywords
  '(("^---$" . 'codex-prompt-delimiter-face)
    ("^User:$" . 'codex-user-prompt-face)
    ("^Codex:$" . 'codex-assistant-prompt-face)))

;;;###autoload
(define-derived-mode codex-mode fundamental-mode "Codex"
  "Major mode for chatting with Codex."
  (setq-local buffer-read-only nil)
  (setq-local truncate-lines (not codex-wrap-lines))
  (setq-local font-lock-defaults '(codex--font-lock-keywords))
  (setq-local codex--busy nil))

(defvar-local codex--prompt-start nil)

(defun codex--ensure-history-dir ()
  (make-directory codex-sessions-history-dir t))

(defun codex--history-file ()
  (when (and codex--session-id (not (string-empty-p codex--session-id)))
    (expand-file-name codex--session-id codex-sessions-history-dir)))

(defun codex--set-prompt-start-from-buffer ()
  (let* ((haystack (buffer-string))
         (needle "---\nUser:\n")
         (index (cl-search needle haystack :from-end t)))
    (when index
      (setq codex--prompt-start (copy-marker (+ (length needle) index)))
      t)))

(defun codex--save-session-history ()
  (with-codex-buffer
   (let ((history-file (codex--history-file)))
     (when history-file
       (codex--ensure-history-dir)
       (write-region (point-min) (point-max) history-file nil 'silent)))))

(defun codex--load-session-history ()
  (with-codex-buffer
   (let ((history-file (codex--history-file)))
     (setq codex--prompt-start nil)
     (erase-buffer)
     (if (and history-file (file-exists-p history-file))
         (progn
           (insert-file-contents history-file)
           (goto-char (point-max))
           (unless (codex--set-prompt-start-from-buffer)
             (user-prompt)))
       (user-prompt)))))

(defun list->hash-set (list &optional test)
  "Return a hash table whose keys are the elements of LIST."
  (let ((ht (make-hash-table :test (or test #'equal))))
    (dolist (x list ht)
      (puthash x t ht))))

(defmacro with-codex-buffer (&rest body)
  `(with-current-buffer (get-buffer-create codex-buffer-name)
     ,@body))

(defun prompt-delimiter ()
  (newline)
  (insert "---")
  (newline))

(defun user-prompt ()
  (prompt-delimiter)
  (insert "User:")
  (newline)
  (setq codex--prompt-start (point-marker)))

(defun codex--write-to-chat (msg)
  (with-codex-buffer
   (newline)
   (insert "Codex:")
   (newline)
   (insert msg)
   (user-prompt)
   (codex--save-session-history)))

(defun codex--parse-session-id (jsons)
  (let ((thread-started-message
         (or (cl-find-if
              (lambda (hm)
                (and (string-equal (gethash "type" hm) "thread.started")
                     (gethash "thread_id" hm)))
              jsons)
             (make-hash-table))))
    (gethash "thread_id" thread-started-message)))

(defun codex--parse-msg-from-response (response-string)
  (let* ((output (string-trim response-string))
         (lines (string-lines output))
         (jsons (delq nil
                      (mapcar (lambda (line)
                                (condition-case nil
                                    (json-parse-string line)
                                  (json-parse-error nil)))
                              lines)))
         (msg (cl-find-if (lambda (hm) (let ((type (gethash "type" hm))
                                             (item (gethash "item" hm)))
                                         (and type
                                              item
                                              (string-equal type "item.completed")
                                              (string-equal (gethash "type" item) "agent_message"))))
                          jsons))
         (item (and msg (gethash "item" msg)))
         (text (and item (gethash "text" item))))
        ;; We should make sure we get one and only one message here, otherwise bail out
    (setf codex--session-id (codex--parse-session-id jsons))
    (or text output)))

(defun codex--ensure-session-in-sessions-file (prompt)
  (let ((sessions (codex--read-sessions-file)))
    (unless (member codex--session-id (mapcar #'cdr sessions))
      (codex--write-sessions-file (cons (cons prompt codex--session-id) sessions)))))

(defun codex--format-response (msg)
  (with-temp-buffer
    (insert msg)
    (fill-region (point-min) (point-max))
    (buffer-string)))

(defun codex--send (prompt)
  "Sends a prompt to codex."
  (let* ((buf (generate-new-buffer "*codex--send*"))
         (proc (make-process
                :name "codex"
                :buffer buf
                :command
                (let* ((resume (when (and codex--session-id (not (string-empty-p codex--session-id)))
                                (progn
                                  (message "Using session %s" codex--session-id)
                                  (list "resume" codex--session-id))))
                       (skip (and codex-skip-git-repo-check (list "--skip-git-repo-check")))
                       (provider (and codex-provider (list "-c" (concat "model_provider=" codex-provider))))
                       (model (and codex-model (list "-m" codex-model)))
                       (command (append (list codex-command "exec") resume skip provider model (list "--json" prompt))))
                  command))))
    (set-process-sentinel
     proc
     (lambda (p event)
       (message event)
       (when (string= event "finished\n")
         (with-current-buffer (process-buffer p)
           (thread-first
             (buffer-string)
             codex--parse-msg-from-response
             codex--format-response
             codex--write-to-chat)
           (codex--ensure-session-in-sessions-file prompt)
           (kill-buffer)))
       (when (string-prefix-p "exited abnormally" event)
         (with-current-buffer (process-buffer p)
           (error (buffer-string))))))))

(defun codex--get-current-user-prompt ()
  (with-codex-buffer
    (let ((prompt-start (and codex--prompt-start
                             (marker-buffer codex--prompt-start)
                             (marker-position codex--prompt-start))))
      ;; fallback in case we somehow lost where our user prompt starts
      (unless prompt-start
        (when (codex--set-prompt-start-from-buffer)
          (setq prompt-start (marker-position codex--prompt-start))))
      (unless prompt-start
        (error "No current prompt start"))
      (buffer-substring-no-properties prompt-start (point-max)))))

(defun codex-send ()
  (interactive)
  (with-current-buffer codex-buffer-name
    (let ((prompt (codex--get-current-user-prompt)))
      (if (string-equal prompt "")
          (error "Prompt is empty!")
        (newline)
        (insert "---")
        (codex--save-session-history)
        (codex--send prompt)))))

;;;###autoload
(defun codex (session)
  "Open a Codex chat buffer."
  (interactive
   (list (completing-read "Session: " (cons '("New") (codex--read-sessions-file)))))
  (let ((session-id (cdr (assoc session (codex--read-sessions-file)))))
   (setf codex--session-id session-id)
   (let ((buf (get-buffer-create codex-buffer-name)))
     (pop-to-buffer buf)
     (unless (derived-mode-p 'codex-mode)
       (codex-mode))
     (codex--load-session-history)
     (goto-char (point-max)))))

(defun codex--read-sessions-file ()
  (if (file-exists-p codex-sessions-file)
    (with-temp-buffer
     (insert-file-contents codex-sessions-file)
     (goto-char (point-min))
     (condition-case nil
         (read (current-buffer))
       (end-of-file nil)))
    (codex--write-sessions-file '())
    nil))

(defun codex--write-sessions-file (list-of-session-ids)
  (let ((filename codex-sessions-file))
    (with-temp-buffer
      (insert ";;; -*- lisp-data -*-\n")
      (let ((print-length nil)
            (print-level nil))
        (pp list-of-session-ids (current-buffer)))
      (write-region nil nil filename nil 'silent))))

(defun codex-switch-sessions (new-session)
  (interactive
   (list (completing-read "Session: " (cons '("New") (codex--read-sessions-file)))))
  (let ((session-id (cdr (assoc new-session (codex--read-sessions-file)))))
    (codex--save-session-history)
    (setf codex--session-id session-id)
    (codex--load-session-history)))

(defun codex-rename-current-session (name)
  (interactive "sNew name: ")
  (let* ((sessions (codex--read-sessions-file))
         (new-sessions
          (mapcar
           (lambda (tup)
             (if (string-equal (cdr tup) codex--session-id)
                 (cons name (cdr tup))
               tup))
           sessions)))
    (codex--write-sessions-file new-sessions)))

(provide 'codex)


;;; codex.el ends here