# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tool to execute bash commands."""

from __future__ import annotations

import asyncio
import dataclasses
import logging
import os
import pathlib
import resource
import shlex
import signal
from typing import Any
from typing import Optional

from google.genai import types

from .base_tool import BaseTool
from .tool_context import ToolContext

logger = logging.getLogger("google_adk." + __name__)


@dataclasses.dataclass(frozen=True)
class BashToolPolicy:
  """Configuration for allowed bash commands and resource limits.

  Set allowed_command_prefixes to ("*",) to allow all commands (default),
  or explicitly list allowed prefixes.

  Values for max_memory_bytes, max_file_size_bytes, and max_child_processes
  will be enforced upon the spawned subprocess.
  """

  allowed_command_prefixes: tuple[str, ...] = ("*",)
  blocked_operators: tuple[str, ...] = ()
  timeout_seconds: Optional[int] = 30
  max_memory_bytes: Optional[int] = None
  max_file_size_bytes: Optional[int] = None
  max_child_processes: Optional[int] = None


def _validate_command(command: str, policy: BashToolPolicy) -> Optional[str]:
  """Validates a bash command against the permitted prefixes."""
  stripped = command.strip()
  if not stripped:
    return "Command is required."

  for op in policy.blocked_operators:
    if op in command:
      return f"Command contains blocked operator: {op}"

  if "*" in policy.allowed_command_prefixes:
    return None

  for prefix in policy.allowed_command_prefixes:
    if stripped.startswith(prefix):
      return None

  allowed = ", ".join(policy.allowed_command_prefixes)
  return f"Command blocked. Permitted prefixes are: {allowed}"


def _set_resource_limits(policy: BashToolPolicy) -> None:
  """Sets resource limits for the subprocess based on the provided policy."""
  try:
    resource.setrlimit(resource.RLIMIT_CORE, (0, 0))
    if policy.max_memory_bytes:
      resource.setrlimit(
          resource.RLIMIT_AS,
          (policy.max_memory_bytes, policy.max_memory_bytes),
      )
    if policy.max_file_size_bytes:
      resource.setrlimit(
          resource.RLIMIT_FSIZE,
          (policy.max_file_size_bytes, policy.max_file_size_bytes),
      )
    if policy.max_child_processes:
      resource.setrlimit(
          resource.RLIMIT_NPROC,
          (policy.max_child_processes, policy.max_child_processes),
      )
  except (ValueError, OSError) as e:
    logger.warning("Failed to set resource limits: %s", e)


class ExecuteBashTool(BaseTool):
  """Tool to execute a validated bash command within a workspace directory."""

  def __init__(
      self,
      *,
      workspace: pathlib.Path | None = None,
      policy: Optional[BashToolPolicy] = None,
  ):
    if workspace is None:
      workspace = pathlib.Path.cwd()
    policy = policy or BashToolPolicy()
    allowed_hint = (
        "any command"
        if "*" in policy.allowed_command_prefixes
        else (
            "commands matching prefixes:"
            f" {', '.join(policy.allowed_command_prefixes)}"
        )
    )
    super().__init__(
        name="execute_bash",
        description=(
            "Executes a bash command with the working directory set to the"
            f" workspace. Allowed: {allowed_hint}. All commands require user"
            " confirmation."
        ),
    )
    self._workspace = workspace
    self._policy = policy

  def _get_declaration(self) -> Optional[types.FunctionDeclaration]:
    return types.FunctionDeclaration(
        name=self.name,
        description=self.description,
        parameters_json_schema={
            "type": "object",
            "properties": {
                "command": {
                    "type": "string",
                    "description": "The bash command to execute.",
                },
            },
            "required": ["command"],
        },
    )

  async def run_async(
      self, *, args: dict[str, Any], tool_context: ToolContext
  ) -> Any:
    command = args.get("command")
    if not command:
      return {"error": "Command is required."}

    # Static validation.
    error = _validate_command(command, self._policy)
    if error:
      return {"error": error}

    # Always request user confirmation.
    if not tool_context.tool_confirmation:
      tool_context.request_confirmation(
          hint=f"Please approve or reject the bash command: {command}",
      )
      tool_context.actions.skip_summarization = True
      return {
          "error": (
              "This tool call requires confirmation, please approve or reject."
          )
      }
    elif not tool_context.tool_confirmation.confirmed:
      return {"error": "This tool call is rejected."}

    stdout = None
    stderr = None
    try:
      process = await asyncio.create_subprocess_exec(
          *shlex.split(command),
          cwd=str(self._workspace),
          stdout=asyncio.subprocess.PIPE,
          stderr=asyncio.subprocess.PIPE,
          start_new_session=True,
          preexec_fn=lambda: _set_resource_limits(self._policy),
      )

      try:
        stdout, stderr = await asyncio.wait_for(
            process.communicate(), timeout=self._policy.timeout_seconds
        )
      except asyncio.TimeoutError:
        try:
          if process.pid:
            os.killpg(process.pid, signal.SIGKILL)
        except ProcessLookupError:
          pass
        stdout, stderr = await process.communicate()
        return {
            "error": (
                f"Command timed out after {self._policy.timeout_seconds}"
                " seconds."
            ),
            "stdout": (
                stdout.decode(errors="replace")
                if stdout
                else "<no stdout captured>"
            ),
            "stderr": (
                stderr.decode(errors="replace")
                if stderr
                else "<no stderr captured>"
            ),
            "returncode": process.returncode,
        }
      finally:
        try:
          if process.pid:
            os.killpg(process.pid, signal.SIGKILL)
        except ProcessLookupError:
          pass
      return {
          "stdout": (
              stdout.decode(errors="replace")
              if stdout
              else "<no stdout captured>"
          ),
          "stderr": (
              stderr.decode(errors="replace")
              if stderr
              else "<no stderr captured>"
          ),
          "returncode": process.returncode,
      }
    except Exception as e:  # pylint: disable=broad-except
      logger.exception("ExecuteBashTool execution failed")

      stdout_res = (
          stdout.decode(errors="replace") if stdout else "<no stdout captured>"
      )
      stderr_res = (
          stderr.decode(errors="replace") if stderr else "<no stderr captured>"
      )

      return {
          "error": f"Execution failed: {str(e)}",
          "stdout": stdout_res,
          "stderr": stderr_res,
      }

  def _detect_error_in_response(self, response: Any) -> Optional[str]:
    """Telemetry hook: returns an error type if the response indicates an error."""
    if isinstance(response, dict) and response.get("error"):
      return "TOOL_ERROR"
    return None
