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

import ntpath
import os
from pathlib import Path
from textwrap import dedent
from typing import Literal
from typing import Type
from unittest import mock

from google.adk.agents import config_agent_utils
from google.adk.agents.agent_config import AgentConfig
from google.adk.agents.base_agent import BaseAgent
from google.adk.agents.base_agent_config import BaseAgentConfig
from google.adk.agents.common_configs import AgentRefConfig
from google.adk.agents.llm_agent import LlmAgent
from google.adk.agents.loop_agent import LoopAgent
from google.adk.agents.parallel_agent import ParallelAgent
from google.adk.agents.sequential_agent import SequentialAgent
from google.adk.models.lite_llm import LiteLlm
import pytest
import yaml


def test_agent_config_discriminator_default_is_llm_agent(tmp_path: Path):
  yaml_content = """\
name: search_agent
model: gemini-2.5-flash
description: a sample description
instruction: a fake instruction
tools:
  - name: google_search
"""
  config_file = tmp_path / "test_config.yaml"
  config_file.write_text(yaml_content)

  config = AgentConfig.model_validate(yaml.safe_load(yaml_content))
  agent = config_agent_utils.from_config(str(config_file))

  assert isinstance(agent, LlmAgent)
  assert config.root.agent_class == "LlmAgent"


@pytest.mark.parametrize(
    "agent_class_value",
    [
        "LlmAgent",
        "google.adk.agents.LlmAgent",
        "google.adk.agents.llm_agent.LlmAgent",
    ],
)
def test_agent_config_discriminator_llm_agent(
    agent_class_value: str, tmp_path: Path
):
  yaml_content = f"""\
agent_class: {agent_class_value}
name: search_agent
model: gemini-2.5-flash
description: a sample description
instruction: a fake instruction
tools:
  - name: google_search
"""
  config_file = tmp_path / "test_config.yaml"
  config_file.write_text(yaml_content)

  config = AgentConfig.model_validate(yaml.safe_load(yaml_content))
  agent = config_agent_utils.from_config(str(config_file))

  assert isinstance(agent, LlmAgent)
  assert config.root.agent_class == agent_class_value


@pytest.mark.parametrize(
    "agent_class_value",
    [
        "LoopAgent",
        "google.adk.agents.LoopAgent",
        "google.adk.agents.loop_agent.LoopAgent",
    ],
)
def test_agent_config_discriminator_loop_agent(
    agent_class_value: str, tmp_path: Path
):
  yaml_content = f"""\
agent_class: {agent_class_value}
name: CodePipelineAgent
description: Executes a sequence of code writing, reviewing, and refactoring.
sub_agents: []
"""
  config_file = tmp_path / "test_config.yaml"
  config_file.write_text(yaml_content)

  config = AgentConfig.model_validate(yaml.safe_load(yaml_content))
  agent = config_agent_utils.from_config(str(config_file))

  assert isinstance(agent, LoopAgent)
  assert config.root.agent_class == agent_class_value


@pytest.mark.parametrize(
    "agent_class_value",
    [
        "ParallelAgent",
        "google.adk.agents.ParallelAgent",
        "google.adk.agents.parallel_agent.ParallelAgent",
    ],
)
def test_agent_config_discriminator_parallel_agent(
    agent_class_value: str, tmp_path: Path
):
  yaml_content = f"""\
agent_class: {agent_class_value}
name: CodePipelineAgent
description: Executes a sequence of code writing, reviewing, and refactoring.
sub_agents: []
"""
  config_file = tmp_path / "test_config.yaml"
  config_file.write_text(yaml_content)

  config = AgentConfig.model_validate(yaml.safe_load(yaml_content))
  agent = config_agent_utils.from_config(str(config_file))

  assert isinstance(agent, ParallelAgent)
  assert config.root.agent_class == agent_class_value


