nixpkgs/modules/home-manager/coding/agents/pi.nix

{
  config,
  lib,
  pkgs,
  ...
}: {
  imports = [
    ./shared/default.nix
  ];

  options.coding.agents.pi = let
    shared = import ./shared/shared-options.nix {inherit lib;};
    mcpCfg = config.programs.mcp or null;
  in
    with lib; {
      enable = mkEnableOption "Pi agent management via canonical agent.toml definitions";

      mcpServers = mkOption {
        type = types.attrsOf types.anything;
        default =
          if mcpCfg != null
          then mcpCfg.servers
          else {};
        defaultText = literalExpression "config.programs.mcp.servers";
        description = ''
          MCP server configurations for Pi (pi-mcp-adapter).
          Written to ~/.pi/agent/mcp.json.
          Automatically inherits from config.programs.mcp.servers.
        '';
      };

      agentsInput = shared.mkAgentsInputOption ''
        The `agents` flake input (your personal AGENTS repo).
        When set, the primary agent's system prompt is rendered as SYSTEM.md,
        all agents are listed in AGENTS.md, and subagent .md files are deployed.
      '';

      modelOverrides = shared.mkModelOverridesOption;

      primaryAgent = mkOption {
        type = types.nullOr types.str;
        default = null;
        description = ''
          Override which canonical agent is used as primary for SYSTEM.md.
          When null, the first agent with mode="primary" is used.
        '';
      };

      codingRules = mkOption {
        type = types.nullOr (types.submodule {
          options = {
            languages = mkOption {
              type = types.listOf types.str;
              default = [];
              description = ''
                Language-specific coding rules to include
                (e.g. [ "python" "typescript" "nix" ]).
                Rule files are read from the AGENTS repo's rules/languages/ directory.
              '';
            };

            concerns = mkOption {
              type = types.listOf types.str;
              default = [
                "coding-style"
                "naming"
                "documentation"
                "testing"
                "git-workflow"
                "project-structure"
              ];
              description = ''
                Concern rules to include from the AGENTS repo's rules/concerns/ directory.
              '';
            };

            frameworks = mkOption {
              type = types.listOf types.str;
              default = [];
              description = ''
                Framework-specific coding rules to include
                (e.g. [ "react" "fastapi" ]).
                Rule files are read from the AGENTS repo's rules/frameworks/ directory.
              '';
            };
          };
        });
        default = null;
        description = ''
          Coding rules to inject into ~/.pi/agent/AGENTS.md.
          Rules are read from the AGENTS repository and appended as markdown sections.
          Requires agentsInput to be set.
        '';
        example = literalExpression ''
          {
            languages = [ "python" "typescript" ];
            concerns = [ "coding-style" "testing" ];
            frameworks = [ "fastapi" ];
          }
        '';
      };

      settings = mkOption {
        type = types.submodule {
          freeformType = types.attrsOf types.anything;
          options = {
            packages = mkOption {
              type = types.listOf types.str;
              default = [];
              description = ''
                Pi packages to install (npm:, git:, or local paths).
                These are written to ~/.pi/agent/settings.json.
              '';
            };

            defaultProvider = mkOption {
              type = types.nullOr types.str;
              default = null;
              description = "Default LLM provider (e.g. 'anthropic', 'openai', 'zai').";
            };

            defaultModel = mkOption {
              type = types.nullOr types.str;
              default = null;
              description = "Default model ID.";
            };

            defaultThinkingLevel = mkOption {
              type = types.nullOr (types.enum ["off" "minimal" "low" "medium" "high" "xhigh"]);
              default = null;
              description = "Default extended thinking level.";
            };

            theme = mkOption {
              type = types.nullOr types.str;
              default = null;
              description = "Pi theme name.";
            };

            hideThinkingBlock = mkOption {
              type = types.nullOr types.bool;
              default = null;
              description = "Hide thinking blocks in output.";
            };

            quietStartup = mkOption {
              type = types.nullOr types.bool;
              default = null;
              description = "Hide startup header.";
            };

            compaction = mkOption {
              type = types.nullOr (types.submodule {
                options = {
                  enabled = mkOption {
                    type = types.nullOr types.bool;
                    default = null;
                  };
                  reserveTokens = mkOption {
                    type = types.nullOr types.int;
                    default = null;
                  };
                  keepRecentTokens = mkOption {
                    type = types.nullOr types.int;
                    default = null;
                  };
                };
              });
              default = null;
              description = "Auto-compaction settings.";
            };

            enabledModels = mkOption {
              type = types.nullOr (types.listOf types.str);
              default = null;
              description = "Model patterns for Ctrl+P cycling.";
            };

            sessionDir = mkOption {
              type = types.nullOr types.str;
              default = null;
              description = "Directory where session files are stored.";
            };

            extensions = mkOption {
              type = types.listOf types.str;
              default = [];
              description = "Local extension file paths or directories.";
            };

            skills = mkOption {
              type = types.listOf types.str;
              default = [];
              description = "Local skill file paths or directories.";
            };
          };
        };
        default = {};
        description = ''
          Pi settings written to ~/.pi/agent/settings.json.
          Only non-null values are included in the generated JSON.
          See pi docs/settings.md for all options.
        '';
      };

      # ── Pi Guardrails ─────────────────────────────────────────────
      guardrails = mkOption {
        type = types.nullOr (types.submodule {
          options = {
            enable =
              mkEnableOption
              "Generate ~/.pi/agent/extensions/guardrails.json for pi-guardrails. "
              + "Adds @aliou/pi-guardrails to packages automatically.";

            config = mkOption {
              type = types.attrsOf types.anything;
              default = {};
              description = ''
                Guardrails configuration written to ~/.pi/agent/extensions/guardrails.json.
                See https://github.com/aliou/pi-guardrails for config schema.

                IMPORTANT: Path access checks are lexical (not symlink-safe).
                Local project .pi/extensions/guardrails.json can override same rule IDs
                (memory > local > global > defaults). For immutable global policies,
                consider a wrapper or upstream patch.
              '';
            };
          };
        });
        default = null;
        description = ''
          Pi Guardrails security configuration.
          Generates ~/.pi/agent/extensions/guardrails.json when enabled.
          The @aliou/pi-guardrails package is added to settings.packages automatically.
        '';
      };
    };

  config = with lib; let
    shared = import ./shared/shared-options.nix {inherit lib;};
    cfg = config.coding.agents.pi;
  in
    mkIf cfg.enable (let
      # Build settings.json by filtering out null values recursively
      filterNulls = attrs:
        lib.filterAttrs (_: v: v != null) (
          builtins.mapAttrs (_: v:
            if builtins.isAttrs v
            then let
              filtered = filterNulls v;
            in
              if filtered == {}
              then null
              else filtered
            else v)
          attrs
        );

      # Base settings (already filtered)
      piSettings = filterNulls cfg.settings;

      # Guardrails package to inject when guardrails is enabled
      guardrailsPackage = "npm:@aliou/pi-guardrails@0.11.1";

      # Guardrails config (only when guardrails is enabled)
      guardrailsJson =
        if (cfg.guardrails != null && cfg.guardrails.enable)
        then builtins.toJSON cfg.guardrails.config
        else null;

      # Merge guardrails package into settings.packages when guardrails is enabled
      piSettingsWithGuardrails = let
        baseSettings = cfg.settings;
        basePackages = baseSettings.packages or [];
        hasGuardrailsPackage =
          lib.any
          (p:
            lib.hasPrefix "npm:@aliou/pi-guardrails" p
            || (lib.hasPrefix "git:" p && lib.hasSuffix "/pi-guardrails" p))
          basePackages;
        packagesWithGuardrails =
          if (cfg.guardrails != null && cfg.guardrails.enable && !hasGuardrailsPackage)
          then basePackages ++ [guardrailsPackage]
          else basePackages;
      in
        if packagesWithGuardrails != basePackages
        then filterNulls (baseSettings // {packages = packagesWithGuardrails;})
        else piSettings;

      # Coding rules config for renderForPi (only when both agentsInput and codingRules are set)
      piCodingRules =
        if cfg.agentsInput != null && cfg.codingRules != null
        then cfg.codingRules // {agents = cfg.agentsInput;}
        else null;

      # Rendered agents (only computed when agentsInput is set)
      rendered =
        if cfg.agentsInput != null
        then
          (import ../../../../lib {inherit lib;}).agents.renderForPi {
            inherit pkgs;
            canonical = cfg.agentsInput.lib.loadAgents;
            modelOverrides = cfg.modelOverrides;
            primaryAgent = cfg.primaryAgent;
            codingRules = piCodingRules;
          }
        else null;

      # Dynamic home.file entries for agent .md files
      agentFiles =
        if cfg.agentsInput != null
        then let
          agentNames = builtins.attrNames cfg.agentsInput.lib.loadAgents;
        in
          builtins.listToAttrs (
            map (name: {
              name = ".pi/agent/agents/${name}.md";
              value = {source = "${rendered}/agents/${name}.md";};
            })
            agentNames
          )
        else {};
    in {
      home.file = mkMerge [
        # ── MCP servers from programs.mcp → ~/.pi/agent/mcp.json ───────
        (mkIf (cfg.mcpServers != {}) {
          ".pi/agent/mcp.json".text = builtins.toJSON {mcpServers = cfg.mcpServers;};
          ".pi/agent/mcp.json".force = true;
        })

        # ── ~/.pi/agent/settings.json ──────────────────────────────────
        {
          ".pi/agent/settings.json".text = builtins.toJSON piSettingsWithGuardrails;
          ".pi/agent/settings.json".force = true;
        }

        # ── pi-guardrails config ─────────────────────────────────────
        (mkIf (guardrailsJson != null) {
          ".pi/agent/extensions/guardrails.json".text = guardrailsJson;
          ".pi/agent/extensions/guardrails.json".force = true;
        })

        # ── AGENTS.md — agent descriptions and specialist listing ──────
        (mkIf (cfg.agentsInput != null) {
          ".pi/agent/AGENTS.md".source = "${rendered}/AGENTS.md";
        })

        # ── SYSTEM.md — primary agent's system prompt ──────────────────
        (mkIf (cfg.agentsInput != null) {
          ".pi/agent/SYSTEM.md".source = "${rendered}/SYSTEM.md";
        })

        # ── Agents — pi-subagents .md files ────────────────────────────
        agentFiles
      ];
    });
}