# 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;

  managedSettingsFile = pkgs.writeText "pi-agent-managed-settings.json" (builtins.toJSON cfg.settings);

  managedEnvFile =
    pkgs.writeText "pi-agent-managed.env"
    (concatStringsSep "\n" (mapAttrsToList (k: v: "${k}=${v}") cfg.environment));

  runtimePath = concatStringsSep ":" (
    [
      "${cfg.package}/bin"
      "${pkgs.nodejs}/bin"
      "${pkgs.git}/bin"
      "${pkgs.coreutils}/bin"
      "${pkgs.findutils}/bin"
      "${pkgs.gnugrep}/bin"
      "${pkgs.gnused}/bin"
      "${pkgs.util-linux}/bin"
      "/run/current-system/sw/bin"
    ]
    ++ map (p: "${p}/bin") cfg.extraPackages
  );

  userPolicyCase = concatStringsSep "\n" (
    mapAttrsToList (
      user: userCfg: ''
        ${escapeShellArg user})
          USER_CONFIG_PATH=${escapeShellArg (
          if userCfg.configPath != null
          then userCfg.configPath
          else cfg.wrapper.hostConfigPath
        )}
          USER_ROOTS=(${concatStringsSep " " (map escapeShellArg userCfg.projectRoots)})
          ;;
      ''
    )
    cfg.hostUsers
  );

  runner = pkgs.writeShellScriptBin cfg.wrapper.runnerName ''
        set -euo pipefail

        if [ "$(id -u)" -ne 0 ]; then
          echo "${cfg.wrapper.runnerName} must run as root" >&2
          exit 1
        fi

        if [ "$#" -lt 2 ]; then
          echo "Usage: ${cfg.wrapper.runnerName} <invoking-user> <cwd> [pi-args...]" >&2
          exit 2
        fi

        invoking_user="$1"
        shift
        cwd="$1"
        shift

        resolve_user_policy() {
          local user="$1"
          USER_CONFIG_PATH=""
          USER_ROOTS=()
          case "$user" in
            ${userPolicyCase}
            *)
              return 1
              ;;
          esac
          return 0
        }

        if ! resolve_user_policy "$invoking_user"; then
          echo "User '$invoking_user' is not allowed to use ${cfg.wrapper.commandName}" >&2
          exit 1
        fi

        user_home="$(eval echo "~$invoking_user")"
        if [ -z "$user_home" ] || [ "$user_home" = "~$invoking_user" ]; then
          echo "Unable to determine home directory for user '$invoking_user'" >&2
          exit 1
        fi

        expand_home_path() {
          local input="$1"
          if [ "$input" = "~" ]; then
            printf '%s\n' "$user_home"
          elif ${pkgs.gnugrep}/bin/grep -q '^~/' <<<"$input"; then
            printf '%s\n' "$user_home/''${input:2}"
          elif ${pkgs.gnugrep}/bin/grep -q '^/' <<<"$input"; then
            printf '%s\n' "$input"
          else
            # Bare relative path → resolve from user's home
            printf '%s\n' "$user_home/$input"
          fi
        }

        cwd_real="$(${pkgs.coreutils}/bin/realpath -m "$cwd")"

        resolved_roots=()
        skipped_roots=()
        is_allowed_cwd=0
        for configured_root in "''${USER_ROOTS[@]}"; do
          expanded_root="$(expand_home_path "$configured_root")"
          resolved_root="$(${pkgs.coreutils}/bin/realpath -m "$expanded_root")"
          if [ ! -d "$resolved_root" ]; then
            skipped_roots+=("$resolved_root")
            continue
          fi
          resolved_roots+=("$resolved_root")
          case "$cwd_real/" in
            "$resolved_root"/*)
              is_allowed_cwd=1
              ;;
          esac
        done

        if [ "''${#resolved_roots[@]}" -eq 0 ]; then
          echo "Denied: no valid existing project roots are configured for user '$invoking_user'." >&2
          if [ "''${#skipped_roots[@]}" -gt 0 ]; then
            echo "Configured but missing roots:" >&2
            for root in "''${skipped_roots[@]}"; do
              echo "  - $root" >&2
            done
          fi
          exit 1
        fi

        if [ "$is_allowed_cwd" -ne 1 ]; then
          echo "Denied: '$cwd_real' is outside allowed project roots for user '$invoking_user'." >&2
          echo "Allowed roots:" >&2
          for root in "''${resolved_roots[@]}"; do
            echo "  - $root" >&2
          done
          exit 1
        fi

        ${pkgs.coreutils}/bin/install -d -m 0750 -o ${escapeShellArg cfg.user} -g ${escapeShellArg cfg.group} \
          ${escapeShellArg cfg.stateDir} \
          ${escapeShellArg "${cfg.stateDir}/.pi"} \
          ${escapeShellArg "${cfg.stateDir}/.pi/agent"} \
          ${escapeShellArg "${cfg.stateDir}/.project-mounts"} \
          ${escapeShellArg "${cfg.stateDir}/projects"} \
          ${escapeShellArg "${cfg.stateDir}/.npm"} \
          ${escapeShellArg "${cfg.stateDir}/.npm-global"} \
          ${escapeShellArg "${cfg.stateDir}/.npm-global/bin"} \
          ${escapeShellArg "${cfg.stateDir}/.npm-global/lib"}

        config_source="$USER_CONFIG_PATH"
        if ${pkgs.gnugrep}/bin/grep -q '^/' <<<"$config_source"; then
          source_dir="$config_source"
        else
          source_dir="$(expand_home_path "$config_source")"
        fi


        if [ "${
      if cfg.wrapper.syncConfigFromHost
      then "1"
      else "0"
    }" = "1" ] && [ -d "$source_dir" ]; then
          ${pkgs.rsync}/bin/rsync -a --delete "$source_dir/" ${escapeShellArg "${cfg.stateDir}/.pi/agent/"}
          ${pkgs.coreutils}/bin/chown -R ${escapeShellArg "${cfg.user}:${cfg.group}"} ${escapeShellArg "${cfg.stateDir}/.pi/agent"}
        fi

        # Merge host settings.json (if any) with Nix-managed settings.
        # Precedence: host settings first, Nix-managed keys override recursively.
        settings_target=${escapeShellArg "${cfg.stateDir}/.pi/agent/settings.json"}
        ${pkgs.python3}/bin/python3 - "$settings_target" ${escapeShellArg managedSettingsFile} <<'PY_PI_SETTINGS_MERGE'
    import json
    import os
    import sys


    def load_obj(path):
        if not os.path.exists(path):
            return {}
        try:
            with open(path, "r", encoding="utf-8") as f:
                data = json.load(f)
                return data if isinstance(data, dict) else {}
        except Exception:
            return {}


    def deep_merge(base, override):
        if isinstance(base, dict) and isinstance(override, dict):
            out = dict(base)
            for key, value in override.items():
                out[key] = deep_merge(out.get(key), value)
            return out
        return override


    def main():
        target = sys.argv[1]
        managed = sys.argv[2]
        base_obj = load_obj(target)
        managed_obj = load_obj(managed)
        merged = deep_merge(base_obj, managed_obj)

        os.makedirs(os.path.dirname(target), exist_ok=True)
        tmp = f"{target}.tmp"
        with open(tmp, "w", encoding="utf-8") as f:
            json.dump(merged, f, indent=2, sort_keys=True)
            f.write("\n")
        os.replace(tmp, target)


    if __name__ == "__main__":
        main()
    PY_PI_SETTINGS_MERGE
        ${pkgs.coreutils}/bin/chown ${escapeShellArg "${cfg.user}:${cfg.group}"} "$settings_target"
        ${pkgs.coreutils}/bin/chmod 0640 "$settings_target"

        # Merge environment into isolated .env with precedence:
        # 1) synced host env (source_dir/.env)
        # 2) Nix-managed environment attrset
        # 3) Nix-managed environmentFiles (appended in declaration order)
        env_target=${escapeShellArg "${cfg.stateDir}/.pi/.env"}
        ${pkgs.coreutils}/bin/install -o ${escapeShellArg cfg.user} -g ${escapeShellArg cfg.group} -m 0640 /dev/null "$env_target"

        if [ -f "$source_dir/.env" ]; then
          ${pkgs.coreutils}/bin/cat "$source_dir/.env" >> "$env_target"
          printf '\n' >> "$env_target"
        fi

        if [ -f ${escapeShellArg managedEnvFile} ]; then
          ${pkgs.coreutils}/bin/cat ${escapeShellArg managedEnvFile} >> "$env_target"
          printf '\n' >> "$env_target"
        fi

        ${concatStringsSep "\n" (map (f: ''
        if [ -f ${escapeShellArg f} ]; then
          ${pkgs.coreutils}/bin/cat ${escapeShellArg f} >> "$env_target"
          printf '\n' >> "$env_target"
        fi
      '')
      cfg.environmentFiles)}

        ${pkgs.coreutils}/bin/chown ${escapeShellArg "${cfg.user}:${cfg.group}"} "$env_target"
        ${pkgs.coreutils}/bin/chmod 0640 "$env_target"

        npm_prefix=${escapeShellArg "${cfg.stateDir}/.npm-global"}
        runtime_path=${escapeShellArg runtimePath}

        project_mount_dir=${escapeShellArg "${cfg.stateDir}/.project-mounts"}
        project_links_dir=${escapeShellArg "${cfg.stateDir}/projects"}
        project_bind_pairs=()

        matched_root=""
        matched_mount=""
        project_index=0

        for root in "''${resolved_roots[@]}"; do
          if [ ! -d "$root" ]; then
            continue
          fi

          root_slug="$(printf '%s' "$root" | ${pkgs.gnused}/bin/sed 's#^/##; s#/#-#g; s#-\{2,\}#-#g; s#-$##; s#^$#root#')"
          root_slug="''${project_index}-''${root_slug}"
          project_index=$((project_index + 1))

          mount_point="''${project_mount_dir}/''${root_slug}"
          link_path="''${project_links_dir}/''${root_slug}"

          ${pkgs.coreutils}/bin/install -d -m 0750 -o ${escapeShellArg cfg.user} -g ${escapeShellArg cfg.group} "$mount_point"
          ${pkgs.coreutils}/bin/ln -sfn "$mount_point" "$link_path"

          project_bind_pairs+=("$root:$mount_point")

          case "$cwd_real/" in
            "$root"/*)
              if [ -z "$matched_root" ] || [ "''${#root}" -gt "''${#matched_root}" ]; then
                matched_root="$root"
                matched_mount="$mount_point"
              fi
              ;;
          esac
        done

        if [ -z "$matched_root" ]; then
          echo "Failed to map cwd '$cwd_real' to an allowed root." >&2
          exit 1
        fi

        if [ "$cwd_real" = "$matched_root" ]; then
          mapped_cwd="$matched_mount"
        else
          rel_path="''${cwd_real#"$matched_root/"}"
          mapped_cwd="$matched_mount/$rel_path"
        fi

        pi_bin=${escapeShellArg "${cfg.package}/bin/${cfg.binaryName}"}

        if [ ! -x "$pi_bin" ]; then
          for candidate in pi pi-agent; do
            alt=${escapeShellArg "${cfg.package}/bin"}/$candidate
            if [ -x "$alt" ]; then
              pi_bin="$alt"
              break
            fi
          done
        fi

        if [ ! -x "$pi_bin" ]; then
          echo "Pi binary not found or not executable: $pi_bin" >&2
          echo "Available executables in ${cfg.package}/bin:" >&2
          ${pkgs.coreutils}/bin/ls -1 ${escapeShellArg "${cfg.package}/bin"} >&2 || true
          exit 127
        fi

        cmd=(
          ${pkgs.systemd}/bin/systemd-run
          --collect
          --wait
          --pty
          --service-type=exec
          -p User=${cfg.user}
          -p Group=${cfg.group}
          -p WorkingDirectory="$mapped_cwd"
          -p NoNewPrivileges=yes
          -p PrivateTmp=yes
          -p ProtectSystem=strict
          -p ProtectHome=false
          -p ProtectControlGroups=yes
          -p ProtectKernelTunables=yes
          -p ProtectKernelModules=yes
          -p RestrictSUIDSGID=yes
          -p LockPersonality=yes
          -p RestrictRealtime=yes
          -p RestrictNamespaces=yes
          -p MemoryDenyWriteExecute=no
          -p UMask=0077
          -p ReadWritePaths=${cfg.stateDir}
          -p EnvironmentFile=${cfg.stateDir}/.pi/.env
          -E HOME=${cfg.stateDir}
          -E PI_HOME=${cfg.stateDir}/.pi
          -E MESSAGING_CWD="$mapped_cwd"
          -E PATH="$runtime_path"
          -E NPM_CONFIG_CACHE=${cfg.stateDir}/.npm
          -E NPM_CONFIG_PREFIX="$npm_prefix"
          -E PI_AGENT_INVOKING_USER="$invoking_user"
        )

        # Only mark existing top-level paths inaccessible; systemd fails namespace
        # setup if InaccessiblePaths points to a non-existent path on this host.
        for p in /home /root /mnt /media /srv; do
          if [ -e "$p" ]; then
            cmd+=( -p "InaccessiblePaths=$p" )
          fi
        done

        for pair in "''${project_bind_pairs[@]}"; do
          src="''${pair%%:*}"
          dst="''${pair#*:}"
          cmd+=( -p "BindPaths=$src:$dst" )
        done

        ${concatStringsSep "\n" (mapAttrsToList (name: value: ''cmd+=( -E ${escapeShellArg "${name}=${value}"} )'') cfg.wrapper.extraEnvironment)}

        cmd+=( "$pi_bin" )
        ${concatStringsSep "\n" (map (arg: ''cmd+=( ${escapeShellArg arg} )'') cfg.wrapper.extraRunArgs)}
        cmd+=( "$@" )

        exec "''${cmd[@]}"
  '';

  wrapper = pkgs.writeShellScriptBin cfg.wrapper.commandName ''
    set -euo pipefail

    user_name="$(id -un)"
    user_home="$(eval echo "~$user_name")"
    if [ -z "$user_home" ] || [ "$user_home" = "~$user_name" ]; then
      user_home="$HOME"
    fi

    resolve_user_policy() {
      local user="$1"
      USER_ROOTS=()
      case "$user" in
        ${concatStringsSep "\n" (
      mapAttrsToList (
        user: userCfg: ''
          ${escapeShellArg user})
            USER_ROOTS=(${concatStringsSep " " (map escapeShellArg userCfg.projectRoots)})
            ;;
        ''
      )
      cfg.hostUsers
    )}
        *)
          return 1
          ;;
      esac
      return 0
    }

    if ! resolve_user_policy "$user_name"; then
      echo "User '$user_name' is not allowed to use ${cfg.wrapper.commandName}" >&2
      exit 1
    fi

    expand_home_path() {
      local input="$1"
      if [ "$input" = "~" ]; then
        printf '%s\n' "$user_home"
      elif ${pkgs.gnugrep}/bin/grep -q '^~/' <<<"$input"; then
        printf '%s\n' "$user_home/''${input:2}"
      elif ${pkgs.gnugrep}/bin/grep -q '^/' <<<"$input"; then
        printf '%s\n' "$input"
      else
        printf '%s\n' "$user_home/$input"
      fi
    }

    cwd_real="$(${pkgs.coreutils}/bin/realpath -m "$PWD")"

    is_allowed_cwd=0
    resolved_roots=()
    skipped_roots=()
    for configured_root in "''${USER_ROOTS[@]}"; do
      expanded_root="$(expand_home_path "$configured_root")"
      resolved_root="$(${pkgs.coreutils}/bin/realpath -m "$expanded_root")"
      if [ ! -d "$resolved_root" ]; then
        skipped_roots+=("$resolved_root")
        continue
      fi
      resolved_roots+=("$resolved_root")
      case "$cwd_real/" in
        "$resolved_root"/*)
          is_allowed_cwd=1
          ;;
      esac
    done

    if [ "''${#resolved_roots[@]}" -eq 0 ]; then
      echo "Denied: no valid existing project roots are configured for user '$user_name'." >&2
      if [ "''${#skipped_roots[@]}" -gt 0 ]; then
        echo "Configured but missing roots:" >&2
        for root in "''${skipped_roots[@]}"; do
          echo "  - $root" >&2
        done
      fi
      exit 1
    fi

    if [ "$is_allowed_cwd" -ne 1 ]; then
      echo "Denied: '$cwd_real' is outside allowed project roots for user '$user_name'." >&2
      echo "Allowed roots:" >&2
      for root in "''${resolved_roots[@]}"; do
        echo "  - $root" >&2
      done
      exit 1
    fi

    exec /run/wrappers/bin/sudo --non-interactive ${runner}/bin/${cfg.wrapper.runnerName} "$user_name" "$cwd_real" "$@"
  '';
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.";
    };

    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;
        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}/.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} - -"
    ];

    # 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"];
          }
        ];
      }
    ];
  };
}