@pytest.mark.parametrize(
    "agent_class_value",
    [
        "SequentialAgent",
        "google.adk.agents.SequentialAgent",
        "google.adk.agents.sequential_agent.SequentialAgent",
    ],
)
def test_agent_config_discriminator_sequential_agent(
    agent_class_value: str, tmp_path: Path
):
  yaml_content = f"""\
agent_class: {agent_class_value}
name: CodePipelineAgent
description: Executes a sequence of code writing, reviewing, and refactoring.
sub_agents: []
"""
  config_file = tmp_path / "test_config.yaml"
  config_file.write_text(yaml_content)

  config = AgentConfig.model_validate(yaml.safe_load(yaml_content))
  agent = config_agent_utils.from_config(str(config_file))

  assert isinstance(agent, SequentialAgent)
  assert config.root.agent_class == agent_class_value


@pytest.mark.parametrize(
    ("agent_class_value", "expected_agent_type"),
    [
        ("LoopAgent", LoopAgent),
        ("google.adk.agents.LoopAgent", LoopAgent),
        ("google.adk.agents.loop_agent.LoopAgent", LoopAgent),
        ("ParallelAgent", ParallelAgent),
        ("google.adk.agents.ParallelAgent", ParallelAgent),
        ("google.adk.agents.parallel_agent.ParallelAgent", ParallelAgent),
        ("SequentialAgent", SequentialAgent),
        ("google.adk.agents.SequentialAgent", SequentialAgent),
        ("google.adk.agents.sequential_agent.SequentialAgent", SequentialAgent),
    ],
)
def test_agent_config_discriminator_with_sub_agents(
    agent_class_value: str, expected_agent_type: Type[BaseAgent], tmp_path: Path
):
  # Create sub-agent config files
  sub_agent_dir = tmp_path / "sub_agents"
  sub_agent_dir.mkdir()
  sub_agent_config = """\
name: sub_agent_{index}
model: gemini-2.5-flash
description: a sub agent
instruction: sub agent instruction
"""
  (sub_agent_dir / "sub_agent1.yaml").write_text(
      sub_agent_config.format(index=1)
  )
  (sub_agent_dir / "sub_agent2.yaml").write_text(
      sub_agent_config.format(index=2)
  )
  yaml_content = f"""\
agent_class: {agent_class_value}
name: main_agent
description: main agent with sub agents
sub_agents:
  - config_path: sub_agents/sub_agent1.yaml
  - config_path: sub_agents/sub_agent2.yaml
"""
  config_file = tmp_path / "test_config.yaml"
  config_file.write_text(yaml_content)

  config = AgentConfig.model_validate(yaml.safe_load(yaml_content))
  agent = config_agent_utils.from_config(str(config_file))

  assert isinstance(agent, expected_agent_type)
  assert config.root.agent_class == agent_class_value


@pytest.mark.parametrize(
    ("agent_class_value", "expected_agent_type"),
    [
        ("LlmAgent", LlmAgent),
        ("google.adk.agents.LlmAgent", LlmAgent),
        ("google.adk.agents.llm_agent.LlmAgent", LlmAgent),
    ],
)
def test_agent_config_discriminator_llm_agent_with_sub_agents(
    agent_class_value: str, expected_agent_type: Type[BaseAgent], tmp_path: Path
):
  # Create sub-agent config files
  sub_agent_dir = tmp_path / "sub_agents"
  sub_agent_dir.mkdir()
  sub_agent_config = """\
name: sub_agent_{index}
model: gemini-2.5-flash
description: a sub agent
instruction: sub agent instruction
"""
  (sub_agent_dir / "sub_agent1.yaml").write_text(
      sub_agent_config.format(index=1)
  )
  (sub_agent_dir / "sub_agent2.yaml").write_text(
      sub_agent_config.format(index=2)
  )
  yaml_content = f"""\
agent_class: {agent_class_value}
name: main_agent
model: gemini-2.5-flash
description: main agent with sub agents
instruction: main agent instruction
sub_agents:
  - config_path: sub_agents/sub_agent1.yaml
  - config_path: sub_agents/sub_agent2.yaml
"""
  config_file = tmp_path / "test_config.yaml"
  config_file.write_text(yaml_content)

  config = AgentConfig.model_validate(yaml.safe_load(yaml_content))
  agent = config_agent_utils.from_config(str(config_file))

  assert isinstance(agent, expected_agent_type)
  assert config.root.agent_class == agent_class_value


