;;; Guix-kernel -- Guix kernel for Jupyter
;;; Copyright (C) 2018 Evgeny Panfilov <epanfilov@gmail.com>
;;; Copyright (C) 2018 Pierre-Antoine Rouby <pierre-antoine.rouby@inria.fr>
;;; Copyright (C) 2018, 2019 Inria
;;;
;;; 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/>.

(use-modules (json)
             (simple-zmq)
	     (srfi srfi-1)
             (srfi srfi-11)
             (srfi srfi-13)
             (srfi srfi-19)
             (rnrs bytevectors)
             (ice-9 vlist)
             (ice-9 futures)
             (ice-9 match)
             (sxml simple)
             (gnu build linux-container)
             (gnu system file-systems)
             (guix build syscalls)
             (guix build utils)
             (jupyter messages)
             (jupyter kernels)
             (jupyter servers)
             (guix-kernel magic)
             (guix-kernel environ))

(define session-id (random (* 255 255)
                           (seed->random-state
                            (time-second (current-time time-utc)))))

;;
;; Container tools.
;;

(define container-context (zmq-create-context))

(define container-path
  (string-replace (car (command-line)) "guix-jupyter-container.scm"
                  (+ (string-rindex (car (command-line)) #\/) 1)))

(define (register-proxy vhash name proxy)
  "Register PROXY, a kernel object for a process running in a container that
acts as a proxy to an actual kernel."
  (vhash-cons name proxy vhash))

(define (proxy-by-name vhash name)
  (match (vhash-assoc name vhash)
    ((_ . proxy) proxy)
    (#f #f)))


;;
;; Html.
;;

(define (error->shtml error)
  (match error
    (('guix-kernel msg errno)
     `(h3 (@ (style "color: red;"))
          ,(string-append "Error: " msg)))
    (_
     `(h3 (@ (style "color: red;"))
          "Error !"))))

(define (error->html error)
  (call-with-output-string
    (lambda (port)
      (sxml->xml (error->shtml error) port))))

;;
;; Handlers.
;;

(define (general-handler kernel kind message state)
  "Handle MESSAGE, which was sent by KERNEL, a client of ours, on KIND (one
of 'kernel-shell', 'kernel-stdin', etc.)"
  (match state
    ((containers count)
     (let ((handler (dispatch (message-type message))))
       (pub-busy kernel message)
       (let-values (((containers count)
                     (handler kernel message containers count)))
         (pub-idle kernel message)
         (list containers count))))))

;; Unknown request type, ignore it.
(define (ignore-request kernel message containers count)
  (values containers count))

(define KERNEL-INFO
  `(("protocol_version" . "5.3.0")
    ("implementation"   . "Guix Jupyter kernel")
    ("implementation_version" . "0.0.2")
    ("language_info" .
     (("name"     . "guile")
      ("version"  . ,(effective-version))
      ("mimetype" . "application/x-scheme")
      ("file_extension"  . ".scm")
      ("pygments_lexer"  . "scheme")
      ("codemirror_mode" . "scheme")))
    ("banner"     . "Guix kernel")
    ("help_links" .
     (("Gitlab Inria" .
       "https://gitlab.inria.fr/guix-hpc/guix-kernel")))))

;; Send kernel-info.
(define (reply-kernel-info-request kernel message containers count)
  (send-message kernel (reply message "kernel_info_reply"
                              (scm->json-string KERNEL-INFO)))
  (values containers count))

(define (reply-execute-request kernel message containers count)
  (let* ((content (message-content message))
         (code    (assoc-ref (json-string->scm content) "code"))
         (magic   (get-magic-line code)))
    (catch 'guix-kernel
      (λ ()
        (cond
         ((magic-env? code)
          (let* ((list     (string-split magic #\ ))
                 (env-name (list-ref list 2))
                 (env      (list-cdr-ref list 4)))
            (match (proxy-by-name containers env-name)
              (#f
               (let* ((id        (start-container kernel env-name env))
                      (container (new-container-connect env-name id))
                      (new       (register-proxy containers env-name
                                                 container)))
                 (relay-message (kernel-shell container) kernel
                                (string->utf8 env-name)
                                message)
                 (values new (+ count 1))))
              ((? kernel? proxy)
               (relay-message (kernel-shell proxy) kernel
                              (string->utf8 env-name)
                              message)
               (values containers (+ count 1))))))
         ((or (magic-run? code)
              (magic-kernel? code))
          (let* ((list     (string-split magic #\ ))
                 (env-name (list-ref list 2))
                 (proxy    (proxy-by-name containers env-name)))
            (relay-message (kernel-shell proxy) kernel
                           (string->utf8 env-name)
                           message)
            (values containers (+ count 1))))
         ((magic-html? code)
          (values containers
                  (reply-html kernel message
                              (delete-magic code)
                              count)))
         (else
          (relay-message (kernel-shell
                          (proxy-by-name containers "default"))
                         kernel (string->utf8 "default")
                         message)
          (values containers (+ count 1)))))
      (λ error
        (values containers
                (reply-html kernel message
                            (error->html error) count))))))

(define (shutdown kernel message containers count)
  (kill-containers containers kernel message)
  (exit #t))

;;
;; Dispatch route.
;;

(define dispatch-route
  `(("kernel_info_request" . ,reply-kernel-info-request)
    ("execute_request"     . ,reply-execute-request)
    ("shutdown_request"    . ,shutdown)
    ("comm_info_request"   . ,ignore-request)))

(define (dispatch msg-type)
  (let ((res (assoc-ref dispatch-route msg-type)))
    (unless res
      (display
       (string-append "\n(WW) unknown message type: "
                      msg-type "\n\n")))
    (if res res ignore-request)))

;;
;; Process exit and signals handler.
;;

(define (kill-containers containers kernel message)
  "Terminate all the proxies listed in the CONTAINERS vhash, sending them
MESSAGE."
  (vlist-for-each (match-lambda
                    ((name . proxy)
                     (format (current-error-port)
                             "terminating proxy ~s (PID ~s)...~%"
                             name (kernel-pid proxy))
                     (relay-message (kernel-shell proxy) kernel
                                    (string->utf8 name) message)
                     (zmq-close-socket (kernel-shell proxy))
                     (false-if-exception (kill (kernel-pid proxy) SIGTERM))))
                  containers))

;;
;; Run.
;;

(define (new-container-connect name pid)
  "Return a kernel object for process PID."
  (let* ((ipc-dir  (string-append "/tmp/guix-kernel/ipc/" name "-"
                                  (number->string session-id) "/"))
         (container-addr (string-append "ipc://" ipc-dir))
         (container-addr-sock (string-append container-addr "shell")))

    (mkdir-p ipc-dir)

    (let ((container-socket (zmq-create-socket container-context ZMQ_DEALER)))
      (zmq-set-socket-option container-socket ZMQ_IDENTITY name)

      (zmq-bind-socket container-socket container-addr-sock)
      (display "+")

      (kernel name pid
              #:shell container-socket
              #:iosub container-socket))))

(define %network-file-systems
  ;; The list of <file-system> objects corresponding to bind-mounts required
  ;; for networking.
  (filter-map (lambda (mapping)
                (let ((fs (file-system-mapping->bind-mount mapping)))
                  (and (file-exists? (file-system-device fs))
                       fs)))
              %network-file-mappings))

(define (start-container kernel name env)
  (let* ((guile-version (if (null? env)
                            (guile-current-version->path)
                            (if (not (package-in-list->path env "guile"))
                                (guile-current-version->path)
                                (guile->bin-path
                                 (package-in-list->path env "guile")))))
         (new-env (if (null? env)
                      (environ)
                      (make-new-environment (environ) name env)))
         (exec (lambda ()
                 (set-network-interface-up "lo") ;up lo interface

                 (apply execle guile-version
                        new-env
                        (append  (list "guile")
                                 (guile-current-load-path->args-list)
                                 (guile-current-load-compiled-path->args-list)
                                 (list "--no-auto-compile" "-s"
                                       container-path
                                       name (number->string session-id)
                                       (utf8->string (kernel-key kernel)))))))
         (root (string-append "/tmp/guix-kernel/container/"
                              name "-" (number->string session-id)))
         (fs   (cons* %immutable-store
                      (file-system
                        (device "/tmp/guix-kernel")
                        (mount-point "/tmp/guix-kernel")
                        (type "none")
                        (check? #f)
                        (flags '(bind-mount)))
                      (append %container-file-systems
                              %network-file-systems))))
    (mkdir-p root)
    (run-container root fs
                   %namespaces 1
                   exec)))


(define (exit-handler kernel)
  (lambda _
    (close-kernel kernel)
    (exit 1)))

;; Start!
(let ((kernel (call-with-input-file (car (last-pair (command-line)))
                connection-file->kernel)))
  (sigaction SIGTERM (exit-handler kernel))
  (sigaction SIGINT  (exit-handler kernel))

  (let* ((pid        (start-container kernel "default" '()))
         (containers (vhash-cons "default"
                                 (new-container-connect "default" pid)
                                 vlist-null)))
   (serve-kernels (list kernel)
                  general-handler
                  (list containers 0))))