# NixOS Module for isolated Pi execution (fresh design)
#
# Goals:
# - Dedicated isolated runtime identity (pi-agent user/group)
# - Host UX via `pi` wrapper command
# - Per-host-user project allowlists (different roots per user)
# - No container mode
# - Merge user Pi config + Nix-managed settings/env into isolated runtime
{
  config,
  lib,
  pkgs,
  ...
}:
with lib; let
  cfg = config.m3ta.pi-agent;

  hostUserNames = attrNames cfg.hostUsers;

  runner = import ./pi-agent-runner.nix {inherit cfg pkgs lib;};
  wrapper = import ./pi-agent-wrapper.nix {inherit cfg pkgs lib runner;};
in {
  options.m3ta.pi-agent = {
    enable = mkEnableOption "isolated Pi execution with dedicated system user and policy-enforced wrapper";

    package = mkOption {
      type = types.package;
      default = pkgs.pi-coding-agent;
      defaultText = literalExpression "pkgs.pi-coding-agent";
      description = "Pi package providing the executable used in isolated runtime.";
    };

    binaryName = mkOption {
      type = types.str;
      default = "pi-agent";
      description = "Preferred executable name inside `${cfg.package}/bin` (falls back to pi/pi-agent auto-detection).";
      example = "pi";
    };

    user = mkOption {
      type = types.str;
      default = "pi-agent";
      description = "System user that executes Pi in isolated mode.";
    };

    group = mkOption {
      type = types.str;
      default = "pi-agent";
      description = "System group for the isolated Pi user.";
    };

    stateDir = mkOption {
      type = types.str;
      default = "/var/lib/pi-agent";
      description = "Writable state/home directory for isolated Pi runtime.";
    };

    createUser = mkOption {
      type = types.bool;
      default = true;
      description = "Whether to create the dedicated Pi user/group automatically.";
    };

    hostUsers = mkOption {
      type = types.attrsOf (types.submodule {
        options = {
          projectRoots = mkOption {
            type = types.listOf types.str;
            default = [];
            description = ''
              Allowed project roots for this host user.
              `~` and `~/...` are expanded relative to that host user's home.
            '';
            example = ["~/p" "~/work/client-a"];
          };

          configPath = mkOption {
            type = types.nullOr types.str;
            default = null;
            description = ''
              Optional host path for this user's Pi config source. If null,
              wrapper.hostConfigPath is used. Relative paths resolve from the
              host user's home.
            '';
            example = ".pi/agent";
          };
        };
      });
      default = {};
      description = ''
        Per-host-user policy map. Keys are host usernames.
        Each user defines their own allowed project roots and optional config source.
      '';
      example = literalExpression ''
        {
          m3tam3re = {
            projectRoots = [ "~/p" "~/src/private" ];
            configPath = ".pi/agent";
          };
          teammate = {
            projectRoots = [ "~/projects" ];
          };
        }
      '';
    };

    settings = mkOption {
      type = types.attrsOf types.anything;
      default = {};
      description = ''
        Nix-managed Pi settings merged into isolated `${cfg.stateDir}/.pi/agent/settings.json`.
        Merge precedence: synced host settings first, Nix-managed values override recursively.
      '';
      example = literalExpression ''
        {
          defaultModel = "anthropic/claude-sonnet-4";
          defaultProvider = "anthropic";
          quietStartup = true;
        }
      '';
    };

    environment = mkOption {
      type = types.attrsOf types.str;
      default = {};
      description = ''
        Non-secret Nix-managed environment variables appended into isolated
        `${cfg.stateDir}/.pi/.env` after synced host values.
      '';
    };

    environmentFiles = mkOption {
      type = types.listOf types.str;
      default = [];
      description = ''
        Paths to env files (secrets/tokens) appended to isolated `${cfg.stateDir}/.pi/.env`
        after `environment` entries.
      '';
    };

    extraPackages = mkOption {
      type = types.listOf types.package;
      default = [];
      description = "Extra packages added to isolated runtime PATH.";
    };

    projectGroup = mkOption {
      type = types.nullOr types.str;
      default = null;
      description = ''
        When set, the pi-agent user is added to this group and the group is
        passed as SupplementaryGroups to the systemd-run sandbox. This allows
        pi-agent to write to project directories that grant group write access.
        The user must ensure project directories have appropriate group ownership
        and permissions (e.g. setgid + group write).
      '';
      example = "users";
    };

    wrapper = {
      enable = mkOption {
        type = types.bool;
        default = true;
        description = "Enable host-side wrapper command that enforces policy and runs isolated Pi.";
      };

      commandName = mkOption {
        type = types.str;
        default = "pi";
        description = "Host wrapper command name.";
      };

      runnerName = mkOption {
        type = types.str;
        default = "m3ta-pi-agent-runner";
        description = "Privileged runner command invoked via scoped sudo rule.";
      };

      hideDirectBinary = mkOption {
        type = types.bool;
        default = true;
        description = ''
          When true and wrapper is enabled, do not add the raw Pi package to host PATH,
          reducing bypass risk by making wrapper the canonical entrypoint.
        '';
      };

      syncConfigFromHost = mkOption {
        type = types.bool;
        default = true;
        description = ''
          Sync host Pi config directory into isolated `${cfg.stateDir}/.pi/agent`
          on each invocation.
        '';
      };

      hostConfigPath = mkOption {
        type = types.str;
        default = ".pi/agent";
        description = ''
          Default source path for host Pi config sync. Relative paths resolve from
          the invoking user's home. Per-user hostUsers.<name>.configPath overrides this.
        '';
      };

      extraRunArgs = mkOption {
        type = types.listOf types.str;
        default = [];
        description = "Extra arguments inserted before user-provided Pi args.";
      };

      extraEnvironment = mkOption {
        type = types.attrsOf types.str;
        default = {};
        description = "Additional environment variables passed to isolated Pi runtime.";
      };
    };
  };

  config = mkIf cfg.enable {
    assertions =
      [
        {
          assertion = cfg.hostUsers != {};
          message = "m3ta.pi-agent.hostUsers must define at least one authorized host user.";
        }
        {
          assertion = (!cfg.wrapper.enable) || (cfg.hostUsers != {});
          message = "m3ta.pi-agent.hostUsers must not be empty when wrapper is enabled.";
        }
      ]
      ++ mapAttrsToList (user: userCfg: {
        assertion = userCfg.projectRoots != [];
        message = "m3ta.pi-agent.hostUsers.${user}.projectRoots must not be empty.";
      })
      cfg.hostUsers;

    users.groups = mkIf cfg.createUser {
      "${cfg.group}" = {};
    };

    users.users = mkIf cfg.createUser {
      "${cfg.user}" = {
        isSystemUser = true;
        group = cfg.group;
        extraGroups = mkIf (cfg.projectGroup != null) [cfg.projectGroup];
        description = "Isolated Pi agent user";
        home = cfg.stateDir;
        createHome = true;
        shell = pkgs.bashInteractive;
      };
    };

    systemd.tmpfiles.rules = [
      "d ${cfg.stateDir} 0750 ${cfg.user} ${cfg.group} - -"
      "d ${cfg.stateDir}/.pi 0750 ${cfg.user} ${cfg.group} - -"
      "d ${cfg.stateDir}/.pi/agent 0750 ${cfg.user} ${cfg.group} - -"
      "d ${cfg.stateDir}/.pi/agent/sessions 0750 ${cfg.user} ${cfg.group} - -"
      "d ${cfg.stateDir}/.project-mounts 0750 ${cfg.user} ${cfg.group} - -"
      "d ${cfg.stateDir}/projects 0750 ${cfg.user} ${cfg.group} - -"
      "d ${cfg.stateDir}/.npm 0750 ${cfg.user} ${cfg.group} - -"
      "d ${cfg.stateDir}/.npm-global 0750 ${cfg.user} ${cfg.group} - -"
      "d ${cfg.stateDir}/.npm-global/bin 0750 ${cfg.user} ${cfg.group} - -"
      "d ${cfg.stateDir}/.npm-global/lib 0750 ${cfg.user} ${cfg.group} - -"
    ];

    # Ensure correct ownership of stateDir after user creation.
    # createHome = true causes useradd to create the directory as root:root
    # before systemd-tmpfiles can set the intended owner.
    system.activationScripts.pi-agent-chown = {
      deps = ["users"];
      text = ''
        chown ${cfg.user}:${cfg.group} ${cfg.stateDir}
      '';
    };

    # Wrapper is canonical when enabled; raw package on PATH is optional and
    # disabled by default to reduce bypass opportunities.
    environment.systemPackages =
      optional cfg.wrapper.enable wrapper
      ++ optional ((!cfg.wrapper.enable) || (!cfg.wrapper.hideDirectBinary)) cfg.package;

    security.sudo.extraRules = mkIf (cfg.wrapper.enable && hostUserNames != []) [
      {
        users = hostUserNames;
        commands = [
          {
            command = "${runner}/bin/${cfg.wrapper.runnerName}";
            options = ["NOPASSWD"];
          }
        ];
      }
    ];
  };
}