def test_agent_config_model_code_resolves_preconfigured_client(tmp_path: Path):
  """model_code references a pre-built model instance by fully qualified name.

  Configured clients (custom api_base, etc.) must be constructed in Python
  and referenced from YAML; YAML cannot pass constructor arguments.
  """
  preconfigured = LiteLlm(
      model="kimi/k2", api_base="https://proxy.litellm.ai/v1"
  )

  yaml_content = """\
name: managed_api_agent
description: Agent using LiteLLM managed endpoint
instruction: Respond concisely.
model_code:
  name: my_library.clients.my_litellm
"""
  config_file = tmp_path / "litellm_agent.yaml"
  config_file.write_text(yaml_content)

  with mock.patch.object(
      config_agent_utils,
      "resolve_code_reference",
      return_value=preconfigured,
  ):
    agent = config_agent_utils.from_config(str(config_file))

  assert isinstance(agent, LlmAgent)
  assert agent.model is preconfigured


def test_agent_config_discriminator_custom_agent():
  class MyCustomAgentConfig(BaseAgentConfig):
    agent_class: Literal["mylib.agents.MyCustomAgent"] = (
        "mylib.agents.MyCustomAgent"
    )
    other_field: str

  yaml_content = """\
agent_class: mylib.agents.MyCustomAgent
name: CodePipelineAgent
description: Executes a sequence of code writing, reviewing, and refactoring.
other_field: other value
"""
  config_data = yaml.safe_load(yaml_content)

  config = AgentConfig.model_validate(config_data)

  # pylint: disable=unidiomatic-typecheck Needs exact class matching.
  assert type(config.root) is BaseAgentConfig
  assert config.root.agent_class == "mylib.agents.MyCustomAgent"
  assert config.root.model_extra == {"other_field": "other value"}

  my_custom_config = MyCustomAgentConfig.model_validate(
      config.root.model_dump()
  )
  assert my_custom_config.other_field == "other value"


def test_from_config_passes_extra_yaml_fields_to_custom_agent_constructor(
    tmp_path: Path,
):
  """Custom agent fields in YAML reach the constructor without a custom config_type.

  Mirrors the 1.x AgentConfigMapper behavior: a custom agent subclass with
  extra Pydantic fields declared on the agent (not on a config_type) can
  populate those fields directly from YAML.
  """

  class MyCustomAgent(BaseAgent):
    custom_field: str = ""

  yaml_content = """\
agent_class: mylib.agents.MyCustomAgent
name: custom_agent
description: a custom agent
custom_field: hello from yaml
"""
  config_file = tmp_path / "custom_agent.yaml"
  config_file.write_text(yaml_content)

  with mock.patch.object(
      config_agent_utils,
      "resolve_fully_qualified_name",
      return_value=MyCustomAgent,
  ):
    agent = config_agent_utils.from_config(str(config_file))

  assert isinstance(agent, MyCustomAgent)
  assert agent.custom_field == "hello from yaml"


def test_from_config_ignores_extra_yaml_fields_not_on_agent(tmp_path: Path):
  """Extra YAML keys that don't map to constructor params are silently dropped."""

  class MyCustomAgent(BaseAgent):
    custom_field: str = ""

  yaml_content = """\
agent_class: mylib.agents.MyCustomAgent
name: custom_agent
description: a custom agent
custom_field: kept
unknown_field: dropped
"""
  config_file = tmp_path / "custom_agent.yaml"
  config_file.write_text(yaml_content)

  with mock.patch.object(
      config_agent_utils,
      "resolve_fully_qualified_name",
      return_value=MyCustomAgent,
  ):
    agent = config_agent_utils.from_config(str(config_file))

  assert agent.custom_field == "kept"
  assert not hasattr(agent, "unknown_field")


