;;; Generic disk image builder for ruthless-guix
;;; Copyright © 2025
;;;
;;; This module provides a generic, reusable wrapper for building disk images
;;; that can work with different partition schemes (LVM, RAID, plain, etc.).

(define-module (ruthless image)
  #:use-module (gnu bootloader)
  #:use-module (gnu bootloader grub)
  #:use-module (gnu build marionette)
  #:use-module (gnu packages base)
  #:use-module (gnu packages bootloaders)
  #:use-module (gnu packages cryptsetup)
  #:use-module (gnu packages disk)
  #:use-module (gnu packages file-systems)
  #:use-module (gnu packages firmware)
  #:use-module (gnu packages gnupg)
  #:use-module (gnu packages guile)
  #:use-module (gnu packages linux)
  #:use-module (gnu packages virtualization)
  #:use-module (gnu services)
  #:use-module (gnu services base)
  #:use-module (gnu system)
  #:use-module (gnu system vm)
  #:use-module (gnu tests)
  #:use-module (gnu tests base)
  #:use-module (guix gexp)
  #:use-module (guix modules)
  #:use-module (guix monads)
  #:use-module (guix packages)
  #:use-module (guix records)
  #:use-module (guix store)
  #:use-module (guix scripts system reconfigure)
  #:use-module ((guix self) #:select (make-config.scm))
  #:use-module (srfi srfi-1)
  #:use-module (ice-9 match)
  #:export (<disk-initializer>
            disk-initializer
            disk-initializer?
            disk-initializer-name
            disk-initializer-packages
            disk-initializer-setup-proc
            disk-initializer-cleanup-proc

            build-disk-image

            make-lvm-disk-initializer
            lvm-disk-initializer
            make-simple-gpt-disk-initializer
            simple-gpt-disk-initializer))

;;;
;;; Disk initializer record
;;;

