{ lib, config, pkgs, ... }:

with lib;

let
  cfg = config.roles.nvidia;
in
{
  options.roles.nvidia = {
    enable = mkEnableOption "Enable the nvidia role";

    # Driver configuration options
    open = mkOption {
      type = types.bool;
      default = true;
      description = "Use the open source nvidia kernel driver (for Turing and newer GPUs).";
    };

    modesetting = mkOption {
      type = types.bool;
      default = true;
      description = "Enable kernel modesetting for nvidia.";
    };

    nvidiaSettings = mkOption {
      type = types.bool;
      default = true;
      description = "Enable the nvidia-settings GUI.";
    };

    package = mkOption {
      type = types.enum [ "stable" "latest" "beta" "vulkan_beta" "production" ];
      default = "stable";
      description = "The nvidia driver package to use.";
    };

    powerManagement = {
      enable = mkOption {
        type = types.bool;
        default = false;
        description = "Enable nvidia power management (useful for laptops, not recommended for desktops).";
      };

      finegrained = mkOption {
        type = types.bool;
        default = false;
        description = "Enable fine-grained power management for Turing and newer GPUs.";
      };
    };

    graphics = {
      enable = mkOption {
        type = types.bool;
        default = true;
        description = "Enable hardware graphics support.";
      };

      enable32Bit = mkOption {
        type = types.bool;
        default = false;
        description = "Enable 32-bit graphics libraries (needed for some games).";
      };

      extraPackages = mkOption {
        type = types.listOf types.package;
        default = [];
        description = "Extra packages to add to hardware.graphics.extraPackages.";
      };
    };
  };

  config = mkIf cfg.enable {
    # Set xserver video driver
    services.xserver.videoDrivers = [ "nvidia" ];

    # Graphics configuration
    hardware.graphics = {
      enable = cfg.graphics.enable;
      enable32Bit = cfg.graphics.enable32Bit;
      extraPackages = cfg.graphics.extraPackages;
    };

    # NVIDIA driver configuration
    hardware.nvidia = {
      modesetting.enable = cfg.modesetting;
      nvidiaSettings = cfg.nvidiaSettings;
      open = cfg.open;
      package = config.boot.kernelPackages.nvidiaPackages.${cfg.package};
      powerManagement.enable = cfg.powerManagement.enable;
      powerManagement.finegrained = cfg.powerManagement.finegrained;
    };

    # Additional packages for nvidia support
    environment.systemPackages = with pkgs; [
      libva-utils
      nvidia-vaapi-driver
      nvtopPackages.nvidia
    ];
  };
}