flake.nix

{
  inputs = {
    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
    nixpkgs_master.url = "github:NixOS/nixpkgs/master";
    systems.url = "github:nix-systems/default";
    flake-utils.url = "github:numtide/flake-utils";
    flake-utils.inputs.systems.follows = "systems";
    git-hooks = {
      url = "github:cachix/git-hooks.nix";
      inputs.nixpkgs.follows = "nixpkgs";
    };
    treefmt-nix = {
      url = "github:numtide/treefmt-nix";
      inputs.nixpkgs.follows = "nixpkgs";
    };
  };

  outputs =
    {
      self,
      nixpkgs,
      flake-utils,
      git-hooks,
      treefmt-nix,
      ...
    }@inputs:
    flake-utils.lib.eachDefaultSystem (
      system:
      let
        pkgs = import nixpkgs {
          inherit system;
          config.allowUnfree = true;
          config.cudaSupport = true;
        };

        mpkgs = import inputs.nixpkgs_master {
          inherit system;
          config.allowUnfree = true;
          config.cudaSupport = true;
        };

        libList = [
          # Add needed packages here
          mpkgs.stdenv.cc.cc.lib
          pkgs.libGL
          pkgs.glib
          pkgs.zlib
        ]
        ++ pkgs.lib.optionals pkgs.stdenv.isLinux (with pkgs; [ ]);

        treefmtEval = treefmt-nix.lib.evalModule pkgs {
          projectRootFile = "flake.nix";
          programs.nixfmt.enable = true;
          programs.ruff-check.enable = true;
          programs.ruff-format.enable = true;
        };

        portablePreCommitConfig = pkgs.writeText "pre-commit-config.yaml" ''
          # DO NOT MODIFY
          # Generated by `nix run .#update-pre-commit-config`
          repos:
            - repo: https://github.com/astral-sh/ruff-pre-commit
              rev: v${pkgs.ruff.version}
              hooks:
                - id: ruff-check
                  args: [--fix]
                  types_or: [python, pyi]
                - id: ruff-format
                  types_or: [python, pyi]
        '';

        updatePreCommitConfig = pkgs.writeShellApplication {
          name = "update-pre-commit-config";
          runtimeInputs = [
            pkgs.coreutils
            pkgs.git
          ];
          text = ''
            repoRoot="$(git rev-parse --show-toplevel)"
            rm -f "$repoRoot/.pre-commit-config.yaml"
            install -m 0644 ${portablePreCommitConfig} "$repoRoot/.pre-commit-config.yaml"
          '';
        };

        pre-commit-check = git-hooks.lib.${system}.run {
          src = ./.;
          package = pkgs.prek;
          install.enable = false;
          hooks.treefmt = {
            enable = true;
            package = treefmtEval.config.build.wrapper;
          };
        };
      in
      with pkgs;
      {
        apps.update-pre-commit-config = {
          type = "app";
          program = "${updatePreCommitConfig}/bin/update-pre-commit-config";
          meta.description = "Regenerate the portable pre-commit configuration";
        };
        checks = {
          inherit pre-commit-check;
          formatting = treefmtEval.config.build.check self;
          pre-commit-config = runCommand "pre-commit-config-check" { nativeBuildInputs = [ diffutils ]; } ''
            diff --unified ${portablePreCommitConfig} ${self}/.pre-commit-config.yaml
            touch $out
          '';
        };
        formatter = treefmtEval.config.build.wrapper;
        devShells = {
          default =
            let
              python_with_pkgs = pkgs.python311.withPackages (pp: [
                # Add python pkgs here that you need from nix repos
              ]);
            in
            mkShell {
              NIX_LD = runCommand "ld.so" { } ''
                ln -s "$(cat '${pkgs.stdenv.cc}/nix-support/dynamic-linker')" $out
              '';
              NIX_LD_LIBRARY_PATH = lib.makeLibraryPath libList;
              packages = [
                python_with_pkgs
                python311Packages.venvShellHook
                mpkgs.uv
                prek

              ]
              ++ libList;
              venvDir = "./.venv";
              postVenvCreation = ''
                unset SOURCE_DATE_EPOCH
              '';
              postShellHook = ''
                unset SOURCE_DATE_EPOCH
              '';
              shellHook = ''
                if ${pkgs.git}/bin/git rev-parse --git-dir > /dev/null 2>&1; then
                  repoRoot="$(${pkgs.git}/bin/git rev-parse --show-toplevel)"
                  (cd "$repoRoot" && prek install --config .pre-commit-config.yaml)
                fi

                export LD_LIBRARY_PATH=$NIX_LD_LIBRARY_PATH:$LD_LIBRARY_PATH
                export PYTHON_KEYRING_BACKEND=keyring.backends.fail.Keyring
                runHook venvShellHook
                uv sync --frozen --all-extras --dev
                export PYTHONPATH=${python_with_pkgs}/${python_with_pkgs.sitePackages}:$PYTHONPATH
              '';
            };
        };
      }
    );
}