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

let inherit (lib) mkOption types;
in {
  options = {
    enableCuda = mkOption {
      type = types.bool;
      default = false;
      description = ''
        Enable CUDA support for Blender and other packages.
        Warning: This can cause very long build times.
      '';
    };
    nvidiaOpen = mkOption {
      type = types.bool;
      default = false;
    };
  };

  config = {
    hardware.graphics.enable = true;

    services.xserver.videoDrivers = [ "nvidia" ];

    hardware.nvidia = {
      modesetting.enable = true;
      powerManagement.enable = false;
      open = config.nvidiaOpen;
      nvidiaSettings = true;
      package = config.boot.kernelPackages.nvidiaPackages.stable;
    };

    # Apply CUDA-specific settings if enabled
    nixpkgs.config.cudaSupport = config.enableCuda;
    hardware.enableAllFirmware = config.enableCuda;

    environment.systemPackages = (if config.enableCuda then
      [ (pkgs.blender.override { cudaSupport = true; }) ]
    else
      [ ]);
  };
}