nixpkgs/nixos/modules/programs/zsh/oh-my-zsh.nix

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

let

  cfg = config.programs.zsh.ohMyZsh;

  mkLinkFarmEntry = name: dir:
    let
      env = pkgs.buildEnv {
        name = "zsh-${name}-env";
        paths = cfg.customPkgs;
        pathsToLink = "/share/zsh/${dir}";
      };
    in
      { inherit name; path = "${env}/share/zsh/${dir}"; };

  mkLinkFarmEntry' = name: mkLinkFarmEntry name name;

  custom =
    if cfg.custom != null then cfg.custom
    else if builtins.length cfg.customPkgs == 0 then null
    else pkgs.linkFarm "oh-my-zsh-custom" [
      (mkLinkFarmEntry' "themes")
      (mkLinkFarmEntry "completions" "site-functions")
      (mkLinkFarmEntry' "plugins")
    ];

in
  {
    imports = [
      (lib.mkRenamedOptionModule [ "programs" "zsh" "oh-my-zsh" "enable" ] [ "programs" "zsh" "ohMyZsh" "enable" ])
      (lib.mkRenamedOptionModule [ "programs" "zsh" "oh-my-zsh" "theme" ] [ "programs" "zsh" "ohMyZsh" "theme" ])
      (lib.mkRenamedOptionModule [ "programs" "zsh" "oh-my-zsh" "custom" ] [ "programs" "zsh" "ohMyZsh" "custom" ])
      (lib.mkRenamedOptionModule [ "programs" "zsh" "oh-my-zsh" "plugins" ] [ "programs" "zsh" "ohMyZsh" "plugins" ])
    ];

    options = {
      programs.zsh.ohMyZsh = {
        enable = lib.mkOption {
          type = lib.types.bool;
          default = false;
          description = ''
            Enable oh-my-zsh.
          '';
        };

        package = lib.mkPackageOption pkgs "oh-my-zsh" { };

        plugins = lib.mkOption {
          default = [];
          type = lib.types.listOf(lib.types.str);
          description = ''
            List of oh-my-zsh plugins
          '';
        };

        custom = lib.mkOption {
          default = null;
          type = with lib.types; nullOr str;
          description = ''
            Path to a custom oh-my-zsh package to override config of oh-my-zsh.
            (Can't be used along with `customPkgs`).
          '';
        };

        customPkgs = lib.mkOption {
          default = [];
          type = lib.types.listOf lib.types.package;
          description = ''
            List of custom packages that should be loaded into `oh-my-zsh`.
          '';
        };

        theme = lib.mkOption {
          default = "";
          type = lib.types.str;
          description = ''
            Name of the theme to be used by oh-my-zsh.
          '';
        };

        cacheDir = lib.mkOption {
          default = "$HOME/.cache/oh-my-zsh";
          type = lib.types.str;
          description = ''
            Cache directory to be used by `oh-my-zsh`.
            Without this option it would default to the read-only nix store.
          '';
        };

        preLoaded = lib.mkOption {
          type = lib.types.lines;
          default = "";
          description = ''
            Shell commands executed before the `oh-my-zsh` is loaded.
            For example, to disable async git prompt write `zstyle ':omz:alpha:lib:git' async-prompt no` (more information https://github.com/ohmyzsh/ohmyzsh?tab=readme-ov-file#async-git-prompt)
          '';
        };
      };
    };

    config = lib.mkIf cfg.enable {

      # Prevent zsh from overwriting oh-my-zsh's prompt
      programs.zsh.promptInit = lib.mkDefault "";

      environment.systemPackages = [ cfg.package ];

      programs.zsh.interactiveShellInit = ''
        # oh-my-zsh configuration generated by NixOS
        export ZSH=${cfg.package}/share/oh-my-zsh

        ${lib.optionalString (builtins.length(cfg.plugins) > 0)
          "plugins=(${builtins.concatStringsSep " " cfg.plugins})"
        }

        ${lib.optionalString (custom != null)
          "ZSH_CUSTOM=\"${custom}\""
        }

        ${lib.optionalString (builtins.stringLength(cfg.theme) > 0)
          "ZSH_THEME=\"${cfg.theme}\""
        }

        ${lib.optionalString (cfg.cacheDir != null) ''
          if [[ ! -d "${cfg.cacheDir}" ]]; then
            mkdir -p "${cfg.cacheDir}"
          fi
          ZSH_CACHE_DIR=${cfg.cacheDir}
        ''}

        ${cfg.preLoaded}
        source $ZSH/oh-my-zsh.sh
      '';

      assertions = [
        {
          assertion = cfg.custom != null -> cfg.customPkgs == [];
          message = "If `cfg.custom` is set for `ZSH_CUSTOM`, `customPkgs` can't be used!";
        }
      ];

    };

    meta.doc = ./oh-my-zsh.md;
  }