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

from unittest.mock import Mock

from google.adk.agents.base_agent import BaseAgent
from google.adk.agents.base_agent import BaseAgentState
from google.adk.agents.invocation_context import InvocationContext
from google.adk.apps import ResumabilityConfig
from google.adk.events.event import Event
from google.adk.events.event_actions import EventActions
from google.adk.sessions.base_session_service import BaseSessionService
from google.adk.sessions.session import Session
from google.genai.types import Content
from google.genai.types import FunctionCall
from google.genai.types import Part
import pytest

from .. import testing_utils


class TestInvocationContext:
  """Test suite for InvocationContext."""

  @pytest.fixture
  def mock_events(self):
    """Create mock events for testing."""
    event1 = Mock(spec=Event)
    event1.invocation_id = 'inv_1'
    event1.branch = 'agent_1'

    event2 = Mock(spec=Event)
    event2.invocation_id = 'inv_1'
    event2.branch = 'agent_2'

    event3 = Mock(spec=Event)
    event3.invocation_id = 'inv_2'
    event3.branch = 'agent_1'

    event4 = Mock(spec=Event)
    event4.invocation_id = 'inv_2'
    event4.branch = 'agent_2'

    return [event1, event2, event3, event4]

  @pytest.fixture
  def mock_invocation_context(self, mock_events):
    """Create a mock invocation context for testing."""
    ctx = InvocationContext(
        session_service=Mock(spec=BaseSessionService),
        agent=Mock(spec=BaseAgent),
        invocation_id='inv_1',
        branch='agent_1',
        session=Mock(spec=Session, events=mock_events),
    )
    return ctx

  def test_get_events_returns_all_events_by_default(
      self, mock_invocation_context, mock_events
  ):
    """Tests that get_events returns all events when no filters are applied."""
    events = mock_invocation_context._get_events()
    assert events == mock_events

  def test_get_events_filters_by_current_invocation(
      self, mock_invocation_context, mock_events
  ):
    """Tests that get_events correctly filters by the current invocation."""
    event1, event2, _, _ = mock_events
    events = mock_invocation_context._get_events(current_invocation=True)
    assert events == [event1, event2]

  def test_get_events_filters_by_current_branch(
      self, mock_invocation_context, mock_events
  ):
    """Tests that get_events correctly filters by the current branch."""
    event1, _, event3, _ = mock_events
    events = mock_invocation_context._get_events(current_branch=True)
    assert events == [event1, event3]

  def test_get_events_filters_by_invocation_and_branch(
      self, mock_invocation_context, mock_events
  ):
    """Tests that get_events filters by invocation and branch."""
    event1, _, _, _ = mock_events
    events = mock_invocation_context._get_events(
        current_invocation=True,
        current_branch=True,
    )
    assert events == [event1]

  def test_get_events_with_no_events_in_session(self, mock_invocation_context):
    """Tests get_events when the session has no events."""
    mock_invocation_context.session.events = []
    events = mock_invocation_context._get_events()
    assert not events

  def test_get_events_with_no_matching_events(self, mock_invocation_context):
    """Tests get_events when no events match the filters."""
    mock_invocation_context.invocation_id = 'inv_3'
    mock_invocation_context.branch = 'branch_C'

    # Filter by invocation
    events = mock_invocation_context._get_events(current_invocation=True)
    assert not events

    # Filter by branch
    events = mock_invocation_context._get_events(current_branch=True)
    assert not events

    # Filter by both
    events = mock_invocation_context._get_events(
        current_invocation=True,
        current_branch=True,
    )
    assert not events


