# llm-nixos: VM 200 on lovey. 2x RTX 3090 via VFIO, model weights on virtiofs /models.
{ config, lib, pkgs, ... }:

{
  imports = [ ./hardware-configuration.nix ./gpu-unlock.nix ];

  boot.loader.systemd-boot.enable = true;
  boot.loader.efi.canTouchEfiVariables = true;
  boot.loader.systemd-boot.configurationLimit = 5;  # initrds carry ~84 MB of gpu-unlock firmware
  boot.kernelPackages = pkgs.linuxPackages;  # LTS-ish default, keeps the NVIDIA module happy

  networking.hostName = "llm-nixos";
  # Static: lovey-fans only accepts GPU temps from this address
  networking.useDHCP = false;
  networking.interfaces.ens18.ipv4.addresses = [ { address = "192.168.1.67"; prefixLength = 24; } ];
  networking.defaultGateway = { address = "192.168.1.1"; interface = "ens18"; };
  networking.nameservers = [ "1.1.1.1" "8.8.8.8" ];
  networking.firewall.enable = true;
  networking.firewall.allowedTCPPorts = [ 22 8080 ];  # SSH, Strata API (key required)

  time.timeZone = "America/Los_Angeles";
  i18n.defaultLocale = "en_US.UTF-8";

  nixpkgs.config.allowUnfree = true;
  nix.settings = {
    experimental-features = [ "nix-command" "flakes" ];
    # prebuilt CUDA packages, so nothing CUDA-heavy compiles locally
    substituters = [ "https://cache.nixos.org" "https://cuda-maintainers.cachix.org" ];
    trusted-public-keys = [
      "cache.nixos.org-1:6NCHdD59X431o0gWypbMrAURkbJ16ZPMQFGspcDShjY="
      "cuda-maintainers.cachix.org-1:0dq3bujKpuEPMCX6U4WylrUDZ9JyUG0VpVZa7CNfq5E="
    ];
  };

  # NVIDIA: open kernel module (Ampere is supported), headless, persistence daemon
  services.xserver.videoDrivers = [ "nvidia" ];
  hardware.graphics.enable = true;
  hardware.nvidia = {
    open = true;
    package = config.boot.kernelPackages.nvidiaPackages.stable;
    nvidiaPersistenced = true;
    nvidiaSettings = false;
  };
  # Cap both 3090s at 275 W (default 350 W); reapplied on every boot. 275 vs 250: +4.6% prompt tok/s, see benchmarks/standard
  systemd.services.nvidia-power-limit = {
    description = "Set NVIDIA GPU power limit";
    after = [ "nvidia-persistenced.service" ];
    requires = [ "nvidia-persistenced.service" ];
    wantedBy = [ "multi-user.target" ];
    serviceConfig = {
      Type = "oneshot";
      RemainAfterExit = true;
      ExecStart = "${config.hardware.nvidia.package.bin}/bin/nvidia-smi -pl 275";
    };
  };
  # lovey's fan daemon (lovey-fans) can't see passed-through GPUs; send it the hottest one every 2 s
  systemd.services.gpu-temp-push = {
    description = "Push hottest GPU temp to lovey fan control";
    after = [ "nvidia-persistenced.service" "network-online.target" ];
    wants = [ "network-online.target" ];
    requires = [ "nvidia-persistenced.service" ];
    wantedBy = [ "multi-user.target" ];
    path = [ config.hardware.nvidia.package.bin pkgs.coreutils ];
    script = ''
      while true; do
        t=$(nvidia-smi --query-gpu=temperature.gpu --format=csv,noheader,nounits | sort -n | tail -1)
        echo "$t" > /dev/udp/192.168.1.53/7777 || true
        sleep 2
      done
    '';
    serviceConfig.Restart = "always";
  };
  hardware.nvidia-container-toolkit.enable = true;  # GPUs inside podman containers (--device nvidia.com/gpu=all)
  virtualisation.podman.enable = true;

  # Model weights shared from the host (lovey:/mnt/models); don't block boot if it's missing
  fileSystems."/models".options = [ "nofail" ];

  # Pinned/locked memory for CUDA, NCCL and vLLM
  security.pam.loginLimits = [
    { domain = "*"; type = "-"; item = "memlock"; value = "unlimited"; }
  ];
  systemd.settings.Manager.DefaultLimitMEMLOCK = "infinity";

  # Prebuilt Python wheels (pip/uv vLLM, torch) need a normal dynamic loader
  programs.nix-ld.enable = true;
  # libcuda/libnvidia-ml for pip/uv wheels (torch, vLLM); on NixOS they live in /run/opengl-driver/lib
  programs.nix-ld.libraries = [ config.hardware.nvidia.package pkgs.stdenv.cc.cc.lib pkgs.zlib ];

  services.qemuGuest.enable = true;
  services.openssh = {
    enable = true;
    settings.PasswordAuthentication = false;
    settings.PermitRootLogin = "prohibit-password";
  };

  users.users.jason = {
    isNormalUser = true;
    extraGroups = [ "wheel" "video" ];
    openssh.authorizedKeys.keys = [
      "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAINZMA1RumOjm198i2Jdaa1lw7vTWnnZKRuYwYhLp6YKx jason@cosmic"
    ];
  };
  # Strata inference engine: source checkout + venv in /var/lib/strata, model data on /models/strata
  users.groups.strata = { };
  users.users.strata = {
    isSystemUser = true;
    group = "strata";
    extraGroups = [ "video" ];
    home = "/var/lib/strata";
    createHome = true;
    shell = pkgs.bashInteractive;  # setup runs as this user
  };
  nix.settings.trusted-users = [ "root" "jason" ];
  systemd.tmpfiles.rules = [ "d /models/strata 0750 strata strata -" ];

  # The API server (setup's run-iq3_s.sh) with offline mode and no downloads. The API key and engine args live in
  # strata-iq3_s.json, which setup wrote; host 0.0.0.0, port 8080.
  systemd.services.strata = {
    description = "Strata inference server (Qwen3.8-Flash-Next IQ3_S)";
    after = [ "network-online.target" "nvidia-persistenced.service" "nvidia-power-limit.service" ];
    wants = [ "network-online.target" ];
    requires = [ "nvidia-persistenced.service" ];
    unitConfig.RequiresMountsFor = [ "/models" ];
    wantedBy = [ "multi-user.target" ];
    environment = {
      HOME = "/var/lib/strata";
      # The engine's RUNPATH points libcuda at the CUDA toolkit's stub; the real driver must come first.
      LD_LIBRARY_PATH = "/run/opengl-driver/lib";
      SSL_CERT_FILE = "/etc/ssl/certs/ca-certificates.crt";
      STRATA_NO_ACQUIRE = "1";
    };
    serviceConfig = {
      User = "strata";
      Group = "strata";
      WorkingDirectory = "/var/lib/strata/Strata";
      ExecStart = "/var/lib/strata/Strata/.venv/bin/python /var/lib/strata/Strata/serve/server.py --engine strata --config /var/lib/strata/Strata/strata-iq3_s.json --port 8080 --offline";
      Restart = "on-failure";
      RestartSec = 10;
      LimitMEMLOCK = "infinity";  # the expert arena and KV are pinned host memory
      # Hardening: image paths can name local files, so the server sees only its own state and the models
      NoNewPrivileges = true;
      ProtectSystem = "strict";
      ProtectHome = true;
      ReadWritePaths = [ "/var/lib/strata" "/models/strata" ];
      PrivateTmp = true;
      ProtectKernelTunables = true;
      ProtectKernelModules = true;
      ProtectControlGroups = true;
      RestrictSUIDSGID = true;
      LockPersonality = true;
    };
  };

  users.users.root.openssh.authorizedKeys.keys = config.users.users.jason.openssh.authorizedKeys.keys;
  security.sudo.wheelNeedsPassword = false;  # key-only login; set a password with `passwd` for console use

  environment.systemPackages = with pkgs; [
    git vim htop btop nvtopPackages.nvidia pciutils tmux uv python3 wget curl
    cudaPackages.cuda_nvcc  # nvcc for building custom kernels; runtime libs come from the wheels
  ];

  system.stateVersion = "26.05";
}