;;; my-async-tests.el --- Scratch/temporary file: some tests about async -*- lexical-binding: t -*-

;; Copyright (C) 2024 Bruno BARBIER

;; Author: Bruno BARBIER
;; Status: Temporary tests.
;; Compatibility: GNU Emacs 30.0.50
;;
;; 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 2 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, write to the Free
;; Software Foundation, Inc., 59 Temple Place, Suite 330, Boston,
;; MA 02111-1307 USA

(require 'cl-lib)
(require 'org)
(require 'org-elib-async)

(require 'ob-shell)
(require 'ob-python)

;;; Shells
;;

;;;; One shell script

;; Standalone direct asynchronous execution.

(defun my-shell-babel-schedule (lang body _params handle-feedback)
  "Execute the bash script BODY.
Execute the shell script BODY using bash.  Use HANDLE-FEEDBACK to report
the outcome (success or failure)."
  (unless (equal "bash" lang)
    (error "Only for bash"))
  (funcall handle-feedback (list :pending "started"))
  (org-elib-async-process (list "bash") :input body :callback handle-feedback))


;;;; Asynchronous using ob-shell

(defun my-org-babel-shell-how-to-execute (body params)
  "Return how to execute BODY using a POSIX shell.
Return how to execute, as expected by
`org-elib-async-comint-queue--execution'."
  ;; Code mostly extracted from ob-shell, following
  ;; `org-babel-execute:shell' and `org-babel-sh-evaluate'.
  ;; Results are expected to differ from ob-shell as we follow the
  ;; same process for all execution paths: asynchronous or not, with
  ;; session or without.
  (let* ((session (org-babel-sh-initiate-session
		   (cdr (assq :session params))))
	 (stdin (let ((stdin (cdr (assq :stdin params))))
                  (when stdin (org-babel-sh-var-to-string
                               (org-babel-ref-resolve stdin)))))
	 (result-params (cdr (assq :result-params params)))
	 (value-is-exit-status
	  (or (and
	       (equal '("replace") result-params)
	       (not org-babel-shell-results-defaults-to-output))
	      (member "value" result-params)))
	 (cmdline (cdr (assq :cmdline params)))
         (shebang (cdr (assq :shebang params)))
         (full-body (concat
		     (org-babel-expand-body:generic
		      body params (org-babel-variable-assignments:shell params))
		     (when value-is-exit-status "\necho $?")))
         (post-process
          (lambda (r)
            (setq r (org-trim r))
            (org-babel-reassemble-table
             (org-babel-result-cond result-params
               r
               (let ((tmp-file (org-babel-temp-file "sh-")))
                 (with-temp-file tmp-file (insert r))
                 (org-babel-import-elisp-from-file tmp-file)))
             (org-babel-pick-name
              (cdr (assq :colname-names params)) (cdr (assq :colnames params)))
             (org-babel-pick-name
              (cdr (assq :rowname-names params)) (cdr (assq :rownames params))))))
         comint-buffer
         finally
         to-run)

    (setq comint-buffer
          (if session session
            ;; No session. We create a temporary one and use 'finally' to
            ;; destroy it once we are done.
            ;;
            ;; FIXME: This session code should be refactored and moved into
            ;;        ob-core.
            (let ((s-buf (org-babel-sh-initiate-session
                          (generate-new-buffer-name (format "*ob-shell-no-session*")))))
              (setq finally (lambda ()
                                ;; We cannot delete it immediately as we are called from it.
                                (run-with-idle-timer
                                 0.1 nil
                                 (lambda ()
                                   (when (buffer-live-p s-buf)
                                     (let ((kill-buffer-query-functions nil)
                                           (kill-buffer-hook nil))
                                       (kill-buffer s-buf)))))))
              s-buf)))

    (org-elib-async-comint-queue-init-if-needed comint-buffer)

    (setq to-run
          (cond
           ((or stdin cmdline)	       ; external shell script w/STDIN
            (let ((script-file (org-babel-temp-file "sh-script-"))
	          (stdin-file (org-babel-temp-file "sh-stdin-"))
	          (padline (not (string= "no" (cdr (assq :padline params))))))
	      (with-temp-file script-file
	        (when shebang (insert shebang "\n"))
	        (when padline (insert "\n"))
	        (insert full-body))
	      (set-file-modes script-file #o755)
	      (with-temp-file stdin-file (insert (or stdin "")))
	      (with-temp-buffer
                (with-connection-local-variables
                 (concat
                  (mapconcat #'shell-quote-argument
                             (cons (if shebang (file-local-name script-file)
                                     shell-file-name)
                                   (if shebang (when cmdline (list cmdline))
                                     (list shell-command-switch
                                           (concat (file-local-name script-file)  " " cmdline))))
                             " ")
                  "<" (shell-quote-argument stdin-file))))))
           (session ; session evaluation
            full-body)
           ;; External shell script, with or without a predefined
           ;; shebang.
           ((org-string-nw-p shebang)
            (let ((script-file (org-babel-temp-file "sh-script-"))
	          (padline (not (equal "no" (cdr (assq :padline params))))))
	      (with-temp-file script-file
	        (insert shebang "\n")
	        (when padline (insert "\n"))
	        (insert full-body))
	      (set-file-modes script-file #o755)
              (if (file-remote-p script-file)
                  ;; Run remote script using its local path as COMMAND.
                  ;; The remote execution is ensured by setting
                  ;; correct `default-directory'.
                  (let ((default-directory (file-name-directory script-file)))
                    (file-local-name script-file)
	            script-file ""))))
           (t
            (let ((script-file (org-babel-temp-file "sh-script-")))
	      (with-temp-file script-file
	        (insert full-body))
	      (set-file-modes script-file #o755)
              (mapconcat #'shell-quote-argument
                         (list shell-file-name
                               shell-command-switch
                               (if (file-remote-p script-file)
                                   (file-local-name script-file)
                                 script-file))
                         " ")))))
    ;; TODO: How to handle `value-is-exit-status'?
    (lambda (&rest q)
      (pcase q
        (`(:instrs-to-enter)
         ;; FIXME: This is wrong.
         "export PS1=''; export PS2='';")
        (`(:instrs-to-exit))
        (`(:finally) (when finally (funcall finally)))
        (`(:instr-to-emit-tag ,tag) (format "printf '%s\\n'" tag))
        (`(:post-process ,r) (when post-process (funcall post-process r)))
        (`(:send-instrs-to-session ,code)
         (with-current-buffer comint-buffer
           (when code
             (goto-char (point-max))
             (insert code) (insert "\n")
             (comint-send-input nil t))))
        (`(:get-code) to-run)
        (`(:get-comint-buffer) comint-buffer)
        (_ (error "Unknown query"))))))





;;; Python
;;

;;;; Asynchronous using ob-python

(defun my-org-babel-python-how-to-execute (body params)
  "Return how to execute BODY using python.
Return how to execute, as expected by
`org-elib-async-comint-queue--execution'."
  ;; Code mostly extracted from ob-python, following
  ;; `org-babel-python-evaluate-session'.
  ;; Results are expected to differ from ob-python as we follow the
  ;; same process for all execution paths: asynchronous or not, with
  ;; session or without.
  (let* ((org-babel-python-command
          (or (cdr (assq :python params))
              org-babel-python-command))
         (session-key (org-babel-python-initiate-session
                       (cdr (assq :session params))))
         (graphics-file (and (member "graphics" (assq :result-params params))
                             (org-babel-graphical-output-file params)))
         (result-params (cdr (assq :result-params params)))
         (result-type (cdr (assq :result-type params)))
         (results-file (when (eq 'value result-type)
                         (or graphics-file
                             (org-babel-temp-file "python-"))))
         (return-val (when (eq result-type 'value)
                       (cdr (assq :return params))))
         (full-body
          (concat
           (org-babel-expand-body:generic
            body params
            (org-babel-variable-assignments:python params))
           (when return-val
             (format "\n%s" return-val))))
         (post-process
          (lambda (r)
            (setq r (string-trim r))
            (when (string-prefix-p "Traceback (most recent call last):" r)
              (signal 'user-error (list r)))
            (when (eq 'value result-type)
              (setq r (org-babel-eval-read-file results-file)))
            (org-babel-reassemble-table
             (org-babel-result-cond result-params
               r
               (org-babel-python-table-or-string r))
             (org-babel-pick-name (cdr (assq :colname-names params))
                                  (cdr (assq :colnames params)))
             (org-babel-pick-name (cdr (assq :rowname-names params))
                                  (cdr (assq :rownames params))))))
         (tmp-src-file (org-babel-temp-file "python-"))
         (session-body
          ;; The real code we evaluate in the session.
          (pcase result-type
            (`output
             (format (string-join
                      (list "with open('%s') as f:\n"
                            "    exec(compile(f.read(), f.name, 'exec'))\n"))
                     (org-babel-process-file-name
                      tmp-src-file 'noquote)))
            (`value
             ;; FIXME: In this case, any output is an error.
             (org-babel-python-format-session-value
              tmp-src-file results-file result-params))))
         comint-buffer
         finally)


    (unless session-key
      ;; No session. We create a temporary one and use 'finally' to
      ;; destroy it once we are done.
      ;;
      ;; FIXME: This session code should be refactored and moved into
      ;;        ob-core.
      (setq session-key (org-babel-python-initiate-session
                         ;; We can't use a simple `generate-new-buffer'
                         ;; due to the earmuffs game.
                         (org-babel-python-without-earmuffs
                          (format "*ob-python-no-session-%s*" (org-id-uuid)))))
      (setq finally (lambda ()
                      (when-let ((s-buf
                                  (get-buffer (org-babel-python-with-earmuffs session-key))))
                        ;; We cannot delete it immediately as we are called from it.
                        (run-with-idle-timer
                         0.1 nil
                         (lambda ()
                           (when (buffer-live-p s-buf)
                             (let ((kill-buffer-query-functions nil)
                                   (kill-buffer-hook nil))
                               (kill-buffer s-buf)))))))))

    (setq comint-buffer
          (get-buffer (org-babel-python-with-earmuffs session-key)))
    (org-elib-async-comint-queue-init-if-needed comint-buffer)
    (with-temp-file tmp-src-file
      (insert (if (and graphics-file (eq result-type 'output))
                  (format org-babel-python--output-graphics-wrapper
                          full-body graphics-file)
                full-body)))

    (lambda (&rest q)
      (pcase q
        (`(:instrs-to-enter)
         ;; FIXME: This is wrong.
         "import sys; sys.ps1=''; sys.ps2=''")
        (`(:instrs-to-exit))
        (`(:finally) (when finally (funcall finally)))
        (`(:instr-to-emit-tag ,tag) (format "print ('%s')" tag))
        (`(:post-process ,r) (when post-process (funcall post-process r)))
        (`(:send-instrs-to-session ,code)
         ;; See org-babel-python-send-string
         (with-current-buffer comint-buffer
           (let ((python-shell-buffer-name
                  (org-babel-python-without-earmuffs session-key)))
             (python-shell-send-string (concat code "\n")))))
        (`(:get-code) session-body)
        (`(:get-comint-buffer) comint-buffer)
        (_ (error "Unknown query"))))))


;;; Org babel 'execute-with'.


;;;; Asynchronous
;;
(defun my-org-babel-schedule (lang body params handle-feedback)
  "Schedule the execution of BODY according to PARAMS.
This function is called by `org-babel-execute-src-block'.  Return a
function that waits and returns the result on success, raise on failure."
  (let ((exec (pcase lang
                ("python" (my-org-babel-python-how-to-execute body params))
                ("bash"   (my-org-babel-shell-how-to-execute body params))
                (_ (error "Not handled (yet): %s" lang)))))
    (org-elib-async-comint-queue--push exec :handle-feedback handle-feedback)))


;;;; Synchronous
;;
(defun my-org-babel-execute (lang body params)
  "Execute Python BODY according to PARAMS.
This function is called by `org-babel-execute-src-block'."
  ;; We just start the asynchronous execution, wait for it, and return
  ;; the result (or raise the exception). No custom code, and,
  ;; synchronous and asynchronous should just mix nicely together.
  (funcall (my-org-babel-schedule lang body params nil)))