"""NVIDIA driver installation and configuration."""

import shutil
from pathlib import Path

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


def emerge_nvidia() -> None:
    """Install NVIDIA drivers."""
    info("=== Installing NVIDIA Drivers ===")

    emerge("x11-drivers/nvidia-drivers")

    success("NVIDIA drivers installed.")


def blacklist_nouveau() -> None:
    """Blacklist nouveau driver."""
    info("=== Blacklisting Nouveau ===")

    modprobe_dir = Path("/etc/modprobe.d")
    blacklist_conf = modprobe_dir / "blacklist-nouveau.conf"

    # Check if already configured (idempotency)
    if blacklist_conf.exists():
        info("Nouveau already blacklisted")
        return

    modprobe_dir.mkdir(parents=True, exist_ok=True)
    blacklist_conf.write_text(
        "# Blacklist nouveau in favor of proprietary nvidia driver\n"
        "blacklist nouveau\n"
        "options nouveau modeset=0\n"
    )

    success(f"Created {blacklist_conf}")


def configure_nvidia_drm() -> None:
    """Configure NVIDIA DRM modeset."""
    info("=== Configuring NVIDIA DRM ===")

    modprobe_dir = Path("/etc/modprobe.d")
    nvidia_conf = modprobe_dir / "nvidia.conf"

    # Check if already configured (idempotency)
    if nvidia_conf.exists():
        info("NVIDIA DRM already configured")
        return

    modprobe_dir.mkdir(parents=True, exist_ok=True)
    nvidia_conf.write_text(
        "# Enable NVIDIA DRM kernel mode setting\n"
        "options nvidia_drm modeset=1 fbdev=1\n"
    )

    success(f"Created {nvidia_conf}")


def copy_dracut_config(source_dir: Path) -> None:
    """Copy NVIDIA dracut configuration."""
    info("=== Copying Dracut NVIDIA Config ===")

    dracut_src = source_dir / "dracut.conf.d" / "nvidia.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 / "nvidia.conf")
        success("Copied nvidia.conf to /etc/dracut.conf.d/")
    else:
        # Create default if not in source
        nvidia_conf = dracut_dst / "nvidia.conf"
        nvidia_conf.write_text(
            "# NVIDIA early KMS\n"
            'add_drivers+=" nvidia nvidia_modeset nvidia_uvm nvidia_drm "\n'
            'force_drivers+=" nvidia nvidia_modeset nvidia_uvm nvidia_drm "\n'
        )
        success("Created default /etc/dracut.conf.d/nvidia.conf")


def rebuild_initramfs() -> None:
    """Rebuild initramfs to include NVIDIA modules."""
    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 verify_initramfs() -> None:
    """Verify NVIDIA modules are in initramfs."""
    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}...")

    # Check for NVIDIA modules
    result = run("lsinitrd", str(initramfs), capture=True, check=False)

    if result.ok and result.stdout:
        nvidia_modules = [
            line for line in result.stdout.split("\n")
            if "nvidia" in line.lower()
        ]

        if nvidia_modules:
            success("NVIDIA modules found in initramfs:")
            for module in nvidia_modules[:10]:  # Show first 10
                print(f"  {module}")
        else:
            error("WARNING: No NVIDIA modules found in initramfs!")
            print("  Initramfs may need to be rebuilt after kernel installation.")
    else:
        error("Could not inspect initramfs")


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

    emerge_nvidia()
    print()
    blacklist_nouveau()
    print()
    configure_nvidia_drm()
    print()
    copy_dracut_config(source_dir)
    print()

    # Ask about rebuilding initramfs
    print()
    info("Initramfs should be rebuilt to include NVIDIA modules.")
    response = prompt("Rebuild initramfs now? (y/n): ").strip().lower()

    if response == "y":
        rebuild_initramfs()
        print()
        verify_initramfs()

    print()
    success("=== NVIDIA Setup Complete ===")
    print()
    info("Next steps:")
    print("  1. Install and configure GRUB bootloader")
    print("  2. Verify initramfs has NVIDIA and crypt modules")
    print("  3. Reboot and test")