# A drop-in NixOS module that reads the nix-facter hardware report and
# automatically configures the correct NVIDIA driver.
#
# Supported driver branches (newest → oldest):
#   stable      Pascal GP1xx / Turing TU1xx / Ampere GA1xx /
#               Ada Lovelace AD1xx / Blackwell GB1xx / future
#   legacy_470  Maxwell GM2xx  (GeForce GTX 9xx)
#   legacy_390  Fermi GF1xx + Kepler GK1xx/GK110 (GTX 4xx–7xx)
#   legacy_340  Tesla G8x/G9x/GT21x              (GeForce 8/9/GT2xx)
{
  config,
  lib,
  ...
}:
let
  cfg = config.hardware.facter.detected.graphics.nvidia;

  # ---------------------------------------------------------------------------
  # Hex parser  (pure Nix, no builtins.fromTOML / lib.toHexInt required)
  # ---------------------------------------------------------------------------
  hexCharVal =
    c:
    if c == "0" then
      0
    else if c == "1" then
      1
    else if c == "2" then
      2
    else if c == "3" then
      3
    else if c == "4" then
      4
    else if c == "5" then
      5
    else if c == "6" then
      6
    else if c == "7" then
      7
    else if c == "8" then
      8
    else if c == "9" then
      9
    else if c == "a" then
      10
    else if c == "b" then
      11
    else if c == "c" then
      12
    else if c == "d" then
      13
    else if c == "e" then
      14
    else if c == "f" then
      15
    else
      0;

  fromHex =
    s:
    let
      chars = lib.stringToCharacters (lib.toLower (lib.removePrefix "0x" s));
    in
    builtins.foldl' (acc: c: acc * 16 + hexCharVal c) 0 chars;

  # ---------------------------------------------------------------------------
  # PCI device extraction from the facter report
  # ---------------------------------------------------------------------------
  # facter.json structure:
  #   { "hardware": { "pci": [ { "vendor": "0x10de",
  #                               "device": "0x2204",
  #                               "class":  "0x030000" } ] } }
  #
  # nixos-facter-modules exposes the parsed JSON as:
  #   config.hardware.facter.report  (attrset)

  facterReport = config.hardware.facter.report;

  pciDevices = facterReport.hardware.pci or [ ];

  # NVIDIA vendor = 0x10de
  # Display classes: 0x030000 VGA, 0x030200 3D controller, etc.
  isNvidiaDisplay =
    device:
    lib.toLower (device.vendor or "") == "0x10de"
    && lib.hasPrefix "0x03" (lib.toLower (device.class or ""));

  nvidiaGpus = builtins.filter isNvidiaDisplay pciDevices;

  # ---------------------------------------------------------------------------
  # Driver-family tables
  # ---------------------------------------------------------------------------
  # Format: list of { min : int; max : int } ranges (device-ID in decimal).
  # A device ID falling in any range of a table → that family.
  # Nothing matches → "stable" (latest, covers all current & future GPUs).
  #
  # Reference: https://www.nvidia.com/en-us/drivers/unix/legacy-gpu/
  # Device-ID ranges are derived from the per-branch README/supportedchips.html

  inRanges = idInt: ranges: builtins.any (r: idInt >= r.min && idInt <= r.max) ranges;

  h = fromHex; # shorthand

  # ── legacy_340  (G8x · G9x · GT21x) ─────────────────────────────────────
  ranges340 = [
    {
      min = h "0190";
      max = h "019F";
    } # GeForce 8800 GTS (G80)
    {
      min = h "01D0";
      max = h "01DF";
    } # GeForce 7900 GS / 7950 GT
    {
      min = h "0390";
      max = h "039F";
    } # GeForce 7900 GX2
    {
      min = h "0400";
      max = h "042F";
    } # GeForce 8xxx G92/G94
    {
      min = h "0600";
      max = h "064F";
    } # GeForce 9xxx G94/G96
    {
      min = h "0750";
      max = h "077F";
    } # GeForce 9xxx G96/G98
    {
      min = h "0840";
      max = h "087F";
    } # GeForce 9xxx M / ION
    {
      min = h "08A0";
      max = h "08AF";
    } # GeForce 9xxx M
    {
      min = h "0A20";
      max = h "0A6F";
    } # GeForce GT 2xx (GT215–GT218)
    {
      min = h "0CA0";
      max = h "0CAF";
    } # GeForce GT 215
  ];

  # ── legacy_390  (Fermi GF1xx · Kepler GK104/GK106/GK107/GK110) ──────────
  ranges390 = [
    {
      min = h "0DC0";
      max = h "0DFF";
    } # Fermi GF106/GF104
    {
      min = h "0E20";
      max = h "0E3F";
    } # Fermi GF104 mobile
    {
      min = h "0E80";
      max = h "0EBF";
    } # Fermi GF108/GF104 mobile
    {
      min = h "1040";
      max = h "107F";
    } # Kepler GK104
    {
      min = h "1080";
      max = h "10BF";
    } # Kepler GK106 / GK107
    {
      min = h "1140";
      max = h "117F";
    } # Kepler GK107M
    {
      min = h "1180";
      max = h "11BF";
    } # Kepler GK110 / GK110B
    {
      min = h "11C0";
      max = h "11FF";
    } # Kepler GK106 desktop
    {
      min = h "1200";
      max = h "121F";
    } # Kepler GK104 (GTX 680)
    {
      min = h "0FC0";
      max = h "0FFF";
    } # Kepler GK107 (GT 640/650)
  ];

  # ── legacy_470  (Maxwell GM2xx · late Kepler GK208) ──────────────────────
  ranges470 = [
    {
      min = h "1280";
      max = h "12BF";
    } # Kepler GK208 (GT 720/730/740)
    {
      min = h "13B0";
      max = h "13FF";
    } # Maxwell GM204 (GTX 970/980)
    {
      min = h "1401";
      max = h "143F";
    } # Maxwell GM206 (GTX 960/950)
    {
      min = h "1617";
      max = h "163F";
    } # Maxwell GM204/GM206 mobile
    {
      min = h "1710";
      max = h "171F";
    } # Maxwell GM204 mobile high-end
    {
      min = h "17C0";
      max = h "17FF";
    } # Maxwell GM200 (Titan X / GTX 980 Ti)
    {
      min = h "1B00";
      max = h "1B0F";
    } # Maxwell GM200 NVS/Quadro
  ];

  # Pascal (0x1B80+) and all future arches → "stable" (no range needed).

  # ---------------------------------------------------------------------------
  # Map a single PCI device-ID string to a driver family
  # ---------------------------------------------------------------------------
  familyForId =
    idStr:
    let
      idInt = fromHex idStr;
    in
    if inRanges idInt ranges340 then
      "legacy_340"
    else if inRanges idInt ranges390 then
      "legacy_390"
    else if inRanges idInt ranges470 then
      "legacy_470"
    else
      "stable"; # Pascal, Turing, Ampere, Ada, Blackwell, future

  # ---------------------------------------------------------------------------
  # Determine the single best driver for the machine
  # ---------------------------------------------------------------------------
  # When multiple GPUs are present (e.g., dGPU + iGPU bypass, or mixed legacy
  # setup), we choose the *newest* driver family — i.e., the one closest to
  # "stable" — because:
  #   1. The stable driver cannot drive a 340-class card, but a 340-class card
  #      typically lives on a machine that also has a modern dGPU only in rare
  #      cases (mixed system).  Users with true mixed legacy systems should
  #      override `hardware.nvidia.package` manually.
  #   2. Unknown / future device IDs always resolve to "stable", which is safe.
  familyOrder = [
    "stable"
    "legacy_470"
    "legacy_390"
    "legacy_340"
  ];

  gpuFamilies = map (dev: familyForId (dev.device or "0x0000")) nvidiaGpus;

  bestFamily =
    let
      present = lib.unique gpuFamilies;
      winner = lib.findFirst (f: builtins.elem f present) null familyOrder;
    in
    if winner == null then "stable" else winner;

  hasNvidiaGpu = nvidiaGpus != [ ];

  selectedPackage =
    config.boot.kernelPackages.nvidiaPackages.${bestFamily}
      or config.boot.kernelPackages.nvidiaPackages.stable;
in
{
  # ---------------------------------------------------------------------------
  # Module options
  # ---------------------------------------------------------------------------
  options.hardware.facter.detected.graphics.nvidia = {
    enable = lib.mkOption {
      type = lib.types.bool;
      default = hasNvidiaGpu;
      description = ''
        Enable the automatic configuration of the NVIDIA driver
        based on the nix-facter hardware report.  Set to false to manage
        hardware.nvidia manually.
      '';
    };
  };

  # ---------------------------------------------------------------------------
  # Module configuration
  # ---------------------------------------------------------------------------
  config = lib.mkIf (config.hardware.facter.enable && cfg.enable && hasNvidiaGpu) {
    services.xserver.videoDrivers = [ "nvidia" ];

    hardware.graphics.enable = lib.mkDefault true;

    hardware.nvidia = {
      # The auto-selected package; users can still override with mkForce.
      package = lib.mkDefault selectedPackage;

      # Open kernel modules only work on Turing and newer (driver ≥ 560).
      # For legacy families they are unavailable, so we hard-disable them.
      open = lib.mkDefault (bestFamily == "stable");

      modesetting.enable = lib.mkDefault true;
      nvidiaSettings = lib.mkDefault true;
    };
  };
}