;;; ob-jupyter.el --- Jupyter integration with org-mode -*- lexical-binding: t -*-

;; Copyright (C) 2018 Nathaniel Nicandro

;; Author: Nathaniel Nicandro <nathanielnicandro@gmail.com>
;; Created: 21 Jan 2018

;; 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, 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 GNU Emacs; see the file COPYING.  If not, write to the
;; Free Software Foundation, Inc., 59 Temple Place - Suite 330,
;; Boston, MA 02111-1307, USA.

;;; Commentary:

;; Interact with a Jupyter kernel via `org-mode' src-block's.

;;; Code:

(defgroup ob-jupyter nil
  "Jupyter integration with org-mode"
  :group 'org-babel)

(require 'jupyter-env)
(require 'jupyter-kernelspec)
(require 'jupyter-org-client)
(require 'jupyter-org-extensions)
(eval-when-compile
  (require 'jupyter-repl) ; For `jupyter-with-repl-buffer'
  (require 'subr-x))

(declare-function org-in-src-block-p "org" (&optional inside))
(declare-function org-element-property "org-element" (property element))
(declare-function org-element-context "org-element" (&optional element))
(declare-function org-babel-variable-assignments:python "ob-python" (params))
(declare-function org-babel-expand-body:generic "ob-core" (body params &optional var-lines))
(declare-function org-export-derived-backend-p "ox" (backend &rest backends))

(declare-function jupyter-server-make-instance "jupyter-server")
(declare-function jupyter-run-server-repl "jupyter-server")
(declare-function jupyter-connect-server-repl "jupyter-server")
(declare-function jupyter-server-kernelspecs "jupyter-server")
(declare-function jupyter-server-kernel-id-from-name "jupyter-server")
(declare-function jupyter-server-name-client-kernel "jupyter-server")
(declare-function jupyter-api-get-kernel "jupyter-rest-api")

(declare-function jupyter-tramp-url-from-file-name "jupyter-tramp")
(declare-function jupyter-tramp-server-from-file-name "jupyter-tramp")
(declare-function jupyter-tramp-file-name-p "jupyter-tramp")

(defvaralias 'org-babel-jupyter-resource-directory
  'jupyter-org-resource-directory)

(defvar org-babel-jupyter-session-clients (make-hash-table :test #'equal)
  "A hash table mapping session names to Jupyter clients.
`org-babel-jupyter-src-block-session' returns a key into this
table for the source block at `point'.")

(defvar org-babel-header-args:jupyter '((kernel . :any)
                                        (async . ((yes no))))
  "Available header arguments for Jupter src-blocks.")

(defvar org-babel-default-header-args:jupyter '((:kernel . "python")
                                                (:async . "no"))
  "Default header arguments for Jupyter src-blocks.")

;;; Helper functions

(defun org-babel-jupyter--src-block-kernel-language ()
  (when (org-in-src-block-p)
    (let ((info (org-babel-get-src-block-info)))
      (save-match-data
        (string-match "^jupyter-\\(.+\\)$" (car info))
        (match-string 1 (car info))))))

(defun org-babel-jupyter-language-p (lang)
  "Return non-nil if LANG src-blocks are executed using Jupyter."
  (or (string-prefix-p "jupyter-" lang)
      ;; Check if the language has been overridden, see
      ;; `org-babel-jupyter-override-src-block'
      (advice-member-p
       'ob-jupyter (intern (concat "org-babel-execute:" lang)))))

(defun org-babel-jupyter-session-key (params)
  "Return a string that is the concatenation of the :session and :kernel PARAMS.
PARAMS is the arguments alist as returned by
`org-babel-get-src-block-info'.  The returned string can then be
used to identify unique Jupyter Org babel sessions."
  (let ((session (alist-get :session params))
        (kernel (alist-get :kernel params)))
    (unless (and session kernel
                 (not (equal session "none")))
      (error "Need a valid session and a kernel to form a key"))
    (concat session "-" kernel)))

(defun org-babel-jupyter-src-block-session ()
  "Return the session key for the current Jupyter source block.
Return nil if the current source block is not a Jupyter block or
if there is no source block at point."
  (let ((info (or (and (org-in-src-block-p)
                       (org-babel-get-src-block-info 'light))
                  (org-babel-lob-get-info))))
    (when info
      (cl-destructuring-bind (lang _ params . rest) info
        (when (org-babel-jupyter-language-p lang)
          (org-babel-jupyter-session-key params))))))

;;; `ob' integration

(defun org-babel-variable-assignments:jupyter (params &optional lang)
  "Assign variables in PARAMS according to the Jupyter kernel language.
LANG is the kernel language of the source block.  If LANG is nil,
get the kernel language from the current source block.

The variables are assigned by looking for the function
`org-babel-variable-assignments:LANG'.  If this function does not
exist or if LANG cannot be determined, assign variables using
`org-babel-variable-assignments:python'."
  (or lang (setq lang (org-babel-jupyter--src-block-kernel-language)))
  (let ((fun (when lang
               (intern (format "org-babel-variable-assignments:%s" lang)))))
    (if (functionp fun) (funcall fun params)
      (require 'ob-python)
      (org-babel-variable-assignments:python params))))

(cl-defgeneric org-babel-jupyter-transform-code (code _changelist)
  "Transform CODE according to CHANGELIST, return the transformed CODE.
CHANGELIST is a property list containing the requested changes.  The default
implementation returns CODE unchanged.

This is useful for kernel languages to extend using the
jupyter-lang method specializer, e.g. to return new code to change
directories before evaluating CODE.

See `org-babel-expand-body:jupyter' for possible changes that can
be in CHANGELIST."
  code)

(defun org-babel-expand-body:jupyter (body params &optional var-lines lang)
  "Expand BODY according to PARAMS.

BODY is the code to expand, PARAMS should be the header arguments
of the src block with BODY as its code, and VAR-LINES should be
the list of strings containing the variables to evaluate before
executing body.  LANG is the kernel language of the source block.

This function is similar to
`org-babel-variable-assignments:jupyter' in that it attempts to
find the kernel language of the source block if LANG is not
provided.

BODY is expanded by calling the function
`org-babel-expand-body:LANG'.  If this function doesn't exist or
if LANG cannot be determined, fall back to
`org-babel-expand-body:generic'.

If PARAMS has a :dir parameter, the expanded code is passed to
`org-babel-jupyter-transform-code' with a changelist that
includes the :dir parameter with the directory being an absolute
path."
  (or lang (setq lang (org-babel-jupyter--src-block-kernel-language)))
  (let* ((expander (when lang
                     (intern (format "org-babel-expand-body:%s" lang))))
         (expanded (if (functionp expander)
                       (funcall expander body params)
                     (org-babel-expand-body:generic body params var-lines)))
         (changelist nil))
    (when-let* ((dir (alist-get :dir params)))
      (setq changelist (plist-put changelist :dir (expand-file-name dir))))
    (if changelist (org-babel-jupyter-transform-code expanded changelist)
      expanded)))

(defun org-babel-edit-prep:jupyter (info)
  "Prepare the edit buffer according to INFO.
Enable `jupyter-repl-interaction-mode' in the edit buffer
associated with the session found in INFO.

If the session is a Jupyter TRAMP file name, the
`default-directory' of the edit buffer is set to the root
directory the notebook serves."
  (let* ((params (nth 2 info))
         (session (alist-get :session params))
         (client-buffer (org-babel-jupyter-initiate-session session params)))
    (jupyter-repl-associate-buffer client-buffer)
    (when (jupyter-tramp-file-name-p session)
      (setq default-directory (concat (file-remote-p session) "/")))))

(defun org-babel-jupyter--insert-variable-assignments (params)
  "Insert variable assignment lines from PARAMS into the `current-buffer'.
Return non-nil if there are variable assignments, otherwise
return nil."
  (let ((var-lines (org-babel-variable-assignments:jupyter params)))
    (prog1 var-lines
      (jupyter-repl-replace-cell-code (mapconcat #'identity var-lines "\n")))))

(defun org-babel-prep-session:jupyter (session params)
  "Prepare a Jupyter SESSION according to PARAMS."
  (with-current-buffer (org-babel-jupyter-initiate-session session params)
    (goto-char (point-max))
    (and (org-babel-jupyter--insert-variable-assignments session params)
         (jupyter-send-execute-request jupyter-current-client))
    (current-buffer)))

(defun org-babel-load-session:jupyter (session body params)
  "In a Jupyter SESSION, load BODY according to PARAMS."
  (save-window-excursion
    (with-current-buffer (org-babel-jupyter-initiate-session session params)
      (goto-char (point-max))
      (when (org-babel-jupyter--insert-variable-assignments session params)
        (insert "\n"))
      (insert (org-babel-expand-body:jupyter (org-babel-chomp body) params))
      (current-buffer))))

;;; Initializing session clients

(cl-defstruct (org-babel-jupyter-session
               (:constructor org-babel-jupyter-session))
  name)

(cl-defstruct (org-babel-jupyter-remote-session
               (:include org-babel-jupyter-session)
               (:constructor org-babel-jupyter-remote-session))
  connect-repl-p)

(cl-defstruct (org-babel-jupyter-server-session
               (:include org-babel-jupyter-remote-session)
               (:constructor org-babel-jupyter-server-session)))

(defun org-babel-jupyter-parse-session (session)
  "Return a session object according to a SESSION string.
If SESSION ends in \".json\" return a
`org-babel-jupyter-remote-session' that indicates an Org Babel
Jupyter client initiates its channels based on a kernel
connection file.

If SESSION is a Jupyter TRAMP file name return a
`org-babel-jupyter-server-session', otherwise if SESSION is a
remote file return an `org-babel-jupyter-remote-session'.  In the
latter case, a kernel will be launched on the remote and a
connection file read via TRAMP and SSH tunnels created to connect
to the kernel.

Otherwise an `org-babel-jupyter-session' is returned which
indicates that the session is local."
  (cond
   ((string-suffix-p ".json" session)
    (org-babel-jupyter-remote-session
     :name session
     :connect-repl t))
   ((file-remote-p session)
    (if (jupyter-tramp-file-name-p session)
        (org-babel-jupyter-server-session :name session)
      (org-babel-jupyter-remote-session :name session)))
   (t
    (org-babel-jupyter-session :name session))))

(cl-defgeneric org-babel-jupyter-initiate-client (session kernel)
  "Launch SESSION's KERNEL, return a `jupyter-org-client' connected to it.
SESSION is the name of the :session header argument of a source
block and KERNEL is the name of the kernelspec used to launch a
kernel.")

(cl-defmethod org-babel-jupyter-initiate-client ((_session org-babel-jupyter-session) kernel)
  (jupyter-run-repl kernel nil nil 'jupyter-org-client))

(cl-defmethod org-babel-jupyter-initiate-client :before ((session org-babel-jupyter-remote-session) _kernel)
  "Raise an error if SESSION a remote file name without a local name.
The local name is used as a unique identifier of a remote
session."
  (unless (not (zerop (length (file-local-name
                               (org-babel-jupyter-session-name session)))))
    (error "No remote session name")))

(cl-defmethod org-babel-jupyter-initiate-client :around (session _kernel)
  (let ((client (cl-call-next-method)))
    (prog1 client
      (jupyter-set client 'jupyter-include-other-output nil)
      ;; Append the name of SESSION to the initiated client REPL's
      ;; `buffer-name'.
      (jupyter-with-repl-buffer client
        (let ((name (buffer-name)))
          (when (string-match "^\\*\\(.+\\)\\*" name)
            (rename-buffer
             (concat "*" (match-string 1 name) "-"
                     (org-babel-jupyter-session-name session)
                     "*")
             'unique)))))))

(cl-defmethod org-babel-jupyter-initiate-client ((session org-babel-jupyter-remote-session) _kernel)
  (let ((session-name (org-babel-jupyter-remote-session-name session)))
    (if (org-babel-jupyter-remote-session-connect-repl-p session)
        (jupyter-connect-repl session-name nil nil 'jupyter-org-client)
      (let ((default-directory (file-remote-p session-name)))
        (org-babel-jupyter-aliases-from-kernelspecs)
        (jupyter-run-repl kernel nil nil 'jupyter-org-client)))))

(cl-defmethod org-babel-jupyter-initiate-client ((session org-babel-jupyter-server-session) kernel)
  (require 'jupyter-server)
  (let* ((session (org-babel-jupyter-server-session-name session))
         (server (or (jupyter-tramp-server-from-file-name session)
                     (jupyter-server-make-instance
                      :url (jupyter-tramp-url-from-file-name session)))))
    (unless (jupyter-server-has-kernelspec-p server kernel)
      (error "No kernelspec matching \"%s\" exists at %s"
             kernel (jupyter-tramp-url-from-file-name session)))
    ;; Language aliases may not exist for the kernels that are accessible on
    ;; the server so ensure they do.
    (org-babel-jupyter-aliases-from-kernelspecs
     nil (jupyter-server-kernelspecs server))
    (let ((session-name (file-local-name session)))
      (if-let ((id (jupyter-server-kernel-id-from-name server session-name)))
          ;; Connecting to an existing kernel
          (cl-destructuring-bind (&key name id &allow-other-keys)
              (or (ignore-errors (jupyter-api-get-kernel server id))
                  (error "Kernel ID, %s, no longer references a kernel @ %s"
                         id (oref server url)))
            (unless (string-match-p kernel name)
              (error "\":kernel %s\" doesn't match \"%s\"" kernel name))
            (jupyter-connect-server-repl server id nil nil 'jupyter-org-client))
        ;; Start a new kernel
        (let ((client (jupyter-run-server-repl
                       server kernel nil nil 'jupyter-org-client)))
          (prog1 client
            ;; TODO: If a kernel gets renamed in the future it doesn't affect
            ;; any source block :session associations because the hash of the
            ;; session name used here is already stored in the
            ;; `org-babel-jupyter-session-clients' variable.  Should that
            ;; variable be updated on a kernel rename?
            ;;
            ;; TODO: Would we always want to do this?
            (jupyter-server-name-client-kernel client session-name)))))))

(defun org-babel-jupyter-initiate-session-by-key (session params)
  "Return the Jupyter REPL buffer for SESSION.
If SESSION does not have a client already, one is created based
on SESSION and PARAMS.  If SESSION ends with \".json\" then
SESSION is interpreted as a kernel connection file and a new
kernel connected to SESSION is created.

Otherwise a kernel is started based on the `:kernel' parameter in
PARAMS which should be either a valid kernel name or a prefix of
one, in which case the first kernel that matches the prefix will
be used.

If SESSION is a remote file name, like /ssh:ec2:jl, then the
kernel starts on the remote host /ssh:ec2: with a session name of
jl.  The remote host must have jupyter installed since the
\"jupyter kernel\" command will be used to start the kernel on
the host."
  (let* ((kernel (alist-get :kernel params))
         (key (org-babel-jupyter-session-key params))
         (client (gethash key org-babel-jupyter-session-clients)))
    (unless client
      (setq client (org-babel-jupyter-initiate-client
                    (org-babel-jupyter-parse-session session)
                    kernel))
      (puthash key client org-babel-jupyter-session-clients)
      (let ((forget-client (lambda () (remhash key org-babel-jupyter-session-clients))))
        (add-hook 'kill-buffer-hook forget-client nil t)))
    (oref client buffer)))

(defun org-babel-jupyter-initiate-session (&optional session params)
  "Initialize a Jupyter SESSION according to PARAMS."
  (if (equal session "none") (error "Need a session to run")
    (org-babel-jupyter-initiate-session-by-key session params)))

;;;###autoload
(defun org-babel-jupyter-scratch-buffer ()
  "Display a scratch buffer connected to the current block's session."
  (interactive)
  (let (buffer)
    (org-babel-do-in-edit-buffer
     (setq buffer (save-window-excursion
                    (jupyter-repl-scratch-buffer))))
    (if buffer (pop-to-buffer buffer)
      (user-error "No source block at point"))))

(defvar org-bracket-link-regexp)

(defun org-babel-jupyter-cleanup-file-links ()
  "Delete the files of image links for the current source block result.
Do this only if the file exists in
`org-babel-jupyter-resource-directory'."
  (when-let* ((result-pos (org-babel-where-is-src-block-result))
              (link-re (format "^[ \t]*%s[ \t]*$" org-bracket-link-regexp)))
    (save-excursion
      (goto-char result-pos)
      (forward-line)
      (let ((bound (org-babel-result-end)))
        ;; This assumes that `jupyter-org-client' only emits bracketed links as
        ;; images
        (while (re-search-forward link-re bound t)
          (when-let* ((link-path
                       (org-element-property :path (org-element-context)))
                      (link-dir
                       (when (file-name-directory link-path)
                         (expand-file-name (file-name-directory link-path))))
                      (resource-dir
                       (expand-file-name org-babel-jupyter-resource-directory)))
            (when (and (equal link-dir resource-dir)
                       (file-exists-p link-path))
              (delete-file link-path))))))))

;; TODO: What is a better way to handle discrepancies between how `org-mode'
;; views header arguments and how `emacs-jupyter' views them? Should the
;; strategy be to always try to emulate the `org-mode' behavior?
(defun org-babel-jupyter--remove-file-param (params)
  "Destructively remove the file result parameter from PARAMS.
These parameters are handled internally."
  (let* ((result-params (assq :result-params params))
         (fresult (member "file" result-params))
         (fparam (assq :file params)))
    (setcar fresult "")
    (delq fparam params)))

(defvar org-babel-jupyter-current-src-block-params nil
  "The block parameters of the most recently executed Jupyter source block.")

(defconst org-babel-jupyter-async-inline-results-pending-indicator "???"
  "A string to disambiguate pending inline results from empty results.")

(defun org-babel-execute:jupyter (body params)
  "Execute BODY according to PARAMS.
BODY is the code to execute for the current Jupyter `:session' in
the PARAMS alist."
  (let* ((org-babel-jupyter-current-src-block-params params)
         (jupyter-current-client
          (thread-first (alist-get :session params)
            (org-babel-jupyter-initiate-session params)
            (thread-last (buffer-local-value 'jupyter-current-client))))
         (kernel-lang (jupyter-kernel-language jupyter-current-client))
         (vars (org-babel-variable-assignments:jupyter params kernel-lang))
         (code (org-babel-expand-body:jupyter body params vars kernel-lang))
         (result-params (assq :result-params params))
         (async-p (or (equal (alist-get :async params) "yes")
                      (plist-member params :async)))
         (req (jupyter-send-execute-request jupyter-current-client :code code)))
    (when (member "replace" (assq :result-params params))
      (org-babel-jupyter-cleanup-file-links))
    ;; KLUDGE: Remove the file result-parameter so that
    ;; `org-babel-insert-result' doesn't attempt to handle it while async
    ;; results are pending.  Do the same in the synchronous case, but not if
    ;; link or graphics are also result-parameters, only in Org >= 9.2, since
    ;; those in combination with file mean to interpret the result as a file
    ;; link, a useful meaning that doesn't interfere with Jupyter style result
    ;; insertion.
    (when (and (member "file" result-params)
               (or async-p
                   (not (or (member "link" result-params)
                            (member "graphics" result-params)))))
      (org-babel-jupyter--remove-file-param params))
    (cond
     (async-p
      (cl-labels
          ((sync-on-export
            ()
            ;; Remove the hook before waiting so it doesn't get called again.
            (remove-hook 'org-babel-after-execute-hook #'sync-on-export t)
            (while (null (jupyter-wait-until-idle req jupyter-long-timeout)))))
        ;; Ensure we convert async blocks to synchronous ones when exporting
        (when (bound-and-true-p org-export-current-backend)
          (add-hook 'org-babel-after-execute-hook #'sync-on-export t t))
        (if (jupyter-org-request-inline-block-p req)
            org-babel-jupyter-async-inline-results-pending-indicator
          (jupyter-org-pending-async-results req))))
     (t
      (while (null (jupyter-wait-until-idle req jupyter-long-timeout)))
      (if (jupyter-org-request-inline-block-p req)
          ;; When evaluating a source block synchronously, only the
          ;; :execute-result will be in `jupyter-org-request-results' since
          ;; stream results and any displayed data will be placed in a separate
          ;; buffer.
          (car (jupyter-org-request-results req))
        (prog1 (jupyter-org-sync-results req)
          ;; KLUDGE: The "raw" result parameter is added to the parameters of
          ;; non-inline synchronous results because `jupyter-org-sync-results'
          ;; already returns an Org formatted string and
          ;; `org-babel-insert-result' should not process it.
          (nconc (alist-get :result-params params) (list "raw"))))))))

;;; Overriding source block languages, language aliases

(defvar org-babel-jupyter--babel-ops
  '("execute" "expand-body" "prep-session" "edit-prep"
    "variable-assignments" "load-session"))

(defun org-babel-jupyter--override-restore-header-args (lang restore)
  "Set `org-babel-header-args:LANG' to its Jupyter equivalent.
`org-babel-header-args:LANG' is set to the value of
`org-babel-header-args:jupyter-LANG', if the latter exists, when
RESTORE is nil.  If `org-babel-header-args:LANG' had a value, save
it as a symbol property of `org-babel-header-args:LANG' for
restoring it later.

If RESTORE is non-nil, set `org-babel-header-args:LANG' to its
saved value before it was overridden.

Do the same for `org-babel-default-header-args:LANG'."
  (dolist (prefix '("org-babel-header-args:"
                    "org-babel-default-header-args:"))
    (when-let* ((jupyter-var (intern-soft (concat prefix "jupyter-" lang))))
      (let ((var (intern-soft (concat prefix lang))))
        (if restore
            (set var (get var 'jupyter-restore-value))
          (if var (put var 'jupyter-restore-value (symbol-value var))
            (setq var (intern (concat prefix lang))))
          (set var (symbol-value jupyter-var)))))))

(defun org-babel-jupyter--override-restore-src-block (lang restore)
  (cl-macrolet ((override-restore
                 (sym jupyter-sym)
                 `(cond
                   (restore
                    (advice-remove ,sym ,jupyter-sym)
                    ;; The function didn't have a definition, so ensure that
                    ;; we restore that fact.
                    (when (eq (symbol-function ,sym) #'ignore)
                      (fmakunbound ,sym)))
                   (t
                    ;; If a language doesn't have a function assigned, set one
                    ;; so it can be overridden
                    (unless (fboundp ,sym)
                      (fset ,sym #'ignore))
                    (advice-add ,sym :override ,jupyter-sym
                                '((name . ob-jupyter)))))))
    (dolist (fn (cl-set-difference
                 org-babel-jupyter--babel-ops
                 '("variable-assignments" "expand-body")
                 :test #'equal))
      (let ((sym (intern (concat "org-babel-" fn ":" lang))))
        (override-restore sym (intern (concat "org-babel-" fn ":jupyter-" lang)))))
    (override-restore (intern (concat "org-babel-" lang "-initiate-session"))
                      #'org-babel-jupyter-initiate-session))
  (org-babel-jupyter--override-restore-header-args lang restore))

(defun org-babel-jupyter-override-src-block (lang)
  "Override the built-in `org-babel' functions for LANG.
This overrides functions like `org-babel-execute:LANG' and
`org-babel-LANG-initiate-session' to use the machinery of
jupyter-LANG source blocks."
  (org-babel-jupyter--override-restore-src-block lang nil))

(defun org-babel-jupyter-restore-src-block (lang)
  "Restore the overridden `org-babel' functions for LANG.
See `org-babel-jupyter-override-src-block'."
  (org-babel-jupyter--override-restore-src-block lang t))

(defun org-babel-jupyter-make-language-alias (kernel lang)
  "Similar to `org-babel-make-language-alias' but for Jupyter src-blocks.
KERNEL should be the name of the default kernel to use for kernel
LANG.  All necessary org-babel functions for a language with the
name jupyter-LANG will be aliased to the Jupyter functions."
  (dolist (fn org-babel-jupyter--babel-ops)
    (let ((sym (intern-soft (concat "org-babel-" fn ":jupyter"))))
      (when (and sym (fboundp sym))
        (defalias (intern (concat "org-babel-" fn ":jupyter-" lang)) sym))))
  (defalias (intern (concat "org-babel-jupyter-" lang "-initiate-session"))
    'org-babel-jupyter-initiate-session)
  (let (var)
    (setq var (concat "org-babel-header-args:jupyter-" lang))
    (unless (intern-soft var)
      (set (intern var) org-babel-header-args:jupyter))
    (put (intern var) 'variable-documentation
         (get 'org-babel-header-args:jupyter 'variable-documentation))
    (setq var (concat "org-babel-default-header-args:jupyter-" lang))
    (unless (and (intern-soft var)
                 (boundp (intern var)))
      (set (intern var) `((:async . "no"))))
    (put (intern var) 'variable-documentation
         (format "Default header arguments for Jupyter %s src-blocks" lang))
    ;; Always set the kernel if there isn't one.  Typically the default header
    ;; args for a language are set by the user in their configurations by
    ;; calling `setq', but the :kernel is typically not something the user
    ;; wants to set directly so make sure its defined in the header args.
    (setq var (intern var))
    (unless (alist-get :kernel (symbol-value var))
      (setf (alist-get :kernel (symbol-value var)) kernel)))
  (when (assoc lang org-babel-tangle-lang-exts)
    (add-to-list 'org-babel-tangle-lang-exts
                 (cons (concat "jupyter-" lang)
                       (cdr (assoc lang org-babel-tangle-lang-exts)))))
  (add-to-list 'org-src-lang-modes
               (cons (concat "jupyter-" lang)
                     (or (cdr (assoc lang org-src-lang-modes))
                         (intern (downcase (replace-regexp-in-string
                                            "[0-9]*" "" lang)))))))

(defun org-babel-jupyter-aliases-from-kernelspecs (&optional refresh specs)
  "Make language aliases based on the available kernelspecs.
For all kernel SPECS, make a language alias for the kernel
language if one does not already exist.  The alias is created with
`org-babel-jupyter-make-language-alias'.

SPECS defaults to `jupyter-available-kernelspecs'.  Optional
argument REFRESH has the same meaning as in
`jupyter-available-kernelspecs'.

Note, spaces or uppercase characters in the kernel language name
are converted into dashes or lowercase characters in the language
alias, e.g.

    Wolfram Language -> jupyter-wolfram-language"
  (cl-loop
   with specs = (or specs
                    (with-demoted-errors "Error retrieving kernelspecs: %S"
                      (jupyter-available-kernelspecs refresh)))
   for (kernel . (_dir . spec)) in specs
   for lang = (jupyter-canonicalize-language-string (plist-get spec :language))
   unless (member lang languages) collect lang into languages and
   do (org-babel-jupyter-make-language-alias kernel lang)))

;;; `ox' integration

(defvar org-latex-minted-langs)

(defun org-babel-jupyter-setup-export (backend)
  "Ensure that Jupyter src-blocks are integrated with BACKEND.
Currently this makes sure that Jupyter src-block languages are
mapped to their appropriate minted language in
`org-latex-minted-langs' if BACKEND is latex."
  (cond
   ((org-export-derived-backend-p backend 'latex)
    (cl-loop
     for (_kernel . (_dir . spec)) in (jupyter-available-kernelspecs)
     for lang = (plist-get spec :language)
     do (cl-pushnew (list (intern (concat "jupyter-" lang)) lang)
                    org-latex-minted-langs :test #'equal)))))

(defun org-babel-jupyter-strip-ansi-escapes (_backend)
  "Remove ANSI escapes from Jupyter src-block results in the current buffer."
  (org-babel-map-src-blocks nil
    (when (org-babel-jupyter-language-p lang)
      (when-let* ((pos (org-babel-where-is-src-block-result))
                  (ansi-color-apply-face-function
                   (lambda (beg end face)
                     ;; Could be useful for export backends
                     (when face
                       (put-text-property beg end 'face face)))))
        (goto-char pos)
        (ansi-color-apply-on-region (point) (org-babel-result-end))))))

;;; Hook into `org'

(org-babel-jupyter-aliases-from-kernelspecs)
(add-hook 'org-export-before-processing-hook #'org-babel-jupyter-setup-export)
(add-hook 'org-export-before-parsing-hook #'org-babel-jupyter-strip-ansi-escapes)

(provide 'ob-jupyter)

;;; ob-jupyter.el ends here