class TestInvocationContextWithAppResumablity:
  """Test suite for InvocationContext regarding app resumability."""

  @pytest.fixture
  def long_running_function_call(self) -> FunctionCall:
    """A long running function call."""
    return FunctionCall(
        id='tool_call_id_1',
        name='long_running_function_call',
        args={},
    )

  @pytest.fixture
  def event_to_pause(self, long_running_function_call) -> Event:
    """An event with a long running function call."""
    return Event(
        invocation_id='inv_1',
        author='agent',
        content=testing_utils.ModelContent(
            [Part(function_call=long_running_function_call)]
        ),
        long_running_tool_ids=[long_running_function_call.id],
    )

  def _create_test_invocation_context(
      self, resumability_config: ResumabilityConfig | None = None
  ) -> InvocationContext:
    """Create a mock invocation context for testing."""
    ctx = InvocationContext(
        session_service=Mock(spec=BaseSessionService),
        agent=Mock(spec=BaseAgent),
        invocation_id='inv_1',
        session=Mock(spec=Session, events=[]),
        resumability_config=resumability_config,
    )
    return ctx

  def test_should_pause_invocation_with_resumable_app(self, event_to_pause):
    """Tests should_pause_invocation with a resumable app."""
    mock_invocation_context = self._create_test_invocation_context(
        ResumabilityConfig(is_resumable=True)
    )

    assert mock_invocation_context.should_pause_invocation(event_to_pause)

  def test_should_pause_invocation_with_non_resumable_app(self, event_to_pause):
    """Tests should_pause_invocation pauses even without resumability."""
    invocation_context = self._create_test_invocation_context(
        ResumabilityConfig(is_resumable=False)
    )

    assert invocation_context.should_pause_invocation(event_to_pause)

  def test_should_not_pause_invocation_with_no_long_running_tool_ids(
      self, event_to_pause
  ):
    """Tests should_pause_invocation with no long running tools."""
    invocation_context = self._create_test_invocation_context(
        ResumabilityConfig(is_resumable=True)
    )
    nonpausable_event = event_to_pause.model_copy(
        update={'long_running_tool_ids': []}
    )

    assert not invocation_context.should_pause_invocation(nonpausable_event)

  def test_should_not_pause_invocation_with_no_function_calls(
      self, event_to_pause
  ):
    """Tests should_pause_invocation with a non-model event."""
    mock_invocation_context = self._create_test_invocation_context(
        ResumabilityConfig(is_resumable=True)
    )
    nonpausable_event = event_to_pause.model_copy(
        update={'content': testing_utils.UserContent('test text part')}
    )

    assert not mock_invocation_context.should_pause_invocation(
        nonpausable_event
    )

  def test_should_not_pause_when_user_resumes_in_sub_branch(
      self, event_to_pause, long_running_function_call
  ):
    """We do not pause the invocation if a subsequent user event belongs to a sub-branch."""
    # Arrange
    mock_invocation_context = self._create_test_invocation_context()
    user_event = Event(
        invocation_id='inv_1',
        author='user',
        branch=f'agent@{long_running_function_call.id}.child',
    )
    mock_invocation_context.session.events = [event_to_pause, user_event]

    # Act
    should_pause = mock_invocation_context.should_pause_invocation(
        event_to_pause
    )

    # Assert
    assert not should_pause

  def test_should_not_pause_when_user_resumes_in_deeply_nested_sub_branch(
      self, event_to_pause, long_running_function_call
  ):
    """We do not pause if the user resumes in a deeply nested sub-branch containing the tool call."""
    # Arrange
    mock_invocation_context = self._create_test_invocation_context()
    user_event = Event(
        invocation_id='inv_1',
        author='user',
        branch=f'parent@other.child@{long_running_function_call.id}.grandchild',
    )
    mock_invocation_context.session.events = [event_to_pause, user_event]

    # Act
    should_pause = mock_invocation_context.should_pause_invocation(
        event_to_pause
    )

    # Assert
    assert not should_pause

  def test_should_pause_when_user_resumes_in_different_branch(
      self, event_to_pause
  ):
    """We still pause the invocation if the subsequent user event belongs to a different branch."""
    # Arrange
    mock_invocation_context = self._create_test_invocation_context()
    user_event = Event(
        invocation_id='inv_1',
        author='user',
        branch='parent@different_id.child',
    )
    mock_invocation_context.session.events = [event_to_pause, user_event]

    # Act
    should_pause = mock_invocation_context.should_pause_invocation(
        event_to_pause
    )

    # Assert
    assert should_pause

  def test_is_resumable_true(self):
    """Tests that is_resumable is True when resumability is enabled."""
    invocation_context = self._create_test_invocation_context(
        ResumabilityConfig(is_resumable=True)
    )
    assert invocation_context.is_resumable

  def test_is_resumable_false(self):
    """Tests that is_resumable is False when resumability is disabled."""
    invocation_context = self._create_test_invocation_context(
        ResumabilityConfig(is_resumable=False)
    )
    assert not invocation_context.is_resumable

  def test_is_resumable_no_config(self):
    """Tests that is_resumable is False when no resumability config is set."""
    invocation_context = self._create_test_invocation_context(None)
    assert not invocation_context.is_resumable

  def test_populate_invocation_agent_states_not_resumable(self):
    """Tests that populate_invocation_agent_states does nothing if not resumable."""
    invocation_context = self._create_test_invocation_context(
        ResumabilityConfig(is_resumable=False)
    )
    event = Event(
        invocation_id='inv_1',
        author='agent1',
        actions=EventActions(end_of_agent=True, agent_state=None),
    )
    invocation_context.session.events = [event]
    invocation_context.populate_invocation_agent_states()
    assert not invocation_context.agent_states
    assert not invocation_context.end_of_agents

  def test_populate_invocation_agent_states_end_of_agent(self):
    """Tests that populate_invocation_agent_states handles end_of_agent."""
    invocation_context = self._create_test_invocation_context(
        ResumabilityConfig(is_resumable=True)
    )
    event = Event(
        invocation_id='inv_1',
        author='agent1',
        actions=EventActions(end_of_agent=True, agent_state=None),
    )
    invocation_context.session.events = [event]
    invocation_context.populate_invocation_agent_states()
    assert not invocation_context.agent_states
    assert invocation_context.end_of_agents == {'agent1': True}

  def test_populate_invocation_agent_states_with_agent_state(self):
    """Tests that populate_invocation_agent_states handles agent_state."""
    invocation_context = self._create_test_invocation_context(
        ResumabilityConfig(is_resumable=True)
    )
    event = Event(
        invocation_id='inv_1',
        author='agent1',
        actions=EventActions(
            end_of_agent=False,
            agent_state=BaseAgentState().model_dump(mode='json'),
        ),
    )
    invocation_context.session.events = [event]
    invocation_context.populate_invocation_agent_states()
    assert invocation_context.agent_states == {'agent1': {}}
    assert invocation_context.end_of_agents == {'agent1': False}

  def test_populate_invocation_agent_states_with_agent_state_and_end_of_agent(
      self,
  ):
    """Tests that populate_invocation_agent_states handles agent_state and end_of_agent."""
    invocation_context = self._create_test_invocation_context(
        ResumabilityConfig(is_resumable=True)
    )
    event = Event(
        invocation_id='inv_1',
        author='agent1',
        actions=EventActions(
            end_of_agent=True,
            agent_state=BaseAgentState().model_dump(mode='json'),
        ),
    )
    invocation_context.session.events = [event]
    invocation_context.populate_invocation_agent_states()
    # When both agent_state and end_of_agent are set, agent_state should be
    # cleared, as end_of_agent is of a higher priority.
    assert not invocation_context.agent_states
    assert invocation_context.end_of_agents == {'agent1': True}

  def test_populate_invocation_agent_states_with_content_no_state(self):
    """Tests that populate_invocation_agent_states creates default state."""
    invocation_context = self._create_test_invocation_context(
        ResumabilityConfig(is_resumable=True)
    )
    event = Event(
        invocation_id='inv_1',
        author='agent1',
        actions=EventActions(end_of_agent=False, agent_state=None),
        content=Content(role='model', parts=[Part(text='hi')]),
    )
    invocation_context.session.events = [event]
    invocation_context.populate_invocation_agent_states()
    assert invocation_context.agent_states == {
        'agent1': BaseAgentState().model_dump(mode='json')
    }
    assert invocation_context.end_of_agents == {'agent1': False}

  def test_populate_invocation_agent_states_user_message_event(self):
    """Tests that populate_invocation_agent_states ignores user message events for default state."""
    invocation_context = self._create_test_invocation_context(
        ResumabilityConfig(is_resumable=True)
    )
    event = Event(
        invocation_id='inv_1',
        author='user',
        actions=EventActions(end_of_agent=False, agent_state=None),
        content=Content(role='user', parts=[Part(text='hi')]),
    )
    invocation_context.session.events = [event]
    invocation_context.populate_invocation_agent_states()
    assert not invocation_context.agent_states
    assert not invocation_context.end_of_agents

  def test_populate_invocation_agent_states_no_content(self):
    """Tests that populate_invocation_agent_states ignores events with no content if no state."""
    invocation_context = self._create_test_invocation_context(
        ResumabilityConfig(is_resumable=True)
    )
    event = Event(
        invocation_id='inv_1',
        author='agent1',
        actions=EventActions(end_of_agent=None, agent_state=None),
        content=None,
    )
    invocation_context.session.events = [event]
    invocation_context.populate_invocation_agent_states()
    assert not invocation_context.agent_states
    assert not invocation_context.end_of_agents

  def test_set_agent_state_with_end_of_agent_true(self):
    """Tests that set_agent_state clears agent_state and sets end_of_agent to True."""
    invocation_context = self._create_test_invocation_context(
        ResumabilityConfig(is_resumable=True)
    )
    invocation_context.agent_states['agent1'] = {}
    invocation_context.end_of_agents['agent1'] = False

    # Set state with end_of_agent=True, which should clear the existing
    # agent_state.
    invocation_context.set_agent_state('agent1', end_of_agent=True)
    assert 'agent1' not in invocation_context.agent_states
    assert invocation_context.end_of_agents['agent1']

  def test_set_agent_state_with_agent_state(self):
    """Tests that set_agent_state sets agent_state and sets end_of_agent to False."""
    agent_state = BaseAgentState()
    invocation_context = self._create_test_invocation_context(
        ResumabilityConfig(is_resumable=True)
    )
    invocation_context.end_of_agents['agent1'] = True

    # Set state with agent_state=agent_state, which should set the agent_state
    # and reset the end_of_agent flag to False.
    invocation_context.set_agent_state('agent1', agent_state=agent_state)
    assert invocation_context.agent_states['agent1'] == agent_state.model_dump(
        mode='json'
    )
    assert invocation_context.end_of_agents['agent1'] is False

  def test_reset_agent_state(self):
    """Tests that set_agent_state clears agent_state and end_of_agent."""
    invocation_context = self._create_test_invocation_context(
        ResumabilityConfig(is_resumable=True)
    )
    invocation_context.agent_states['agent1'] = {}
    invocation_context.end_of_agents['agent1'] = True

    # Reset state, which should clear the agent_state and end_of_agent flag.
    invocation_context.set_agent_state('agent1')
    assert 'agent1' not in invocation_context.agent_states
    assert 'agent1' not in invocation_context.end_of_agents

  def test_reset_sub_agent_states(self):
    """Tests that reset_sub_agent_states resets sub-agent states."""
    sub_sub_agent_1 = BaseAgent(name='sub_sub_agent_1')
    sub_agent_1 = BaseAgent(name='sub_agent_1', sub_agents=[sub_sub_agent_1])
    sub_agent_2 = BaseAgent(name='sub_agent_2')
    root_agent = BaseAgent(
        name='root_agent', sub_agents=[sub_agent_1, sub_agent_2]
    )

    invocation_context = self._create_test_invocation_context(
        ResumabilityConfig(is_resumable=True)
    )
    invocation_context.agent = root_agent
    invocation_context.set_agent_state(
        'sub_agent_1', agent_state=BaseAgentState()
    )
    invocation_context.set_agent_state('sub_agent_2', end_of_agent=True)
    invocation_context.set_agent_state(
        'sub_sub_agent_1', agent_state=BaseAgentState()
    )

    assert 'sub_agent_1' in invocation_context.agent_states
    assert 'sub_agent_2' in invocation_context.end_of_agents
    assert 'sub_sub_agent_1' in invocation_context.agent_states

    invocation_context.reset_sub_agent_states('root_agent')

    assert 'sub_agent_1' not in invocation_context.agent_states
    assert 'sub_agent_1' not in invocation_context.end_of_agents
    assert 'sub_agent_2' not in invocation_context.agent_states
    assert 'sub_agent_2' not in invocation_context.end_of_agents
    assert 'sub_sub_agent_1' not in invocation_context.agent_states
    assert 'sub_sub_agent_1' not in invocation_context.end_of_agents


