{nixpkgs ? import <nixpkgs> {}}:
nixpkgs.mkShell {
  nativeBuildInputs = with nixpkgs; [
    ruff
    python3
  ];

  LD_LIBRARY_PATH = "${nixpkgs.stdenv.cc.cc.lib}/lib:${nixpkgs.zstd.out}/lib:${nixpkgs.zlib.out}/lib";
  shellHook = ''
    if [[ ! -d ".venv" ]]; then
      python -m venv .venv
      source .venv/bin/activate
      pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm7.2/ --upgrade --force-reinstall
      pip install "unsloth[amd] @ git+https://github.com/unslothai/unsloth"
      pip install "unsloth-zoo[main] @ git+https://github.com/unslothai/unsloth-zoo"
      pip install marimo ipython
    else
      source .venv/bin/activate
    fi
  '';
}