@pytest.mark.parametrize(
    ("config_rel_path", "child_rel_path", "child_name", "instruction"),
    [
        (
            Path("main.yaml"),
            Path("sub_agents/child.yaml"),
            "child_agent",
            "I am a child agent",
        ),
        (
            Path("level1/level2/nested_main.yaml"),
            Path("sub/nested_child.yaml"),
            "nested_child",
            "I am nested",
        ),
    ],
)
def test_resolve_agent_reference_resolves_relative_paths(
    config_rel_path: Path,
    child_rel_path: Path,
    child_name: str,
    instruction: str,
    tmp_path: Path,
):
  """Verify resolve_agent_reference resolves relative sub-agent paths."""
  config_file = tmp_path / config_rel_path
  config_file.parent.mkdir(parents=True, exist_ok=True)

  child_config_path = config_file.parent / child_rel_path
  child_config_path.parent.mkdir(parents=True, exist_ok=True)
  child_config_path.write_text(dedent(f"""
          agent_class: LlmAgent
          name: {child_name}
          model: gemini-2.5-flash
          instruction: {instruction}
          """).lstrip())

  config_file.write_text(dedent(f"""
          agent_class: LlmAgent
          name: main_agent
          model: gemini-2.5-flash
          instruction: I am the main agent
          sub_agents:
            - config_path: {child_rel_path.as_posix()}
          """).lstrip())

  ref_config = AgentRefConfig(config_path=child_rel_path.as_posix())
  agent = config_agent_utils.resolve_agent_reference(
      ref_config, str(config_file)
  )

  assert agent.name == child_name

  config_dir = os.path.dirname(str(config_file.resolve()))
  assert config_dir == str(config_file.parent.resolve())

  expected_child_path = os.path.join(config_dir, *child_rel_path.parts)
  assert os.path.exists(expected_child_path)


def test_resolve_agent_reference_uses_windows_dirname():
  """Ensure Windows-style config references resolve via os.path.dirname."""
  ref_config = AgentRefConfig(config_path="sub\\child.yaml")
  recorded: dict[str, str] = {}

  def fake_from_config(path: str):
    recorded["path"] = path
    return "sentinel"

  with (
      mock.patch.object(
          config_agent_utils,
          "from_config",
          autospec=True,
          side_effect=fake_from_config,
      ),
      mock.patch.object(config_agent_utils.os, "path", ntpath),
  ):
    referencing = r"C:\workspace\agents\main.yaml"
    result = config_agent_utils.resolve_agent_reference(ref_config, referencing)

  expected_path = ntpath.join(
      ntpath.dirname(referencing), ref_config.config_path
  )
  assert result == "sentinel"
  assert recorded["path"] == expected_path


def test_resolve_agent_reference_blocks_absolute_path():
  """Verify resolve_agent_reference raises ValueError for absolute paths."""
  ref_config = AgentRefConfig(config_path="/etc/passwd")
  with pytest.raises(
      ValueError,
      match="Absolute paths are not allowed in AgentRefConfig config_path",
  ):
    config_agent_utils.resolve_agent_reference(
        ref_config, "/workspace/main.yaml"
    )


def test_resolve_agent_reference_blocks_path_traversal():
  """Verify resolve_agent_reference raises ValueError for path traversal."""
  ref_config = AgentRefConfig(config_path="../outside.yaml")
  with pytest.raises(ValueError, match="Path traversal detected"):
    config_agent_utils.resolve_agent_reference(
        ref_config, "/workspace/agents/main.yaml"
    )


# --- Security tests: module blocklist for YAML agent config code references ---


def test_resolve_code_reference_blocks_os_when_enforced():
  """Verify resolve_code_reference blocks os module directly."""
  from google.adk.agents.common_configs import CodeConfig

  with pytest.raises(ValueError, match="Blocked module reference"):
    config_agent_utils.resolve_code_reference(CodeConfig(name="os.system"))


def test_resolve_fully_qualified_name_blocks_subprocess_when_enforced():
  """Verify resolve_fully_qualified_name blocks subprocess module.

  resolve_fully_qualified_name wraps all exceptions in
  ValueError("Invalid fully qualified name: ..."), so we check the wrapper
  and verify the __cause__ carries the blocklist message.
  """
  with pytest.raises(
      ValueError, match="Invalid fully qualified name"
  ) as exc_info:
    config_agent_utils.resolve_fully_qualified_name("subprocess.Popen")
  assert "Blocked module reference" in str(exc_info.value.__cause__)


