gentoo-legion-python/install/bootloader.py

"""GRUB bootloader installation and configuration."""

import re
import shutil
from pathlib import Path

from .utils import info, success, error, warn, run, emerge


def emerge_grub() -> None:
    """Install GRUB bootloader."""
    info("=== Installing GRUB ===")

    emerge("sys-boot/grub")

    success("GRUB installed.")


def install_grub_efi() -> None:
    """Install GRUB to EFI system partition."""
    info("=== Installing GRUB to EFI ===")

    run(
        "grub-install",
        "--target=x86_64-efi",
        "--efi-directory=/boot",
        "--bootloader-id=Gentoo",
    )

    success("GRUB installed to EFI partition.")


def configure_grub() -> None:
    """Configure GRUB for LUKS + NVIDIA."""
    info("=== Configuring GRUB ===")

    grub_default = Path("/etc/default/grub")

    if not grub_default.exists():
        error("/etc/default/grub not found")
        return

    content = grub_default.read_text()
    changes_made = False

    # Settings to apply (use regex for robust matching of commented/uncommented lines)
    settings = {
        "GRUB_CMDLINE_LINUX_DEFAULT": '"nvidia_drm.modeset=1 acpi_backlight=native"',
        "GRUB_ENABLE_CRYPTODISK": "y",
        "GRUB_GFXMODE": "1920x1080",
    }

    for key, value in settings.items():
        full_setting = f'{key}={value}'
        pattern = rf'^#?\s*{key}=.*$'

        if re.search(pattern, content, re.MULTILINE):
            # Replace existing line (commented or not)
            new_content = re.sub(pattern, full_setting, content, flags=re.MULTILINE)
            if new_content != content:
                content = new_content
                changes_made = True
        else:
            # Append if not found at all
            content += f"\n{full_setting}"
            changes_made = True

    if changes_made:
        grub_default.write_text(content)
        success("GRUB configured for LUKS + NVIDIA + HiDPI")
        print("  - GRUB_CMDLINE_LINUX_DEFAULT: nvidia_drm.modeset=1 acpi_backlight=native")
        print("  - GRUB_ENABLE_CRYPTODISK: y")
        print("  - GRUB_GFXMODE: 1920x1080")
    else:
        info("GRUB already configured")


def generate_grub_config() -> None:
    """Generate GRUB configuration."""
    info("=== Generating GRUB Config ===")

    run("grub-mkconfig", "-o", "/boot/grub/grub.cfg")

    success("GRUB config generated.")


def copy_dracut_crypt_config(source_dir: Path) -> None:
    """Copy LUKS dracut configuration."""
    info("=== Copying Dracut Crypt Config ===")

    dracut_src = source_dir / "dracut.conf.d" / "crypt.conf"
    dracut_dst = Path("/etc/dracut.conf.d")
    dracut_dst.mkdir(parents=True, exist_ok=True)

    if dracut_src.exists():
        shutil.copy2(dracut_src, dracut_dst / "crypt.conf")
        success("Copied crypt.conf to /etc/dracut.conf.d/")
    else:
        # Create default if not in source
        crypt_conf = dracut_dst / "crypt.conf"
        crypt_conf.write_text(
            "# LUKS support for encrypted root\n"
            'add_dracutmodules+=" crypt "\n'
        )
        success("Created default /etc/dracut.conf.d/crypt.conf")


def verify_initramfs() -> None:
    """Verify initramfs has required modules for boot."""
    info("=== Verifying Initramfs ===")

    # Find latest initramfs
    boot = Path("/boot")
    initramfs_files = sorted(boot.glob("initramfs-*.img"), reverse=True)

    if not initramfs_files:
        error("No initramfs found in /boot")
        return

    initramfs = initramfs_files[0]
    info(f"Checking {initramfs.name}...")

    result = run("lsinitrd", str(initramfs), capture=True, check=False)

    if not result.ok or not result.stdout:
        error("Could not inspect initramfs")
        return

    output = result.stdout

    # Check for crypt/LUKS modules
    print()
    info("LUKS/Crypt modules:")
    crypt_found = False
    for keyword in ["crypt", "luks", "dm-crypt"]:
        matches = [line for line in output.split("\n") if keyword in line.lower()]
        if matches:
            crypt_found = True
            for match in matches[:5]:
                print(f"  {match}")

    if crypt_found:
        success("LUKS support found in initramfs")
    else:
        error("WARNING: No LUKS/crypt modules found!")
        print("  Boot from encrypted root may fail.")

    # Check for NVIDIA modules
    print()
    info("NVIDIA modules:")
    nvidia_found = False
    nvidia_modules = [line for line in output.split("\n") if "nvidia" in line.lower()]

    if nvidia_modules:
        nvidia_found = True
        for module in nvidia_modules[:5]:
            print(f"  {module}")
        success("NVIDIA modules found in initramfs")
    else:
        warn("No NVIDIA modules in initramfs (may be OK if not using early KMS)")

    # Summary
    print()
    if crypt_found and nvidia_found:
        success("Initramfs verification PASSED - LUKS and NVIDIA modules present")
    elif crypt_found:
        success("Initramfs verification PASSED - system should boot (NVIDIA optional)")
    else:
        error("Initramfs verification FAILED - fix before rebooting!")


def rebuild_initramfs() -> None:
    """Rebuild initramfs with current configuration."""
    info("=== Rebuilding Initramfs ===")

    # Find installed kernel version
    modules_dir = Path("/lib/modules")
    if not modules_dir.exists():
        error("No kernel modules found. Install kernel first.")
        return

    kernels = sorted(modules_dir.iterdir(), reverse=True)
    if not kernels:
        error("No kernel versions found in /lib/modules")
        return

    kernel_version = kernels[0].name
    info(f"Rebuilding initramfs for kernel {kernel_version}")

    run("dracut", "--force", f"/boot/initramfs-{kernel_version}.img", kernel_version)

    success("Initramfs rebuilt.")


def setup_bootloader(source_dir: Path | None = None) -> None:
    """Full bootloader setup workflow."""
    if source_dir is None:
        source_dir = Path("/root/gentoo")

    # Ensure dracut crypt config is in place
    copy_dracut_crypt_config(source_dir)
    print()

    emerge_grub()
    print()
    install_grub_efi()
    print()
    configure_grub()
    print()
    generate_grub_config()
    print()

    # Verify initramfs
    verify_initramfs()

    print()
    success("=== Bootloader Setup Complete ===")
    print()
    info("Pre-reboot checklist:")
    print("  1. Verify /etc/fstab is correct")
    print("  2. Verify /etc/conf.d/dmcrypt has swap configured")
    print("  3. Ensure root password is set")
    print("  4. Ensure user account exists")
    print()
    info("Ready to reboot!")
    print("  exit        # Leave chroot")
    print("  reboot      # Restart system")