ps4nixos/nixos/default.nix

{ lib, config, pkgs, modulesPath, ... }: {
  options.ps4 = with lib; {
    enable = mkEnableOption "Build NixOS for Sony PlayStation 4.";

    usbImage = {
      imageName = mkOption {
        type = types.str;
        default = "${config.ps4.usbImage.imageBaseName}-${config.system.nixos.label}.img";
        description = "Name of the generated image file.";
      };
      
      imageBaseName = mkOption {
        type = types.str;
        default = "ps4nixos";
        description = "Prefix of the name of the generated image file.";
      };
    
      compressImage = mkOption {
        type = types.bool;
        default = false;
        description = "Compress the resulting image with zstd.";
      };

      # TODO this is unnecessary. just make the partition big enough to fit the boot files automatically.
      bootPartitionSize = mkOption {
        type = types.int;
        default = 64;
        description = "Size of the boot partition, in mebibytes (1024x1024 bytes).";
      };

      bootPartitionLabel = mkOption {
        type = types.str;
        default = "nixos-boot";
        description = "Label to give the boot partition of the image.";
      };

      rootPartitionLabel = mkOption {
        type = types.str;
        default = "nixos-root";
        description = "Label to give the root partition of the image.";
      };

      expandOnBoot = mkOption {
        type = types.bool;
        default = true;
        description = "Whether the image should expand the root partition on first boot.";
      };

      firstbootFile = mkOption {
        type = types.str;
        default = "/firstboot";
        description = ''
          Location of the file that allows commands to be run on first boot.
          If overriding fileSystems."/" then you should to set this to the root mount + /firstboot
        '';
      };
    };
  };

  config = let cfg = config.ps4; in lib.mkIf cfg.enable {
    # TODO write definitions for the other southbridges
    boot.kernelPackages = pkgs.recurseIntoAttrs (pkgs.linuxPackagesFor (pkgs.callPackage ../kernel/kernel-aeolia-5.3.nix {}));
  
    fileSystems = {
      "/boot" = {
        device = "/dev/disk/by-label/${cfg.usbImage.bootPartitionLabel}";
        fsType = "vfat";
        options = [ "nofail" "noauto" ];
      };
      "/" = {
        device = "/dev/disk/by-label/${cfg.usbImage.rootPartitionLabel}";
        fsType = "ext4";
      };
    };

    boot.loader.grub.enable = lib.mkForce false;
    boot.loader.initScript.enable = lib.mkForce true;

    # from https://github.com/NixOS/nixpkgs/blob/e405f30513169feedb64b5c25e7b00242010af58/nixos/modules/installer/sd-card/sd-image.nix#L267
    boot.postBootCommands = let
      expandOnBoot = lib.optionalString cfg.usbImage.expandOnBoot ''
        # Figure out device names for the boot device and root filesystem.
        rootPart=$(${pkgs.util-linux}/bin/findmnt -n -o SOURCE /)
        bootDevice=$(lsblk -npo PKNAME $rootPart)
        partNum=$(lsblk -npo MAJ:MIN $rootPart | ${pkgs.gawk}/bin/awk -F: '{print $2}')

        # Resize the root partition and the filesystem to fit the disk
        echo ",+," | sfdisk -N$partNum --no-reread $bootDevice
        ${pkgs.parted}/bin/partprobe
        ${pkgs.e2fsprogs}/bin/resize2fs $rootPart
      '';
      inherit (cfg.usbImage) firstbootFile;
    in ''
      # On the first boot do some maintenance tasks
      if [ -f ${firstbootFile} ]; then
        set -euo pipefail
        set -x

        ${expandOnBoot}

        # TODO what does all this do?

        # Register the contents of the initial Nix store
        # ${config.nix.package.out}/bin/nix-store --load-db < ${firstbootFile}

        # nixos-rebuild also requires a "system" profile and an /etc/NIXOS tag.
        # touch /etc/NIXOS
        ${config.nix.package.out}/bin/nix-env -p /nix/var/nix/profiles/system --set /run/current-system

        # Prevents this from running on later boots.
        rm -f ${firstbootFile}
      fi
    '';
    
    system.build.rootfsImage = pkgs.callPackage (modulesPath + "/../lib/make-ext4-fs.nix") {
      storePaths = [ config.system.build.toplevel ];
      volumeLabel = cfg.usbImage.rootPartitionLabel;
      populateImageCommands = ''
        cp ${config.system.build.toplevel}/init ./files/init
        touch ./files/firstboot
      '';
    };

    # slightly modified version of https://github.com/NixOS/nixpkgs/blob/e405f30513169feedb64b5c25e7b00242010af58/nixos/modules/installer/sd-card/sd-image.nix#L182
    system.build.usbImage = pkgs.stdenv.mkDerivation {
      name = cfg.usbImage.imageName;

      nativeBuildInputs = with pkgs; [ dosfstools parted mtools e2fsprogs util-linux ]
        ++ lib.optional cfg.usbImage.compressImage zstd;

      inherit (cfg.usbImage) compressImage;

      buildCommand = ''
        # mkdir -p $out/usb-image
        img=$out
        rootfs=${config.system.build.rootfsImage}

        # Create the image to fit rootfs + bootfs + 1 MiB
        rootfsSizeBlocks=$(du -B 512 --apparent-size $rootfs | awk '{ print $1 }')
        bootfsSizeBlocks=$((${toString cfg.usbImage.bootPartitionSize} * 1024 * 1024 / 512))
        imageSize=$((rootfsSizeBlocks * 512 + bootfsSizeBlocks * 512 + 1 * 1024 * 1024))
        truncate -s $imageSize $img

        # Partition the image
        parted -s $img -- mklabel msdos
        parted -s $img -- mkpart primary fat32 1M ${toString (cfg.usbImage.bootPartitionSize + 1)}M
        parted -s $img -- mkpart primary ext4 ${toString (cfg.usbImage.bootPartitionSize + 1)}M 100%

        # Copy the rootfs to the image
        eval $(partx $img -o START,SECTORS --nr 2 --pairs)
        dd conv=notrunc if=$rootfs of=$img seek=$START count=$SECTORS

        # Create a temporary image for the boot partition
        eval $(partx $img -o START,SECTORS --nr 1 --pairs)
        truncate -s $((SECTORS * 512)) boot.img
        mkfs.vfat -F 32 -n ${cfg.usbImage.bootPartitionLabel} boot.img

        # Copy the files that the PS4 Linux Loader payload expects to the boot partition
        mcopy -i boot.img ${config.system.build.kernel}/bzImage ::/bzImage
        mcopy -i boot.img ${config.system.build.initialRamdisk}/initrd.gz ::/initramfs.cpio.gz

        # Verify the boot partition before copying it
        fsck.vfat -vn boot.img
        dd conv=notrunc if=boot.img of=$img seek=$START count=$SECTORS

        if test -n "$compressImage"; then
          zstd -T$NIX_BUILD_CORES --rm $img
        fi
      '';
    };
  };
}