def test_allowed_module_passes_when_enforced(tmp_path: Path):
  """Verify that google.adk modules are NOT blocked by the module denylist."""
  # This should NOT raise — google.adk modules must remain allowed
  result = config_agent_utils.resolve_fully_qualified_name(
      "google.adk.agents.llm_agent.LlmAgent"
  )
  assert result is LlmAgent


@pytest.mark.parametrize(
    "blocked_module",
    [
        "os.system",
        "subprocess.call",
        "builtins.exec",
    ],
)
def test_resolve_agent_code_reference_blocks_when_enforced(
    blocked_module: str,
):
  """Verify _resolve_agent_code_reference blocks dangerous modules."""
  with pytest.raises(ValueError, match="Blocked module reference"):
    config_agent_utils._resolve_agent_code_reference(blocked_module)


@pytest.mark.parametrize(
    "blocked_ref",
    [
        "os.system",
        "subprocess.call",
        "builtins.exec",
        "pickle.loads",
    ],
)
def test_resolve_tools_blocks_dangerous_modules(blocked_ref: str):
  """Verify _resolve_tools blocks dangerous modules for user-defined tools."""
  from google.adk.agents.llm_agent import LlmAgent
  from google.adk.tools.tool_configs import ToolConfig

  tool_config = ToolConfig(name=blocked_ref)
  with pytest.raises(ValueError, match="Blocked module reference"):
    LlmAgent._resolve_tools([tool_config], "/fake/path.yaml")


def test_resolve_tools_allows_builtin_adk_tools():
  """Verify _resolve_tools allows ADK built-in tools (no dot in name)."""
  from google.adk.agents.llm_agent import LlmAgent
  from google.adk.tools.tool_configs import ToolConfig

  # Built-in tools have no dot — they import from google.adk.tools
  tool_config = ToolConfig(name="google_search")
  # Should NOT raise — this is a safe, hardcoded import path
  resolved = LlmAgent._resolve_tools([tool_config], "/fake/path.yaml")
  assert len(resolved) == 1


@pytest.mark.parametrize(
    "blocked_ref",
    [
        "ftplib.FTP",
        "smtplib.SMTP",
        "xmlrpc.client",
        "telnetlib.Telnet",
        "poplib.POP3",
        "imaplib.IMAP4",
        "asyncio.run",
        "pathlib.Path",
    ],
)
def test_newly_blocked_network_modules_are_rejected(blocked_ref: str):
  """Verify newly added network-capable modules are blocked.

  resolve_fully_qualified_name wraps errors, so we check the cause.
  """
  with pytest.raises(
      ValueError, match="Invalid fully qualified name"
  ) as exc_info:
    config_agent_utils.resolve_fully_qualified_name(blocked_ref)
  assert "Blocked module reference" in str(exc_info.value.__cause__)


def test_denylist_can_be_disabled():
  """Verify _set_enforce_denylist(False) disables module blocking."""
  config_agent_utils._set_enforce_denylist(False)
  try:
    # os.getcwd is a real, importable reference — should succeed
    result = config_agent_utils.resolve_fully_qualified_name("os.getcwd")
    assert callable(result)
  finally:
    config_agent_utils._set_enforce_denylist(True)


def test_load_config_from_path_blocks_args_when_enforced(tmp_path: Path):
  """_load_config_from_path blocks the 'args' key when enforcement is on."""
  config_file = tmp_path / "agent.yaml"
  config_file.write_text("name: my_agent\nargs:\n  key: value\n")
  config_agent_utils._set_enforce_yaml_key_denylist(True)
  try:
    with pytest.raises(ValueError) as exc_info:
      config_agent_utils._load_config_from_path(str(config_file))
    assert "Blocked key 'args' found" in str(exc_info.value)
  finally:
    config_agent_utils._set_enforce_yaml_key_denylist(False)