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

with lib;

let
  # Docker subcommands that are BLOCKED for ai-worker
  # These commands allow file modification inside containers or data exfiltration.
  blockedCommands = [
    "exec"    # Execute arbitrary commands in containers (FILE MODIFICATION)
    "cp"      # Copy files between containers and host (FILE ACCESS)
    "commit"  # Create images from running containers (DATA EXFIL)
    "diff"    # Inspect filesystem changes of containers (INFO LEAK)
    "export"  # Export container filesystem as tar archive (DATA EXFIL)
    "import"  # Import a tar archive to create filesystem (FILE INJECTION)
    "load"    # Load images from tar archive (FILE INJECTION)
    "save"    # Save images to tar archive (DATA EXFIL)
    "attach"  # Attach to running container's stdio (INTERACTIVE ACCESS)
    "push"    # Push images to remote registries (DATA EXFIL)
    "tag"     # Tag/rename images (used with push)
  ];

  blockedDockerArgs = lib.concatStringsSep "|" blockedCommands;

  # Docker wrapper script that blocks dangerous subcommands
  # Must handle: docker exec, docker compose exec, docker cp, etc.
  restrictedDockerScript = pkgs.writeShellScriptBin "docker" ''
    set -e

    # Blocklist pattern
    BLOCKED_PATTERN="^(${blockedDockerArgs})$"

    # Parse the first non-flag argument to find the docker subcommand
    # Flags: -H, --host, -D, --debug, --config, --context, --log-level, -l
    # Also handle: docker compose <subcommand> (subcommand may be after 'compose')
    SUBCOMMAND=""
    COMPOSE_MODE=false
    FOUND_ARG=false

    for arg in "$@"; do
      # Skip flags and their values
      case "$arg" in
        -H|--host|-l|--log-level|--config|--context|-D|--debug)
          FOUND_ARG=true
          continue
          ;;
        --tls|--tlsverify|--tlscacert|--tlscert|--tlskey)
          if $FOUND_ARG; then FOUND_ARG=false; else continue; fi
          ;;
        # Skip flag values (the next arg after a flag that takes a value)
        -*)
          continue
          ;;
        *)
          # This is a positional argument — first one is the subcommand (or 'compose')
          if [ -z "$SUBCOMMAND" ]; then
            if [ "$arg" = "compose" ]; then
              COMPOSE_MODE=true
              continue
            fi
            SUBCOMMAND="$arg"
            break
          fi
          ;;
      esac
      FOUND_ARG=false
    done

    # If in compose mode, the subcommand is after 'compose'
    if $COMPOSE_MODE; then
      # In compose mode, we check the sub-subcommand
      NEXT_GOT=""
      for arg in "$@"; do
        if [ "$NEXT_GOT" = "true" ]; then
          if echo "$arg" | grep -qE "$BLOCKED_PATTERN"; then
            echo "ERROR: docker compose '$arg' is blocked by security policy" >&2
            echo "This command can modify files inside containers." >&2
            exit 1
          fi
          break
        fi
        if [ "$arg" = "compose" ]; then
          NEXT_GOT="true"
        fi
      done
    fi

    # Check if the subcommand is blocked
    if [ -n "$SUBCOMMAND" ]; then
      if echo "$SUBCOMMAND" | grep -qE "$BLOCKED_PATTERN"; then
        echo "ERROR: docker '$SUBCOMMAND' is blocked by security policy" >&2
        echo "This command can modify files inside containers." >&2
        echo "" >&2
        echo "Allowed commands: ps, images, inspect, logs, start, stop, restart," >&2
        echo "  rm, rmi, pull, build, run, compose, system, network ls, volume ls" >&2
        exit 1
      fi
    fi

    # Execute the real docker binary
    exec ${pkgs.docker}/bin/docker "$@"
  '';
in
{
  options.services.aiWorkerAccess = mkOption {
    type = types.bool;
    default = false;
    description = "Enable AI worker SSH access with restricted docker commands";
  };

  config = mkIf config.services.aiWorkerAccess {
    # ai-worker is in docker group for normal docker operations
    users.groups.docker.members = [ "ai-worker" ];

    # Install the docker wrapper for ai-worker
    # This puts a filtered 'docker' script in ai-worker's PATH that blocks
    # dangerous commands like exec, cp, commit, etc.
    # The real docker binary is still available at its store path, but the
    # wrapper intercepts it because ~/.nix-profile/bin/ comes before /run/.../sw/bin/ in PATH.
    users.users.ai-worker.packages = [ restrictedDockerScript ];

    # Also install the wrapper system-wide for consistency
    environment.systemPackages = [ restrictedDockerScript ];
  };
}