class TestFindMatchingFunctionCall:
  """Test suite for find_matching_function_call."""

  @pytest.fixture
  def test_invocation_context(self):
    """Create a mock invocation context for testing."""

    def _create_invocation_context(events):
      return InvocationContext(
          session_service=Mock(spec=BaseSessionService),
          agent=Mock(spec=BaseAgent, name='agent'),
          invocation_id='inv_1',
          session=Mock(spec=Session, events=events),
      )

    return _create_invocation_context

  def test_find_matching_function_call_found(self, test_invocation_context):
    """Tests that a matching function call is found."""
    fc = Part.from_function_call(name='some_tool', args={})
    fc.function_call.id = 'test_function_call_id'
    fc_event = Event(
        invocation_id='inv_1',
        author='agent',
        content=testing_utils.ModelContent([fc]),
    )
    fr = Part.from_function_response(
        name='some_tool', response={'result': 'ok'}
    )
    fr.function_response.id = 'test_function_call_id'
    fr_event = Event(
        invocation_id='inv_1',
        author='agent',
        content=Content(role='user', parts=[fr]),
    )
    invocation_context = test_invocation_context([fc_event, fr_event])
    matching_fc_event = invocation_context._find_matching_function_call(
        fr_event
    )
    assert testing_utils.simplify_content(
        matching_fc_event.content
    ) == testing_utils.simplify_content(fc_event.content)

  def test_find_matching_function_call_not_found(self, test_invocation_context):
    """Tests that no matching function call is returned if id doesn't match."""
    fc = Part.from_function_call(name='some_tool', args={})
    fc.function_call.id = 'another_function_call_id'
    fc_event = Event(
        invocation_id='inv_1',
        author='agent',
        content=testing_utils.ModelContent([fc]),
    )
    fr = Part.from_function_response(
        name='some_tool', response={'result': 'ok'}
    )
    fr.function_response.id = 'test_function_call_id'
    fr_event = Event(
        invocation_id='inv_1',
        author='agent',
        content=Content(role='user', parts=[fr]),
    )
    invocation_context = test_invocation_context([fc_event, fr_event])
    match = invocation_context._find_matching_function_call(fr_event)
    assert match is None

  def test_find_matching_function_call_no_call_events(
      self, test_invocation_context
  ):
    """Tests that no matching function call is returned if there are no call events."""
    fr = Part.from_function_response(
        name='some_tool', response={'result': 'ok'}
    )
    fr.function_response.id = 'test_function_call_id'
    fr_event = Event(
        invocation_id='inv_1',
        author='agent',
        content=Content(role='user', parts=[fr]),
    )
    invocation_context = test_invocation_context([fr_event])
    match = invocation_context._find_matching_function_call(fr_event)
    assert match is None

  def test_find_matching_function_call_no_response_in_event(
      self, test_invocation_context
  ):
    """Tests result is None if function_response_event has no function response."""
    fr_event_no_fr = Event(
        author='agent',
        content=Content(role='user', parts=[Part(text='user message')]),
    )
    fc = Part.from_function_call(name='some_tool', args={})
    fc.function_call.id = 'test_function_call_id'
    fc_event = Event(
        invocation_id='inv_1',
        author='agent',
        content=testing_utils.ModelContent([fc]),
    )
    fr = Part.from_function_response(
        name='some_tool', response={'result': 'ok'}
    )
    fr.function_response.id = 'test_function_call_id'
    fr_event = Event(
        invocation_id='inv_1',
        author='agent',
        content=Content(role='user', parts=[Part(text='user message')]),
    )
    invocation_context = test_invocation_context([fc_event, fr_event])
    match = invocation_context._find_matching_function_call(fr_event_no_fr)
    assert match is None

  def test_stamp_event_branch_context_preserves_isolation_scope(
      self, test_invocation_context
  ):
    """Tests stamp_event_branch_context does not overwrite existing isolation_scope with None."""
    fc = Part.from_function_call(name='some_tool', args={})
    fc.function_call.id = 'test_function_call_id'
    fc_event = Event(
        invocation_id='inv_1',
        author='agent',
        branch='root@1',
        isolation_scope=None,  # Coordinator FC has None scope
        content=testing_utils.ModelContent([fc]),
    )
    fr = Part.from_function_response(
        name='some_tool', response={'result': 'ok'}
    )
    fr.function_response.id = 'test_function_call_id'
    fr_event = Event(
        invocation_id='inv_1',
        author='agent',
        isolation_scope='task_123',  # Pre-populated active task scope
        content=Content(role='user', parts=[fr]),
    )
    invocation_context = test_invocation_context([fc_event, fr_event])

    invocation_context.stamp_event_branch_context(fr_event)
    assert fr_event.branch == 'root@1'
    assert fr_event.isolation_scope == 'task_123'

  def test_stamp_event_branch_context_does_not_overwrite_existing_scope(
      self, test_invocation_context
  ):
    """Tests stamp_event_branch_context does not overwrite existing isolation_scope if set."""
    fc = Part.from_function_call(name='some_tool', args={})
    fc.function_call.id = 'test_function_call_id'
    fc_event = Event(
        invocation_id='inv_1',
        author='agent',
        branch='root@1',
        isolation_scope='task_456',  # Function call has isolation scope
        content=testing_utils.ModelContent([fc]),
    )
    fr = Part.from_function_response(
        name='some_tool', response={'result': 'ok'}
    )
    fr.function_response.id = 'test_function_call_id'
    fr_event = Event(
        invocation_id='inv_1',
        author='agent',
        isolation_scope='task_123',  # Pre-populated active task scope
        content=Content(role='user', parts=[fr]),
    )
    invocation_context = test_invocation_context([fc_event, fr_event])

    invocation_context.stamp_event_branch_context(fr_event)
    assert fr_event.branch == 'root@1'
    assert fr_event.isolation_scope == 'task_123'

  def test_find_matching_function_call_when_response_is_not_last_event(
      self, test_invocation_context
  ):
    """Tests that matching function call is found even when response is not the last event in history."""
    fc = Part.from_function_call(name='some_tool', args={})
    fc.function_call.id = 'test_function_call_id'
    fc_event = Event(
        invocation_id='inv_1',
        author='agent',
        content=testing_utils.ModelContent([fc]),
    )
    fr = Part.from_function_response(
        name='some_tool', response={'result': 'ok'}
    )
    fr.function_response.id = 'test_function_call_id'
    fr_event = Event(
        invocation_id='inv_1',
        author='agent',
        content=Content(role='user', parts=[fr]),
    )
    # Add a subsequent event to the history so that fr_event is NOT the last one
    subsequent_event = Event(
        invocation_id='inv_1',
        author='user',
        content=Content(role='user', parts=[Part(text='next user message')]),
    )
    invocation_context = test_invocation_context(
        [fc_event, fr_event, subsequent_event]
    )

    matching_fc_event = invocation_context._find_matching_function_call(
        fr_event
    )
    assert testing_utils.simplify_content(
        matching_fc_event.content
    ) == testing_utils.simplify_content(fc_event.content)