(define-module (nbde system kexec-initrd)
  #:use-module (gnu packages bash)
  #:use-module (gnu packages base)
  #:use-module (gnu packages compression)
  #:use-module (gnu packages linux)
  #:use-module ((gnu packages make-bootstrap)
                #:select (%guile-static-initrd))
  #:use-module (gnu system file-systems)
  #:use-module (gnu system keyboard)
  #:use-module (gnu system linux-initrd)
  #:use-module (guix gexp)
  #:use-module (guix modules)
  #:export (kexec-installer-initrd))

(define flat-linux-module-directory*
  (@@ (gnu system linux-initrd) flat-linux-module-directory))

(define* (expression->initrd-with-proc exp
                                       #:key
                                       (guile %guile-static-initrd)
                                       (name "guile-initrd"))
  (define real-init
    (program-file "real-init" exp #:guile guile))

  (define init-wrapper
    (computed-file
     "init"
     #~(begin
         (call-with-output-file #$output
           (lambda (port)
             (display "#!" port)
             (display #$(file-append static-bash "/bin/sh") port)
             (newline port)
             (display #$(file-append util-linux "/bin/mount") port)
             (display " -t proc proc /proc\n" port)
             (display "exec " port)
             (display #$(file-append guile "/bin/guile") port)
             (display " --no-auto-compile " port)
             (display #$real-init port)
             (newline port)))
         (chmod #$output #o555))))

  (define (import-module? module)
    (and (guix-module-name? module)
         (not (equal? module '(guix store deduplication)))))

  (define builder
    (with-imported-modules
        (source-module-closure '((gnu build linux-initrd))
                               #:select? import-module?)
      #~(begin
          (use-modules (gnu build linux-initrd)
                       (guix build store-copy)
                       (guix build utils)
                       (system base compile)
                       (system foreign)
                       (rnrs bytevectors))

          (define (cache-compiled-file-name file)
            (format #f ".cache/guile/ccache/~a-~a-~a-~a/~a"
                    (effective-version)
                    (if (eq? (native-endianness) (endianness little))
                        "LE"
                        "BE")
                    (sizeof '*)
                    (effective-version)
                    file))

          (define (compile-to-cache file)
            (let ((compiled-file (cache-compiled-file-name file)))
              (mkdir-p (dirname compiled-file))
              (compile-file file
                            #:opts %auto-compilation-options
                            #:output-file compiled-file)))

          (mkdir #$output)
          (mkdir "contents")
          (populate-store '("closure") "contents"
                          #:deduplicate? #f)

          (with-directory-excursion "contents"
            (symlink #$init-wrapper "init")
            (symlink #$real-init "real-init")
            (compile-to-cache "real-init")

            ;; Allow Guile to find out where it is.
            (mkdir-p "proc/self")
            (symlink (string-append #$guile "/bin/guile") "proc/self/exe")
            (readlink "proc/self/exe")

            (write-cpio-archive (string-append #$output "/initrd.cpio.gz")
                                "."
                                #:gzip (string-append #+gzip "/bin/gzip"))))
      ))

  (file-append
   (computed-file name builder
                  #:options
                  `(#:references-graphs (("closure" ,init-wrapper))))
   "/initrd.cpio.gz"))

(define* (kexec-installer-initrd file-systems
                                 #:key
                                 (linux linux-libre)
                                 (linux-modules '())
                                 (mapped-devices '())
                                 keyboard-layout
                                 (helper-packages (list util-linux))
                                 qemu-networking?
                                 volatile-root?
                                 (on-error 'backtrace)
                                 #:allow-other-keys)
  "Return an initrd for kexec-style live systems whose root file system is a
fresh tmpfs.  The initrd bind-mounts its own /gnu tree into the mounted root
before switching to it, so 'gnu.load' remains reachable after 'switch-root'."
  (define kodir
    (flat-linux-module-directory* linux linux-modules))

  (define helper-packages*
    (append helper-packages
            (if keyboard-layout
                (list loadkeys-static)
                '())))

  (expression->initrd-with-proc
   (with-imported-modules
       (source-module-closure
        '((gnu build linux-boot)
          (gnu build linux-modules)
          (gnu build file-systems)
          (gnu system file-systems)
          (guix build syscalls)
          (guix build utils)
          (guix build bournish)
          (system repl error-handling)
          (system repl repl)))
     #~(begin
         (use-modules (system repl error-handling)
                      (system repl repl)
                      (srfi srfi-1)
                      (srfi srfi-13)
                      (ice-9 match)
                      (ice-9 popen)
                      (ice-9 rdelim)
                      (gnu build linux-boot)
                      (gnu build linux-modules)
                      (gnu build file-systems)
                      (gnu system file-systems)
                      ((guix build syscalls)
                       #:hide (file-system-type))
                      ((guix build utils) #:hide (delete))
                      (guix build bournish))

         (with-output-to-port (%make-void-port "w")
           (lambda ()
             (set-path-environment-variable "PATH" '("bin" "sbin")
                                            '#$helper-packages*)))

         (define mount-root-file-system*
           (@@ (gnu build linux-boot) mount-root-file-system))

         (define move-essential-file-systems*
           (@@ (gnu build linux-boot) move-essential-file-systems))

         (define (device-string->file-system-device device-string)
           (cond ((string-prefix? "/" device-string) device-string)
                 ((string-contains device-string ":/") device-string)
                 ((uuid device-string) => identity)
                 (else (file-system-label device-string))))

         (define (root-mount-point? fs)
           (string=? (file-system-mount-point fs) "/"))

         (define (populate-root-from-initrd-paths paths)
           (for-each
            (lambda (path)
              (when (file-exists? path)
                (let ((target (string-append "/root" path)))
                  (if (directory-exists? path)
                      (begin
                        (mkdir-p target)
                        ;; Make PATH its own mount point first so we can move it
                        ;; under /root and keep it alive across 'switch-root'.
                        (mount path path "" MS_BIND)
                        (mount path target "" MS_MOVE))
                      (begin
                        (mkdir-p (dirname target))
                        (copy-file path target)))
                  (format #t "staged '~a' at '~a'~%" path target))))
            paths))

         (define (dump-file path)
           (when (file-exists? path)
             (format #t "--- ~a ---~%" path)
             (call-with-input-file path
               (lambda (port)
                 (let loop ()
                   (let ((line (read-line port 'concat)))
                     (unless (eof-object? line)
                       (display line)
                       (newline)
                       (loop))))))
             (format #t "--- end ~a ---~%" path)))

         (define (log-path-state label path)
           (format #t "~a: ~a exists? ~a dir? ~a~%"
                   label path
                   (file-exists? path)
                   (directory-exists? path)))

         (define (move-initrd-directory-under-root path)
           (let ((target (string-append "/root" path)))
             (mkdir-p target)
             ;; Make PATH its own mount point so nested mounts move with it.
             (mount path path "" MS_BIND)
             (mount path target "" MS_MOVE)
             (format #t "moved initrd mount '~a' to '~a'~%" path target)))

         (define (attach-loop-device image)
           (let* ((pipe   (open-pipe* OPEN_READ "losetup" "--find" "--show"
                                      "--read-only" image))
                  (device (string-trim-right (read-line pipe))))
             (close-pipe pipe)
             (unless (and (string? device)
                          (string-prefix? "/dev/loop" device))
               (error "failed to attach loop device" image device))
             (format #t "attached loop device '~a' for '~a'~%" device image)
             device))

         (define (enter-root root)
           ;; For the kexec installer we only need a new logical root, not a
           ;; full initrd teardown.  Avoiding 'switch-root' sidesteps the
           ;; disappearing-/gnu problem while still giving the boot program a
           ;; coherent "/" view.
           (display "before enter-root\n")
           (dump-file "/proc/self/mountinfo")
           (move-essential-file-systems* root)
           (chdir root)
           (chroot ".")
           (chdir "/")
           (display "after enter-root\n")
           (dump-file "/proc/self/mountinfo"))

         (define* (boot-system-with-initrd-binds #:key
                                                 (linux-modules '())
                                                 linux-module-directory
                                                 keymap-file
                                                 qemu-guest-networking?
                                                 volatile-root?
                                                 (mounts '())
                                                 (squashfs-store-image
                                                  "/gnu-store.squashfs")
                                                 (initrd-bind-mounts
                                                  '("/etc/guix-kexec"))
                                                 (on-error 'debug))
           (display "Welcome, this is GNU's early boot Guile.\n")
           (display "Use 'gnu.repl' for an initrd REPL.\n\n")

           (call-with-error-handling
             (lambda ()
               (mount-essential-file-systems)
               (let* ((args    (linux-command-line))
                      (to-load (find-long-option "gnu.load" args))
                      (root-device (and=> (find-long-option "root" args)
                                          device-string->file-system-device))
                      (rootfstype  (find-long-option "rootfstype" args))
                      (rootflags   (find-long-option "rootflags" args))
                      (root-fs*    (find root-mount-point? mounts))
                      (root-fs
                       (if root-fs*
                           (file-system
                             (inherit root-fs*)
                             (device (or root-device
                                         (file-system-device root-fs*)))
                             (type (or rootfstype
                                       (file-system-type root-fs*)))
                             (options (or rootflags
                                          (file-system-options root-fs*))))
                           (file-system
                             (device root-device)
                             (mount-point "/")
                             (type rootfstype)
                             (options rootflags)))))

                 (unless (or root-fs* (and root-device rootfstype))
                   (error "no root file system or 'root' and 'rootfstype' parameters"))

                 (when (member "gnu.repl" args)
                   (start-repl))

                 (display "loading kernel modules...\n")
                 (load-linux-modules-from-directory linux-modules
                                                    linux-module-directory)

                 (when keymap-file
                   (let ((status (system* "loadkeys" keymap-file)))
                     (unless (zero? status)
                       (format (current-error-port)
                               "warning: 'loadkeys' exited with status ~a~%"
                               status))))

                 (when qemu-guest-networking?
                   (unless (configure-qemu-networking)
                     (display "network interface is DOWN\n")))

                 (unless (file-exists? "/root")
                   (mkdir "/root"))

                 (setenv "EXT2FS_NO_MTAB_OK" "1")

                 (mount-root-file-system* (canonicalize-device-spec
                                           (file-system-device root-fs))
                                          (file-system-type root-fs)
                                          #:volatile-root? volatile-root?
                                          #:flags (mount-flags->bit-mask
                                                   (file-system-flags root-fs))
                                          #:options (file-system-options root-fs)
                                          #:check? #f
                                          #:skip-check-if-clean? #f
                                          #:repair #f)

                 (for-each (lambda (fs)
                             (mount-file-system fs
                                                #:check? #f
                                                #:skip-check-if-clean? #f
                                                #:repair #f))
                           (remove root-mount-point? mounts))

                 (setenv "EXT2FS_NO_MTAB_OK" #f)

                 ;; Mount the squashfs store image (embedded in initrd at
                 ;; /gnu-store.squashfs) directly under /root so the mounts
                 ;; remain visible after we chroot into the target tmpfs root.
                 ;; Items inside are at the squashfs root, stripped of the
                 ;; /gnu/store/ prefix.
                 (when (and squashfs-store-image
                            (file-exists? squashfs-store-image))
                   (display "mounting squashfs store image...\n")
                   (let ((loop-device (attach-loop-device
                                       squashfs-store-image)))
                     (mkdir-p "/root/gnu/.ro-store")
                     (mount loop-device "/root/gnu/.ro-store" "squashfs"
                            MS_RDONLY "")
                     (mkdir-p "/root/gnu/store")
                     (mkdir-p "/root/gnu/.rw-store/store")
                     (mkdir-p "/root/gnu/.rw-store/work")
                     (mount "overlay" "/root/gnu/store" "overlay" 0
                            (string-append
                             "lowerdir=/root/gnu/.ro-store"
                             ",upperdir=/root/gnu/.rw-store/store"
                             ",workdir=/root/gnu/.rw-store/work"))
                     (log-path-state "pre-root" "/root/gnu")
                     (log-path-state "pre-root" "/root/gnu/store")))

                 (populate-root-from-initrd-paths initrd-bind-mounts)

                 (if to-load
                     (begin
                       (enter-root "/root")
                       (format #t "post-root: /gnu exists? ~a, /gnu/store exists? ~a, boot exists? ~a~%"
                               (file-exists? "/gnu")
                               (file-exists? "/gnu/store")
                               (file-exists? to-load))
                       (format #t "loading '~a'...\n" to-load)
                       (primitive-load to-load)
                       (format (current-error-port)
                               "boot program '~a' terminated, rebooting~%"
                               to-load)
                       (sleep 2)
                       (reboot))
                     (begin
                       (display "no boot file passed via 'gnu.load'\n")
                       (display "entering a warm and cozy REPL\n")
                       (start-repl)))))
             #:on-error on-error))

         (boot-system-with-initrd-binds
          #:mounts
          (map spec->file-system
               '#$(map file-system->spec file-systems))
          #:linux-modules '#$linux-modules
          #:linux-module-directory '#$kodir
          #:keymap-file #+(and=> keyboard-layout
                                 keyboard-layout->console-keymap)
          #:qemu-guest-networking? #$qemu-networking?
          #:volatile-root? '#$volatile-root?
          #:on-error '#$on-error)))
   #:name "kexec-installer-initrd"))