(define-record-type* <disk-initializer> disk-initializer
  make-disk-initializer
  disk-initializer?
  (name          disk-initializer-name)           ; symbol
  (packages      disk-initializer-packages        ; list of packages
                 (default '()))
  (setup-proc    disk-initializer-setup-proc)     ; gexp: sets up disk, mounts to /mnt
  (cleanup-proc  disk-initializer-cleanup-proc))  ; gexp: unmounts, deactivates

;;;
;;; Helper utilities
;;;

;; Guile-Gcrypt, Guile-SQLite3, and their propagated inputs.
(define gcrypt-sqlite3&co
  (append-map
   (lambda (package)
     (cons package
           (match (package-transitive-propagated-inputs package)
             (((labels packages) ...)
              packages))))
   (list guile-gcrypt guile-sqlite3)))

(define neither-config-nor-git?
  ;; Select (guix …) and (gnu …) modules, except (guix config) and (guix git).
  (match-lambda
    (('guix 'config) #f)
    (('guix 'git) #f)
    (('guix rest ...) #t)
    (('gnu rest ...) #t)
    (rest #f)))

(define-syntax-rule (with-imported-modules* gexp* ...)
  "Import necessary modules for image building operations."
  (with-extensions gcrypt-sqlite3&co
    (with-imported-modules `(,@(source-module-closure
                                '((gnu build image)
                                  (gnu build bootloader)
                                  (gnu build hurd-boot)
                                  (gnu build linux-boot)
                                  (guix store database))
                                #:select? neither-config-nor-git?)
                             ((guix config) => ,(make-config.scm)))
      #~(begin
          (use-modules (gnu build image)
                       (gnu build bootloader)
                       (gnu build hurd-boot)
                       (gnu build linux-boot)
                       (guix store database)
                       (guix build utils))
          gexp* ...))))

;;;
;;; Main disk image builder
;;;

(define* (build-disk-image os
                           #:key
                           disk-initializer
                           (disk-size (* 2048 1024 1024))  ; 2 GiB
                           (memory-size 1024)
                           (uefi? #t)
                           (image-format 'qcow2)
                           (name "disk-image"))
  "Return a derivation that builds a disk image with OS installed.

DISK-INITIALIZER is a <disk-initializer> record that defines:
  - Packages needed for disk setup
  - How to partition and format the disk
  - How to mount filesystems to /mnt
  - How to cleanup/unmount after initialization

The procedure handles:
  - Creating a builder VM with necessary packages
  - Creating the target disk
  - Calling the disk initializer's setup procedure
  - Running initialize-root-partition on the mounted filesystem
  - Installing the bootloader
  - Calling the cleanup procedure
  - Converting to the requested image format"

  (define builder-os
    ;; The helper OS that prepares the disk image.
    (marionette-operating-system
     (operating-system
       (inherit %simple-os)
       (packages
        (append (disk-initializer-packages disk-initializer)
                (list parted e2fsprogs dosfstools)
                (operating-system-packages %simple-os))))
     #:imported-modules '((gnu services herd)
                          (guix combinators))))

  (define builder-vm
    (virtual-machine
     (operating-system builder-os)
     (memory-size memory-size)
     (disk-image-size (* 1024 1024 1024))))  ; 1 GiB for builder

  (define bootcfg
    (operating-system-bootcfg os))

  (define bootloader-config
    (operating-system-bootloader os))

  (define bootloader
    (bootloader-configuration-bootloader bootloader-config))

  (define graph "system-graph")

  (define schema
    (local-file (search-path %load-path "guix/store/schema.sql")))

  ;; Computed file that contains the references graph for the OS closure.
  ;; This is needed because program-file doesn't support #:references-graphs.
  (define references-graph
    (computed-file
     "references-graph"
     #~(copy-file #$graph #$output)
     #:options `(#:references-graphs ((,graph ,os)))))

  ;; Program that initializes the root partition.
  ;; This will be loaded inside the VM via primitive-load.
  ;; Note: bootloader installation is handled separately after this step.
  (define initialize-root-program
    (program-file
     "initialize-root.scm"
     (with-imported-modules*
      (parameterize ((sql-schema #$schema))
        (initialize-root-partition
         "/mnt"
         #:references-graphs (list #$references-graph)
         #:copy-closures? #t
         #:register-closures? #t
         #:deduplicate? #f
         #:system-directory #$os
         #:bootcfg #$bootcfg
         #:bootcfg-location #$(bootloader-configuration-file bootloader)
         #:bootloader-package #$(bootloader-package bootloader)
         ;; Bootloader installer is called separately via install-bootloader
         #:bootloader-installer #f)))))

  (define setup-proc (disk-initializer-setup-proc disk-initializer))
  (define cleanup-proc (disk-initializer-cleanup-proc disk-initializer))

  ;; Main build procedure
  (define build
    (with-imported-modules (source-module-closure
                            '((guix build utils)
                              (gnu build marionette)))
      #~(begin
          (use-modules (guix build utils)
                       (gnu build marionette)
                       (ice-9 format))

          (define target-disk
            (string-append #$output ".raw"))

          ;; Create the target disk image
          (invoke #$(file-append qemu-minimal "/bin/qemu-img")
                  "create" "-f" "qcow2" target-disk
                  (number->string #$disk-size))

          ;; Boot the builder VM with UEFI if requested
          (let ((marionette
                 (make-marionette
                  (list #$builder-vm
                        #$@(if uefi?
                               #~("-bios" #$(file-append ovmf-x86-64
                                                         "/share/firmware/ovmf_x64.bin"))
                               #~())
                        "-drive"
                        (format #f "file=~a,if=virtio,cache=writeback,werror=report"
                                target-disk)))))

            ;; Wait for the system to be ready
            (unless (marionette-eval
                     '(begin
                        (use-modules (gnu services herd))
                        (start-service 'term-tty1)
                        #t)
                     marionette)
              (error "Failed to start term-tty1 service in VM"))

            ;; Step 1: Run disk setup (partition, format, mount to /mnt)
            (display "Setting up disk...\n")
            (unless (marionette-eval
                     '(begin
                        #$setup-proc
                        #t)
                     marionette)
              (error "Could not set up disks."))

            ;; Step 2: Initialize root partition (copy store, create structure, etc.)
            (display "Initializing root filesystem...\n")
            (unless (marionette-eval
                     '(begin
                        (primitive-load #$initialize-root-program)
                        #t)
                     marionette)
              (error "Failed to initialize root partition in VM"))

            ;; Step 3: Install bootloader
            (display "Installing bootloader...\n")
            (unless (#$(install-bootloader
                        (lambda (exp)
                          #~(lambda (marionette)
                              (marionette-eval '(begin #$exp #t) marionette)))
                        bootloader-config
                        bootcfg
                        #:target "/mnt")
                     marionette)
              (error "Failed to install bootloader"))

            ;; Step 4: Cleanup (unmount, deactivate LVM, etc.)
            (display "Cleaning up...\n")
            (unless (marionette-eval
                     '(begin
                        #$cleanup-proc
                        #t)
                     marionette)
              (error "Cleanup failed"))

            ;; Shutdown builder VM
            (display "Shutting down builder VM...\n")
            (marionette-eval
             '(system* "/run/current-system/profile/sbin/halt")
             marionette)

            ;; Wait for VM to terminate
            (false-if-exception
             (let loop ()
               (let ((status (waitpid (marionette-pid marionette) WNOHANG)))
                 (when (zero? (car status))
                   (sleep 1)
                   (loop))))))

          ;; Convert to final format or rename
          #$(match image-format
              ('qcow2
               #~(rename-file target-disk #$output))
              ('raw
               #~(begin
                   (invoke #$(file-append qemu-minimal "/bin/qemu-img")
                           "convert" "-f" "qcow2" "-O" "raw"
                           target-disk #$output)
                   (delete-file target-disk)))))))

  (computed-file (string-append name "." (symbol->string image-format))
                 build))

;;;
;;; Pre-built disk initializers
;;;

(define* (make-lvm-disk-initializer #:key
                                    (volume-group "vg0")
                                    (root-volume "root")
                                    (esp-size "259MiB")
                                    (esp-label "ESP")
                                    (root-label "root"))
  "Create a disk-initializer for LVM root with UEFI boot.

The partition layout is:
  - 1-3 MiB: BIOS boot partition (for legacy GRUB)
  - 3-ESP-SIZE: EFI System Partition
  - ESP-SIZE-100%: LVM physical volume containing VOLUME-GROUP/ROOT-VOLUME"
  (disk-initializer
   (name 'lvm-uefi)
   (packages (list lvm2))
   (setup-proc
    #~(begin
        (use-modules (guix build utils))

        ;; Partition: GPT with BIOS boot, ESP, and LVM PV
        (unless (zero?
                 (system* #$(file-append parted "/sbin/parted")
                          "-s" "/dev/vdb"
                          "mklabel" "gpt"
                          "mkpart" "bios_grub" "1MiB" "3MiB"
                          "set" "1" "bios_grub" "on"
                          "mkpart" "ESP" "fat32" "3MiB" #$esp-size
                          "set" "2" "esp" "on"
                          "mkpart" "lvm" #$esp-size "100%"
                          "set" "3" "lvm" "on"))
          (error "Failed to partition disk"))

        ;; Create ESP filesystem
        (unless (zero?
                 (system* #$(file-append dosfstools "/sbin/mkfs.fat")
                          "-F" "32" "-n" #$esp-label "/dev/vdb2"))
          (error "Failed to create ESP filesystem"))

        ;; Setup LVM
        (unless (zero?
                 (system* #$(file-append lvm2 "/sbin/pvcreate") "-ff" "-y" "/dev/vdb3"))
          (error "Failed to create LVM physical volume"))
        (unless (zero?
                 (system* #$(file-append lvm2 "/sbin/vgcreate") #$volume-group "/dev/vdb3"))
          (error "Failed to create LVM volume group"))
        (unless (zero?
                 (system* #$(file-append lvm2 "/sbin/lvcreate")
                          "-l" "100%FREE" "-n" #$root-volume #$volume-group))
          (error "Failed to create LVM root volume"))

        ;; Format root
        (unless (zero?
                 (system* #$(file-append e2fsprogs "/sbin/mkfs.ext4")
                          "-L" #$root-label
                          #$(string-append "/dev/mapper/" volume-group "-" root-volume)))
          (error "Failed to format root filesystem"))

        ;; Mount filesystems
        (mkdir-p "/mnt")
        (unless (zero?
                 (system* "mount"
                          #$(string-append "/dev/mapper/" volume-group "-" root-volume)
                          "/mnt"))
          (error "Failed to mount root filesystem"))
        (mkdir-p "/mnt/boot/efi")
        (unless (zero? (system* "mount" "/dev/vdb2" "/mnt/boot/efi"))
          (error "Failed to mount ESP"))))

   (cleanup-proc
    #~(begin
        (sync)
        (unless (zero? (system* "umount" "/mnt/boot/efi"))
          (error "Failed to unmount ESP"))
        (unless (zero? (system* "umount" "/mnt"))
          (error "Failed to unmount root"))
        (unless (zero? (system* #$(file-append lvm2 "/sbin/vgchange") "-an" #$volume-group))
          (error "Failed to deactivate LVM volume group"))))))

(define lvm-disk-initializer
  (make-lvm-disk-initializer))

(define* (make-simple-gpt-disk-initializer #:key
                                           (esp-size "259MiB")
                                           (esp-label "ESP")
                                           (root-label "root"))
  "Create a disk-initializer for simple GPT with UEFI boot (no LVM).

The partition layout is:
  - 1-3 MiB: BIOS boot partition (for legacy GRUB)
  - 3-ESP-SIZE: EFI System Partition
  - ESP-SIZE-100%: Root partition (ext4)"
  (disk-initializer
   (name 'simple-gpt-uefi)
   (packages '())
   (setup-proc
    #~(begin
        (use-modules (guix build utils))

        ;; Partition: GPT with BIOS boot, ESP, and root
        (unless (zero?
                 (system* #$(file-append parted "/sbin/parted")
                          "-s" "/dev/vdb"
                          "mklabel" "gpt"
                          "mkpart" "bios_grub" "1MiB" "3MiB"
                          "set" "1" "bios_grub" "on"
                          "mkpart" "ESP" "fat32" "3MiB" #$esp-size
                          "set" "2" "esp" "on"
                          "mkpart" "root" "ext4" #$esp-size "100%"))
          (error "Failed to partition disk"))

        ;; Create filesystems
        (unless (zero?
                 (system* #$(file-append dosfstools "/sbin/mkfs.fat")
                          "-F" "32" "-n" #$esp-label "/dev/vdb2"))
          (error "Failed to create ESP filesystem"))
        (unless (zero?
                 (system* #$(file-append e2fsprogs "/sbin/mkfs.ext4")
                          "-L" #$root-label "/dev/vdb3"))
          (error "Failed to format root filesystem"))

        ;; Mount filesystems
        (mkdir-p "/mnt")
        (unless (zero? (system* "mount" "/dev/vdb3" "/mnt"))
          (error "Failed to mount root filesystem"))
        (mkdir-p "/mnt/boot/efi")
        (unless (zero? (system* "mount" "/dev/vdb2" "/mnt/boot/efi"))
          (error "Failed to mount ESP"))))

   (cleanup-proc
    #~(begin
        (sync)
        (unless (zero? (system* "umount" "/mnt/boot/efi"))
          (error "Failed to unmount ESP"))
        (unless (zero? (system* "umount" "/mnt"))
          (error "Failed to unmount root"))))))

(define simple-gpt-disk-initializer
  (make-simple-gpt-disk-initializer))