# SPDX-FileCopyrightText: 2025 Ethan Reece <contact@ethanreece.com>
#
# SPDX-License-Identifier: MIT
#
# Uses unlicensed snippets from https://github.com/nix-community/disko/issues/550 meant for public use.

{ pkgs, ... }:
let
  # SPDX-SnippetBegin
  # SPDX-SnippetCopyrightText: 2024 MatthewCroughan <https://github.com/nix-community/disko/issues/550#issuecomment-2503736973>
  configTxt = pkgs.writeText "config.txt" ''
    [pi4]
    kernel=u-boot-rpi4.bin
    enable_gic=1

    # Otherwise the resolution will be weird in most cases, compared to
    # what the pi3 firmware does by default.
    disable_overscan=1

    # Supported in newer board revisions
    arm_boost=1

    [cm4]
    # Enable host mode on the 2711 built-in XHCI USB controller.
    # This line should be removed if the legacy DWC2 controller is required
    # (e.g. for USB device mode) or if USB support is not required.
    otg_mode=1

    [all]
    # Boot in 64-bit mode.
    arm_64bit=1

    # U-Boot needs this to work, regardless of whether UART is actually used or not.
    # Look in arch/arm/mach-bcm283x/Kconfig in the U-Boot tree to see if this is still
    # a requirement in the future.
    enable_uart=1

    # Prevent the firmware from smashing the framebuffer setup done by the mainline kernel
    # when attempting to show low-voltage or overtemperature warnings.
    avoid_warnings=1
  '';
  # SPDX-SnippetEnd
  device = "/dev/sda";
  labels = {
    boot = "NIXOS_BOOT";
    firmware = "NIXOS_FW";
    luks = "NIXOS_LUKS";
    root = "NIXOS_ROOT";
  };
in
{
  boot.postBootCommands = ''
    # On the first boot, resize the disk
    if [ -f /persistent/disko-first-boot ]; then
      set -euo pipefail
      set -x
      # Figure out device names for the boot device and root filesystem.
      rootPart=$(${pkgs.util-linux}/bin/findmnt -n -o SOURCE /)
      cryptDevice=$(lsblk -npo PKNAME "$rootPart")
      bootDevice=$(lsblk -npo PKNAME "/dev/mapper/$cryptDevice")
      partNum=$(lsblk -npo MAJ:MIN "/dev/mapper/$cryptDevice" | ${pkgs.gawk}/bin/awk -F: '{print $2}')

      # Resize partition to fill disk
      echo ",+," | sfdisk -N$partNum --no-reread "$bootDevice"
      ${pkgs.parted}/bin/partprobe

      # Resize cryptographic layer
      ${pkgs.cryptsetup}/bin/cryptsetup resize "$cryptDevice"

      # Resize physical volume
      ${pkgs.lvm2}/bin/pvresize "/dev/mapper/$cryptDevice"

      # Resize logical volume
      ${pkgs.lvm2}/bin/lvextend -l +100%FREE "$rootPart"

      # Resize Btrfs filesystem
      ${pkgs.btrfs-progs}/bin/btrfs filesystem resize max /

      # Prevents this from running on later boots.
      rm -f /persistent/disko-first-boot
    fi
  '';
  disko = {
    # SPDX-SnippetBegin
    # SPDX-License-Identifier: LicenseRef-Unspecified
    # SPDX-SnippetCopyrightText: 2024 MatthewCroughan <https://github.com/nix-community/disko/issues/550#issuecomment-2503736973>
    imageBuilder = {
      kernelPackages = pkgs.linuxPackages_latest;
      extraPostVM = ''
        ${pkgs.zstd}/bin/zstd --compress $out/*raw
        rm $out/*raw
      '';
    };
    devices = {
      disk = {
        disk1 = {
          inherit device;
          imageSize = "20G";
          type = "disk";
          postCreateHook = ''
            lsblk
            sgdisk -A 1:set:2 ${device}
          '';
          content = {
            type = "gpt";
            partitions = {
              firmware = {
                size = "30M";
                priority = 1;
                type = "0700";
                label = labels.firmware;
                content = {
                  type = "filesystem";
                  format = "vfat";
                  mountpoint = "/firmware";
                  postMountHook = toString (
                    pkgs.writeScript "postMountHook.sh" ''
                      (cd ${pkgs.raspberrypifw}/share/raspberrypi/boot && cp bootcode.bin fixup*.dat start*.elf *.dtb /mnt/firmware/)
                      cp ${pkgs.ubootRaspberryPi4_64bit}/u-boot.bin /mnt/firmware/u-boot-rpi4.bin
                      cp ${configTxt} /mnt/firmware/config.txt
                    ''
                  );
                };
              };
              # SPDX-SnippetEnd
              boot = {
                size = "1G";
                type = "EF00";
                label = labels.boot;
                content = {
                  type = "filesystem";
                  format = "vfat";
                  mountpoint = "/boot";
                };
              };
              luks = {
                size = "100%";
                content = {
                  types = "luks";
                  name = "encrypted";
                  label = labels.luks;
                  settings = {
                    allowDiscards = true;
                  };
                  content = {
                    type = "btrfs";
                    extraArgs = [ "-f" ];
                    subvolumes = {
                      "root" = {
                        mountpoint = "/";
                        mountOptions = [ "compress=zstd" ];
                      };
                      "persistent" = {
                        mountpoint = "/persistent";
                        mountOptions = [ "compress=zstd" ];
                      };
                      "nix" = {
                        mountpoint = "/nix";
                        mountOptions = [ "compress=zstd" ];
                      };
                    };
                    postMountHook = toString (
                      pkgs.writeScript "postMountHook.sh" ''
                        touch /mnt/persistent/disko-first-boot
                      ''
                    );
                  };
                };
              };
            };
          };
        };
      };
    };
  };
}