# 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 asyncio
from contextlib import aclosing
from typing import Any
from typing import AsyncGenerator
from typing import Awaitable

from google.adk.agents.live_request_queue import LiveRequestQueue
from google.adk.agents.llm_agent import Agent
from google.adk.models.llm_response import LlmResponse
from google.genai import types
import pytest

from .. import testing_utils


async def _wait_for_queue_empty(queue: LiveRequestQueue):
  """Wait until the queue is empty and the background consumer has finished."""
  while not queue._queue.empty():
    await asyncio.sleep(0)
  # Give opportunity for _send_to_model to finish processing (e.g. append_event)
  for _ in range(10):
    await asyncio.sleep(0)


class StreamingTestRunner(testing_utils.InMemoryRunner):
  """A robust runner for streaming tests that avoids resource leaks."""

  def __init__(self, *args, max_responses=3, **kwargs):
    super().__init__(*args, **kwargs)
    self.max_responses = max_responses

  def _run_with_loop(self, coro):
    try:
      old_loop = asyncio.get_event_loop()
    except RuntimeError:
      old_loop = None

    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
      loop.run_until_complete(coro)
    except (asyncio.TimeoutError, asyncio.CancelledError):
      pass
    finally:
      # Cancel all pending tasks to prevent leaks and warnings
      pending = asyncio.all_tasks(loop)
      for task in pending:
        task.cancel()
      if pending:
        try:
          loop.run_until_complete(
              asyncio.gather(*pending, return_exceptions=True)
          )
        except Exception:  # pylint: disable=broad-except
          pass
      loop.close()
      asyncio.set_event_loop(old_loop)

  def run_live(
      self,
      live_request_queue: LiveRequestQueue,
      run_config: testing_utils.RunConfig = None,
  ) -> list[testing_utils.Event]:
    collected_responses = []

    async def consume_responses(session: testing_utils.Session):
      run_res = self.runner.run_live(
          session=session,
          live_request_queue=live_request_queue,
          run_config=run_config or testing_utils.RunConfig(),
      )

      async with aclosing(run_res) as agen:
        async for response in agen:
          collected_responses.append(response)
          if len(collected_responses) >= self.max_responses:
            await _wait_for_queue_empty(live_request_queue)
            return

    self._run_with_loop(
        asyncio.wait_for(consume_responses(self.session), timeout=5.0)
    )

    return collected_responses

  def run_live_and_get_session(
      self,
      live_request_queue: LiveRequestQueue,
      run_config: testing_utils.RunConfig = None,
  ) -> tuple[list[testing_utils.Event], testing_utils.Session]:
    events = self.run_live(live_request_queue, run_config)
    return events, self.session


def test_streaming():
  response1 = LlmResponse(
      turn_complete=True,
  )

  mock_model = testing_utils.MockModel.create([response1])

  root_agent = Agent(
      name="root_agent",
      model=mock_model,
      tools=[],
  )

  runner = testing_utils.InMemoryRunner(
      root_agent=root_agent, response_modalities=["AUDIO"]
  )
  live_request_queue = LiveRequestQueue()
  live_request_queue.send_realtime(
      blob=types.Blob(data=b"\x00\xFF", mime_type="audio/pcm")
  )
  res_events = runner.run_live(live_request_queue)

  assert res_events is not None, "Expected a list of events, got None."
  assert (
      len(res_events) > 0
  ), "Expected at least one response, but got an empty list."


def test_live_streaming_function_call_single():
  """Test live streaming with a single function call response."""
  # Create a function call response
  function_call = types.Part.from_function_call(
      name="get_weather", args={"location": "San Francisco", "unit": "celsius"}
  )

  # Create LLM responses: function call followed by turn completion
  response1 = LlmResponse(
      content=types.Content(role="model", parts=[function_call]),
      turn_complete=False,
  )
  response2 = LlmResponse(
      turn_complete=True,
  )

  mock_model = testing_utils.MockModel.create([response1, response2])

  # Mock function that would be called
  def get_weather(location: str, unit: str = "celsius") -> dict:
    return {
        "temperature": 22,
        "condition": "sunny",
        "location": location,
        "unit": unit,
    }

  root_agent = Agent(
      name="root_agent",
      model=mock_model,
      tools=[get_weather],
  )

  runner = StreamingTestRunner(root_agent=root_agent, max_responses=3)

  live_request_queue = LiveRequestQueue()
  live_request_queue.send_realtime(
      blob=types.Blob(
          data=b"What is the weather in San Francisco?", mime_type="audio/pcm"
      )
  )

  res_events = runner.run_live(live_request_queue)

  assert res_events is not None, "Expected a list of events, got None."
  assert len(res_events) >= 1, "Expected at least one event."

  # Check that we got a function call event
  function_call_found = False
  function_response_found = False

  for event in res_events:
    if event.content and event.content.parts:
      for part in event.content.parts:
        if part.function_call and part.function_call.name == "get_weather":
          function_call_found = True
          assert part.function_call.args["location"] == "San Francisco"
          assert part.function_call.args["unit"] == "celsius"
        elif (
            part.function_response
            and part.function_response.name == "get_weather"
        ):
          function_response_found = True
          assert part.function_response.response["temperature"] == 22
          assert part.function_response.response["condition"] == "sunny"

  assert function_call_found, "Expected a function call event."
  # Note: In live streaming, function responses might be handled differently,
  # so we check for the function call which is the primary indicator of function calling working


def test_live_streaming_function_call_multiple():
  """Test live streaming with multiple function calls in sequence."""
  # Create multiple function call responses
  function_call1 = types.Part.from_function_call(
      name="get_weather", args={"location": "San Francisco"}
  )
  function_call2 = types.Part.from_function_call(
      name="get_time", args={"timezone": "PST"}
  )

  # Create LLM responses: two function calls followed by turn completion
  response1 = LlmResponse(
      content=types.Content(role="model", parts=[function_call1]),
      turn_complete=False,
  )
  response2 = LlmResponse(
      content=types.Content(role="model", parts=[function_call2]),
      turn_complete=False,
  )
  response3 = LlmResponse(
      turn_complete=True,
  )

  mock_model = testing_utils.MockModel.create([response1, response2, response3])

  # Mock functions
  def get_weather(location: str) -> dict:
    return {"temperature": 22, "condition": "sunny", "location": location}

  def get_time(timezone: str) -> dict:
    return {"time": "14:30", "timezone": timezone}

  root_agent = Agent(
      name="root_agent",
      model=mock_model,
      tools=[get_weather, get_time],
  )

  # Use the custom runner
  runner = StreamingTestRunner(root_agent=root_agent, max_responses=3)
  live_request_queue = LiveRequestQueue()
  live_request_queue.send_realtime(
      blob=types.Blob(
          data=b"What is the weather and current time?", mime_type="audio/pcm"
      )
  )

  res_events = runner.run_live(live_request_queue)

  assert res_events is not None, "Expected a list of events, got None."
  assert len(res_events) >= 1, "Expected at least one event."

  # Check function calls
  weather_call_found = False
  time_call_found = False

  for event in res_events:
    if event.content and event.content.parts:
      for part in event.content.parts:
        if part.function_call:
          if part.function_call.name == "get_weather":
            weather_call_found = True
            assert part.function_call.args["location"] == "San Francisco"
          elif part.function_call.name == "get_time":
            time_call_found = True
            assert part.function_call.args["timezone"] == "PST"

  # In live streaming, we primarily check that function calls are generated correctly
  assert (
      weather_call_found or time_call_found
  ), "Expected at least one function call."


def test_live_streaming_function_call_parallel():
  """Test live streaming with parallel function calls."""
  # Create parallel function calls in the same response
  function_call1 = types.Part.from_function_call(
      name="get_weather", args={"location": "San Francisco"}
  )
  function_call2 = types.Part.from_function_call(
      name="get_weather", args={"location": "New York"}
  )

  # Create LLM response with parallel function calls
  response1 = LlmResponse(
      content=types.Content(
          role="model", parts=[function_call1, function_call2]
      ),
      turn_complete=False,
  )
  response2 = LlmResponse(
      turn_complete=True,
  )

  mock_model = testing_utils.MockModel.create([response1, response2])

  # Mock function
  def get_weather(location: str) -> dict:
    temperatures = {"San Francisco": 22, "New York": 15}
    return {"temperature": temperatures.get(location, 20), "location": location}

  root_agent = Agent(
      name="root_agent",
      model=mock_model,
      tools=[get_weather],
  )

  # Use the custom runner
  runner = StreamingTestRunner(root_agent=root_agent, max_responses=3)
  live_request_queue = LiveRequestQueue()
  live_request_queue.send_realtime(
      blob=types.Blob(
          data=b"Compare weather in SF and NYC", mime_type="audio/pcm"
      )
  )

  res_events = runner.run_live(live_request_queue)

  assert res_events is not None, "Expected a list of events, got None."
  assert len(res_events) >= 1, "Expected at least one event."

  # Check parallel function calls
  sf_call_found = False
  nyc_call_found = False

  for event in res_events:
    if event.content and event.content.parts:
      for part in event.content.parts:
        if part.function_call and part.function_call.name == "get_weather":
          location = part.function_call.args["location"]
          if location == "San Francisco":
            sf_call_found = True
          elif location == "New York":
            nyc_call_found = True

  assert (
      sf_call_found and nyc_call_found
  ), "Expected both location function calls."


def test_live_streaming_function_call_with_error():
  """Test live streaming with function call that returns an error."""
  # Create a function call response
  function_call = types.Part.from_function_call(
      name="get_weather", args={"location": "Invalid Location"}
  )

  # Create LLM responses
  response1 = LlmResponse(
      content=types.Content(role="model", parts=[function_call]),
      turn_complete=False,
  )
  response2 = LlmResponse(
      turn_complete=True,
  )

  mock_model = testing_utils.MockModel.create([response1, response2])

  # Mock function that returns an error for invalid locations
  def get_weather(location: str) -> dict:
    if location == "Invalid Location":
      return {"error": "Location not found"}
    return {"temperature": 22, "condition": "sunny", "location": location}

  root_agent = Agent(
      name="root_agent",
      model=mock_model,
      tools=[get_weather],
  )

  # Use the custom runner
  runner = StreamingTestRunner(root_agent=root_agent, max_responses=3)
  live_request_queue = LiveRequestQueue()
  live_request_queue.send_realtime(
      blob=types.Blob(
          data=b"What is weather in Invalid Location?", mime_type="audio/pcm"
      )
  )

  res_events = runner.run_live(live_request_queue)

  assert res_events is not None, "Expected a list of events, got None."
  assert len(res_events) >= 1, "Expected at least one event."

  # Check that we got the function call (error handling happens at execution time)
  function_call_found = False
  for event in res_events:
    if event.content and event.content.parts:
      for part in event.content.parts:
        if part.function_call and part.function_call.name == "get_weather":
          function_call_found = True
          assert part.function_call.args["location"] == "Invalid Location"

  assert function_call_found, "Expected function call event with error case."


def test_live_streaming_function_call_sync_tool():
  """Test live streaming with synchronous function call."""
  # Create a function call response
  function_call = types.Part.from_function_call(
      name="calculate", args={"x": 5, "y": 3}
  )

  # Create LLM responses
  response1 = LlmResponse(
      content=types.Content(role="model", parts=[function_call]),
      turn_complete=False,
  )
  response2 = LlmResponse(
      turn_complete=True,
  )

  mock_model = testing_utils.MockModel.create([response1, response2])

  # Mock sync function
  def calculate(x: int, y: int) -> dict:
    return {"result": x + y, "operation": "addition"}

  root_agent = Agent(
      name="root_agent",
      model=mock_model,
      tools=[calculate],
  )

  # Use the custom runner
  runner = StreamingTestRunner(root_agent=root_agent, max_responses=3)
  live_request_queue = LiveRequestQueue()
  live_request_queue.send_realtime(
      blob=types.Blob(data=b"Calculate 5 plus 3", mime_type="audio/pcm")
  )

  res_events = runner.run_live(live_request_queue)

  assert res_events is not None, "Expected a list of events, got None."
  assert len(res_events) >= 1, "Expected at least one event."

  # Check function call
  function_call_found = False
  for event in res_events:
    if event.content and event.content.parts:
      for part in event.content.parts:
        if part.function_call and part.function_call.name == "calculate":
          function_call_found = True
          assert part.function_call.args["x"] == 5
          assert part.function_call.args["y"] == 3

  assert function_call_found, "Expected calculate function call event."


def test_live_streaming_simple_streaming_tool():
  """Test live streaming with a simple streaming tool (non-video)."""
  # Create a function call response for the streaming tool
  function_call = types.Part.from_function_call(
      name="monitor_stock_price", args={"stock_symbol": "AAPL"}
  )

  # Create LLM responses
  response1 = LlmResponse(
      content=types.Content(role="model", parts=[function_call]),
      turn_complete=False,
  )
  response2 = LlmResponse(
      turn_complete=True,
  )

  mock_model = testing_utils.MockModel.create([response1, response2])

  # Mock simple streaming tool (without return type annotation to avoid parsing issues)
  async def monitor_stock_price(stock_symbol: str):
    """Mock streaming tool that monitors stock prices."""
    # Simulate some streaming updates
    yield f"Stock {stock_symbol} price: $150"
    await asyncio.sleep(0.1)
    yield f"Stock {stock_symbol} price: $155"
    await asyncio.sleep(0.1)
    yield f"Stock {stock_symbol} price: $160"

  def stop_streaming(function_name: str):
    """Stop the streaming tool."""
    pass

  root_agent = Agent(
      name="root_agent",
      model=mock_model,
      tools=[monitor_stock_price, stop_streaming],
  )

  # Use the custom runner
  runner = StreamingTestRunner(root_agent=root_agent, max_responses=3)
  live_request_queue = LiveRequestQueue()
  live_request_queue.send_realtime(
      blob=types.Blob(data=b"Monitor AAPL stock price", mime_type="audio/pcm")
  )

  res_events = runner.run_live(live_request_queue)

  assert res_events is not None, "Expected a list of events, got None."
  assert len(res_events) >= 1, "Expected at least one event."

  # Check that we got the streaming tool function call
  function_call_found = False
  for event in res_events:
    if event.content and event.content.parts:
      for part in event.content.parts:
        if (
            part.function_call
            and part.function_call.name == "monitor_stock_price"
        ):
          function_call_found = True
          assert part.function_call.args["stock_symbol"] == "AAPL"

  assert (
      function_call_found
  ), "Expected monitor_stock_price function call event."


def test_live_streaming_video_streaming_tool():
  """Test live streaming with a video streaming tool."""
  # Create a function call response for the video streaming tool
  function_call = types.Part.from_function_call(
      name="monitor_video_stream", args={}
  )

  # Create LLM responses
  response1 = LlmResponse(
      content=types.Content(role="model", parts=[function_call]),
      turn_complete=False,
  )
  response2 = LlmResponse(
      turn_complete=True,
  )

  mock_model = testing_utils.MockModel.create([response1, response2])

  # Mock video streaming tool (without return type annotation to avoid parsing issues)
  async def monitor_video_stream(input_stream: LiveRequestQueue):
    """Mock video streaming tool that processes video frames."""
    # Simulate processing a few frames from the input stream
    frame_count = 0
    while frame_count < 3:  # Process a few frames
      try:
        # Try to get a frame from the queue with timeout
        live_req = await asyncio.wait_for(input_stream.get(), timeout=0.1)
        if live_req.blob and live_req.blob.mime_type == "image/jpeg":
          frame_count += 1
          yield f"Processed frame {frame_count}: detected 2 people"
      except asyncio.TimeoutError:
        # No more frames, simulate detection anyway for testing
        frame_count += 1
        yield f"Simulated frame {frame_count}: detected 1 person"
      await asyncio.sleep(0.1)

  def stop_streaming(function_name: str):
    """Stop the streaming tool."""
    pass

  root_agent = Agent(
      name="root_agent",
      model=mock_model,
      tools=[monitor_video_stream, stop_streaming],
  )

  # Use the custom runner
  runner = StreamingTestRunner(root_agent=root_agent, max_responses=3)
  live_request_queue = LiveRequestQueue()

  # Send some mock video frames
  live_request_queue.send_realtime(
      blob=types.Blob(data=b"fake_jpeg_data_1", mime_type="image/jpeg")
  )
  live_request_queue.send_realtime(
      blob=types.Blob(data=b"fake_jpeg_data_2", mime_type="image/jpeg")
  )
  live_request_queue.send_realtime(
      blob=types.Blob(data=b"Monitor video stream", mime_type="audio/pcm")
  )

  res_events = runner.run_live(live_request_queue)

  assert res_events is not None, "Expected a list of events, got None."
  assert len(res_events) >= 1, "Expected at least one event."

  # Check that we got the video streaming tool function call
  function_call_found = False
  for event in res_events:
    if event.content and event.content.parts:
      for part in event.content.parts:
        if (
            part.function_call
            and part.function_call.name == "monitor_video_stream"
        ):
          function_call_found = True

  assert (
      function_call_found
  ), "Expected monitor_video_stream function call event."


def test_live_streaming_stop_streaming_tool():
  """Test live streaming with stop_streaming functionality."""
  # Create function calls for starting and stopping a streaming tool
  start_function_call = types.Part.from_function_call(
      name="monitor_stock_price", args={"stock_symbol": "TSLA"}
  )
  stop_function_call = types.Part.from_function_call(
      name="stop_streaming", args={"function_name": "monitor_stock_price"}
  )

  # Create LLM responses: start streaming, then stop streaming
  response1 = LlmResponse(
      content=types.Content(role="model", parts=[start_function_call]),
      turn_complete=False,
  )
  response2 = LlmResponse(
      content=types.Content(role="model", parts=[stop_function_call]),
      turn_complete=False,
  )
  response3 = LlmResponse(
      turn_complete=True,
  )

  mock_model = testing_utils.MockModel.create([response1, response2, response3])

  # Mock streaming tool and stop function
  async def monitor_stock_price(stock_symbol: str):
    """Mock streaming tool that monitors stock prices."""
    yield f"Started monitoring {stock_symbol}"
    while True:  # Infinite stream (would be stopped by stop_streaming)
      yield f"Stock {stock_symbol} price update"
      await asyncio.sleep(0.1)

  def stop_streaming(function_name: str):
    """Stop the streaming tool."""
    return f"Stopped streaming for {function_name}"

  root_agent = Agent(
      name="root_agent",
      model=mock_model,
      tools=[monitor_stock_price, stop_streaming],
  )

  # Use the custom runner
  runner = StreamingTestRunner(root_agent=root_agent, max_responses=3)
  live_request_queue = LiveRequestQueue()
  live_request_queue.send_realtime(
      blob=types.Blob(data=b"Monitor TSLA and then stop", mime_type="audio/pcm")
  )

  res_events = runner.run_live(live_request_queue)

  assert res_events is not None, "Expected a list of events, got None."
  assert len(res_events) >= 1, "Expected at least one event."

  # Check that we got both function calls
  monitor_call_found = False
  stop_call_found = False

  for event in res_events:
    if event.content and event.content.parts:
      for part in event.content.parts:
        if part.function_call:
          if part.function_call.name == "monitor_stock_price":
            monitor_call_found = True
            assert part.function_call.args["stock_symbol"] == "TSLA"
          elif part.function_call.name == "stop_streaming":
            stop_call_found = True
            assert (
                part.function_call.args["function_name"]
                == "monitor_stock_price"
            )

  assert monitor_call_found, "Expected monitor_stock_price function call event."
  assert stop_call_found, "Expected stop_streaming function call event."


def test_live_streaming_multiple_streaming_tools():
  """Test live streaming with multiple streaming tools running simultaneously."""
  # Create function calls for multiple streaming tools
  stock_function_call = types.Part.from_function_call(
      name="monitor_stock_price", args={"stock_symbol": "NVDA"}
  )
  video_function_call = types.Part.from_function_call(
      name="monitor_video_stream", args={}
  )

  # Create LLM responses: start both streaming tools
  response1 = LlmResponse(
      content=types.Content(
          role="model", parts=[stock_function_call, video_function_call]
      ),
      turn_complete=False,
  )
  response2 = LlmResponse(
      turn_complete=True,
  )

  mock_model = testing_utils.MockModel.create([response1, response2])

  # Mock streaming tools
  async def monitor_stock_price(stock_symbol: str):
    """Mock streaming tool that monitors stock prices."""
    yield f"Stock {stock_symbol} price: $800"
    await asyncio.sleep(0.1)
    yield f"Stock {stock_symbol} price: $805"

  async def monitor_video_stream(input_stream: LiveRequestQueue):
    """Mock video streaming tool."""
    yield "Video monitoring started"
    await asyncio.sleep(0.1)
    yield "Detected motion in video stream"

  def stop_streaming(function_name: str):
    """Stop the streaming tool."""
    pass

  root_agent = Agent(
      name="root_agent",
      model=mock_model,
      tools=[monitor_stock_price, monitor_video_stream, stop_streaming],
  )

  # Use the custom runner
  runner = StreamingTestRunner(root_agent=root_agent, max_responses=3)
  live_request_queue = LiveRequestQueue()
  live_request_queue.send_realtime(
      blob=types.Blob(
          data=b"Monitor both stock and video", mime_type="audio/pcm"
      )
  )

  res_events = runner.run_live(live_request_queue)

  assert res_events is not None, "Expected a list of events, got None."
  assert len(res_events) >= 1, "Expected at least one event."

  # Check that we got both streaming tool function calls
  stock_call_found = False
  video_call_found = False

  for event in res_events:
    if event.content and event.content.parts:
      for part in event.content.parts:
        if part.function_call:
          if part.function_call.name == "monitor_stock_price":
            stock_call_found = True
            assert part.function_call.args["stock_symbol"] == "NVDA"
          elif part.function_call.name == "monitor_video_stream":
            video_call_found = True

  assert stock_call_found, "Expected monitor_stock_price function call event."
  assert video_call_found, "Expected monitor_video_stream function call event."


def test_live_streaming_function_call_yielded_before_finished_transcription():
  """Test that function calls arriving during live transcription are yielded immediately.

  This verifies that tool call events are not buffered and are permitted to
  arrive in the stream before the final completed transcription event.
  """
  function_call = types.Part.from_function_call(
      name="get_weather", args={"location": "San Francisco"}
  )

  response1 = LlmResponse(
      input_transcription=types.Transcription(text="Show"),
      partial=True,  # ← Triggers is_transcribing = True
  )
  response2 = LlmResponse(
      content=types.Content(
          role="model", parts=[function_call]
      ),  # ← Gets buffered
      turn_complete=False,
  )
  response3 = LlmResponse(
      input_transcription=types.Transcription(text="Show me the weather"),
      partial=False,  # ← Transcription ends, buffered events yielded
  )
  response4 = LlmResponse(
      turn_complete=True,
  )

  mock_model = testing_utils.MockModel.create(
      [response1, response2, response3, response4]
  )

  def get_weather(location: str) -> dict:
    return {"temperature": 22, "location": location}

  root_agent = Agent(
      name="root_agent",
      model=mock_model,
      tools=[get_weather],
  )

  runner = StreamingTestRunner(root_agent=root_agent, max_responses=5)
  live_request_queue = LiveRequestQueue()
  live_request_queue.send_realtime(
      blob=types.Blob(data=b"Show me the weather", mime_type="audio/pcm")
  )

  res_events = runner.run_live(live_request_queue)

  assert res_events is not None, "Expected a list of events, got None."
  assert len(res_events) >= 1, "Expected at least one event."

  function_call_index = -1
  finished_transcription_index = -1

  for idx, event in enumerate(res_events):
    if event.content and event.content.parts:
      for part in event.content.parts:
        if part.function_call and part.function_call.name == "get_weather":
          function_call_index = idx
          assert part.function_call.args["location"] == "San Francisco"
        if (
            part.function_response
            and part.function_response.name == "get_weather"
        ):
          assert part.function_response.response["temperature"] == 22
    if (
        event.input_transcription
        and event.input_transcription.text == "Show me the weather"
    ):
      finished_transcription_index = idx

  assert function_call_index != -1, "Function call event was not yielded."
  assert (
      finished_transcription_index != -1
  ), "Finished transcription event was not yielded."
  assert function_call_index < finished_transcription_index, (
      f"Expected function call (at index {function_call_index}) to arrive"
      " before finished transcription (at index"
      f" {finished_transcription_index})."
  )


def test_live_streaming_text_content_persisted_in_session():
  """Test that user text content sent via send_content is persisted in session."""
  response1 = LlmResponse(
      content=types.Content(
          role="model", parts=[types.Part(text="Hello! How can I help you?")]
      ),
      turn_complete=True,
  )

  mock_model = testing_utils.MockModel.create([response1])

  root_agent = Agent(
      name="root_agent",
      model=mock_model,
      tools=[],
  )

  runner = StreamingTestRunner(root_agent=root_agent, max_responses=1)
  live_request_queue = LiveRequestQueue()

  # Send text content (not audio blob)
  user_text = "Hello, this is a test message"
  live_request_queue.send_content(
      types.Content(role="user", parts=[types.Part(text=user_text)])
  )

  res_events, session = runner.run_live_and_get_session(live_request_queue)

  assert res_events is not None, "Expected a list of events, got None."

  # Check that user text content was persisted in the session
  user_content_found = False
  for event in session.events:
    if event.author == "user" and event.content:
      for part in event.content.parts:
        if part.text and user_text in part.text:
          user_content_found = True
          break

  assert user_content_found, (
      f'Expected user text content "{user_text}" to be persisted in session. '
      f"Session events: {[e.content for e in session.events]}"
  )


def _collect_function_call_names(events):
  """Extract the set of function call names from a list of events."""
  return {fc.name for event in events for fc in event.get_function_calls()}


class _LiveTestRunner(testing_utils.InMemoryRunner):
  """Test runner with custom event loop management for live streaming tests."""

  def _run_with_loop(self, coro: Awaitable[Any]) -> None:
    """Run a coroutine in a new event loop, suppressing timeouts."""
    try:
      old_loop = asyncio.get_event_loop()
    except RuntimeError:
      old_loop = None
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
      loop.run_until_complete(coro)
    except (asyncio.TimeoutError, asyncio.CancelledError):
      pass
    finally:
      # Cancel all pending tasks to prevent leaks and warnings
      pending = asyncio.all_tasks(loop)
      for task in pending:
        task.cancel()
      if pending:
        try:
          loop.run_until_complete(
              asyncio.gather(*pending, return_exceptions=True)
          )
        except Exception:  # pylint: disable=broad-except
          pass
      loop.close()
      asyncio.set_event_loop(old_loop)

  def run_live(
      self,
      live_request_queue: LiveRequestQueue,
      max_responses: int = 3,
  ) -> list[testing_utils.Event]:
    """Run live and collect up to max_responses events."""
    collected = []

    async def consume(session: testing_utils.Session):
      run_res = self.runner.run_live(
          session=session,
          live_request_queue=live_request_queue,
      )
      async with aclosing(run_res) as agen:
        async for response in agen:
          collected.append(response)
          if len(collected) >= max_responses:
            await _wait_for_queue_empty(live_request_queue)
            return

    self._run_with_loop(asyncio.wait_for(consume(self.session), timeout=5.0))
    return collected


def test_input_streaming_tool_registered_lazily_with_stream():
  """Test that input-streaming tools are registered lazily when called and receive a stream."""
  # A text response before the function call lets us observe that the
  # tool is NOT registered before the model calls it.
  text_response = LlmResponse(
      content=types.Content(
          role="model",
          parts=[types.Part(text="Processing...")],
      ),
      turn_complete=False,
  )
  function_call = types.Part.from_function_call(
      name="monitor_video_stream", args={}
  )
  call_response = LlmResponse(
      content=types.Content(role="model", parts=[function_call]),
      turn_complete=False,
  )
  done_response = LlmResponse(turn_complete=True)

  mock_model = testing_utils.MockModel.create(
      [text_response, call_response, done_response]
  )

  stream_state_during_call = None

  async def monitor_video_stream(
      input_stream: LiveRequestQueue,
  ) -> AsyncGenerator[str, None]:
    """Record whether input_stream was provided."""
    nonlocal stream_state_during_call
    stream_state_during_call = input_stream is not None
    yield "monitoring started"

  root_agent = Agent(
      name="root_agent",
      model=mock_model,
      tools=[monitor_video_stream],
  )

  runner = _LiveTestRunner(root_agent=root_agent)

  # Capture the invocation context to inspect registration state.
  captured_context = None
  original_method = runner.runner._new_invocation_context_for_live

  def capturing_method(*args, **kwargs) -> Any:
    nonlocal captured_context
    ctx = original_method(*args, **kwargs)
    captured_context = ctx
    return ctx

  runner.runner._new_invocation_context_for_live = capturing_method

  live_request_queue = LiveRequestQueue()
  live_request_queue.send_realtime(
      blob=types.Blob(data=b"test_data", mime_type="audio/pcm")
  )

  # Collect events and check that the tool is NOT registered before
  # the model calls it.
  collected = []
  not_registered_before_call = None

  async def consume(session: testing_utils.Session):
    nonlocal not_registered_before_call
    run_res = runner.runner.run_live(
        session=session,
        live_request_queue=live_request_queue,
    )
    async with aclosing(run_res) as agen:
      async for response in agen:
        collected.append(response)
        # On the first non-function-call event, verify the tool is not
        # yet registered (lazy registration).
        active = (
            captured_context.active_streaming_tools
            if captured_context
            else None
        )
        if (
            not_registered_before_call is None
            and not response.get_function_calls()
        ):
          not_registered_before_call = (
              active is None or "monitor_video_stream" not in active
          )
        if len(collected) >= 4:
          await _wait_for_queue_empty(live_request_queue)
          return

  runner._run_with_loop(asyncio.wait_for(consume(runner.session), timeout=5.0))

  # Tool should not be registered before the model calls it.
  assert (
      not_registered_before_call is True
  ), "Expected tool to NOT be registered before the model calls it"
  # When the model calls the tool, input_stream should be provided.
  assert (
      stream_state_during_call is True
  ), "Expected input_stream to be provided to the streaming tool when called"


def test_stop_streaming_resets_stream_to_none():
  """Test that stop_streaming sets stream back to None."""
  start_call = types.Part.from_function_call(
      name="monitor_stock_price", args={"stock_symbol": "GOOG"}
  )
  stop_call = types.Part.from_function_call(
      name="stop_streaming", args={"function_name": "monitor_stock_price"}
  )

  response1 = LlmResponse(
      content=types.Content(role="model", parts=[start_call]),
      turn_complete=False,
  )
  response2 = LlmResponse(
      content=types.Content(role="model", parts=[stop_call]),
      turn_complete=False,
  )
  response3 = LlmResponse(turn_complete=True)

  mock_model = testing_utils.MockModel.create([response1, response2, response3])

  async def monitor_stock_price(
      stock_symbol: str,
  ) -> AsyncGenerator[str, None]:
    """Yield periodic price updates for the given stock symbol."""
    yield f"Monitoring {stock_symbol}"
    while True:
      await asyncio.sleep(0.1)
      yield f"{stock_symbol} price update"

  def stop_streaming(function_name: str) -> None:
    """Stop a running streaming tool by name."""
    pass

  root_agent = Agent(
      name="root_agent",
      model=mock_model,
      tools=[monitor_stock_price, stop_streaming],
  )

  runner = _LiveTestRunner(root_agent=root_agent)

  # Capture the child invocation context (created by _create_invocation_context
  # inside base_agent.run_live) to inspect active_streaming_tools.
  # We cannot use the parent context from _new_invocation_context_for_live
  # because model_copy creates a separate child object.
  captured_child_context = None
  original_create = root_agent._create_invocation_context

  def capturing_create(*args, **kwargs) -> Any:
    nonlocal captured_child_context
    ctx = original_create(*args, **kwargs)
    captured_child_context = ctx
    return ctx

  root_agent._create_invocation_context = capturing_create

  live_request_queue = LiveRequestQueue()
  live_request_queue.send_realtime(
      blob=types.Blob(data=b"Monitor GOOG then stop", mime_type="audio/pcm")
  )

  res_events = runner.run_live(live_request_queue, max_responses=4)

  # Verify both function calls were processed.
  call_names = _collect_function_call_names(res_events)
  assert (
      "monitor_stock_price" in call_names
  ), "Expected monitor_stock_price function call."
  assert (
      "stop_streaming" in call_names
  ), "Expected stop_streaming function call."

  # Verify that stop_streaming reset the stream to None.
  assert (
      captured_child_context is not None
  ), "Expected child invocation context to be captured"
  active_tools = captured_child_context.active_streaming_tools or {}
  assert (
      "monitor_stock_price" in active_tools
  ), "Expected monitor_stock_price in active_streaming_tools"
  assert (
      active_tools["monitor_stock_price"].stream is None
  ), "Expected stream to be reset to None after stop_streaming"


def test_output_streaming_tool_registered_lazily_without_stream():
  """Test that output-streaming tools are registered lazily when called, with stream=None."""
  function_call = types.Part.from_function_call(
      name="monitor_stock_price", args={"stock_symbol": "GOOG"}
  )
  response1 = LlmResponse(
      content=types.Content(role="model", parts=[function_call]),
      turn_complete=False,
  )
  response2 = LlmResponse(turn_complete=True)

  mock_model = testing_utils.MockModel.create([response1, response2])

  async def monitor_stock_price(
      stock_symbol: str,
  ) -> AsyncGenerator[str, None]:
    """Yield periodic price updates."""
    yield f"price for {stock_symbol}"

  root_agent = Agent(
      name="root_agent",
      model=mock_model,
      tools=[monitor_stock_price],
  )

  runner = _LiveTestRunner(root_agent=root_agent)

  # Capture the child invocation context (created by _create_invocation_context
  # inside base_agent.run_live) to inspect active_streaming_tools.
  captured_child_context = None
  original_create = root_agent._create_invocation_context

  def capturing_create(*args, **kwargs) -> Any:
    nonlocal captured_child_context
    ctx = original_create(*args, **kwargs)
    captured_child_context = ctx
    return ctx

  root_agent._create_invocation_context = capturing_create

  live_request_queue = LiveRequestQueue()
  live_request_queue.send_realtime(
      blob=types.Blob(data=b"test", mime_type="audio/pcm")
  )

  runner.run_live(live_request_queue, max_responses=3)

  # After the model calls the tool, it should be registered with
  # stream=None (output-streaming tools don't consume the live stream).
  assert captured_child_context is not None
  active_tools = captured_child_context.active_streaming_tools or {}
  assert (
      "monitor_stock_price" in active_tools
  ), "Expected output-streaming tool to be registered when called"
  assert (
      active_tools["monitor_stock_price"].stream is None
  ), "Expected stream to be None for output-streaming tool"


def _run_single_tool_live(
    tool_func,
    func_name: str,
    func_args: dict[str, Any] | None = None,
    max_responses: int = 3,
) -> dict[str, Any]:
  """Run a live session that invokes a single tool and return active_streaming_tools.

  Sets up a mock model that issues one function call then completes,
  creates an agent with the given tool, captures the invocation context,
  and returns the ``active_streaming_tools`` dict after execution.
  """
  function_call = types.Part.from_function_call(
      name=func_name, args=func_args or {}
  )
  response1 = LlmResponse(
      content=types.Content(role="model", parts=[function_call]),
      turn_complete=False,
  )
  response2 = LlmResponse(turn_complete=True)

  mock_model = testing_utils.MockModel.create([response1, response2])

  root_agent = Agent(
      name="root_agent",
      model=mock_model,
      tools=[tool_func],
  )

  runner = _LiveTestRunner(root_agent=root_agent)

  captured_child_context = None
  original_create = root_agent._create_invocation_context

  def capturing_create(*args, **kwargs) -> Any:
    nonlocal captured_child_context
    ctx = original_create(*args, **kwargs)
    captured_child_context = ctx
    return ctx

  root_agent._create_invocation_context = capturing_create

  live_request_queue = LiveRequestQueue()
  live_request_queue.send_realtime(
      blob=types.Blob(data=b"test", mime_type="audio/pcm")
  )

  runner.run_live(live_request_queue, max_responses=max_responses)

  assert captured_child_context is not None
  return captured_child_context.active_streaming_tools or {}


def test_input_streaming_tool_has_stream_set_at_registration():
  """Test that input-streaming tools get .stream set to a LiveRequestQueue during registration."""

  async def monitor_video_stream(
      input_stream: LiveRequestQueue,
  ) -> AsyncGenerator[str, None]:
    """Simulate an input-streaming tool."""
    yield "started"

  active_tools = _run_single_tool_live(
      monitor_video_stream, "monitor_video_stream"
  )

  assert (
      "monitor_video_stream" in active_tools
  ), "Expected input-streaming tool to be registered when called"
  # Stream should be a LiveRequestQueue, not None.
  assert (
      active_tools["monitor_video_stream"].stream is not None
  ), "Expected .stream to be set for input-streaming tool"
  assert isinstance(
      active_tools["monitor_video_stream"].stream, LiveRequestQueue
  ), "Expected .stream to be a LiveRequestQueue instance"


def test_input_streaming_tool_stream_recreated_after_stop():
  """Test that re-invoking an input-streaming tool after stop creates a new stream."""
  start_call = types.Part.from_function_call(name="monitor_video", args={})
  stop_call = types.Part.from_function_call(
      name="stop_streaming", args={"function_name": "monitor_video"}
  )
  restart_call = types.Part.from_function_call(name="monitor_video", args={})

  response1 = LlmResponse(
      content=types.Content(role="model", parts=[start_call]),
      turn_complete=False,
  )
  response2 = LlmResponse(
      content=types.Content(role="model", parts=[stop_call]),
      turn_complete=False,
  )
  response3 = LlmResponse(
      content=types.Content(role="model", parts=[restart_call]),
      turn_complete=False,
  )
  response4 = LlmResponse(turn_complete=True)

  mock_model = testing_utils.MockModel.create(
      [response1, response2, response3, response4]
  )

  call_count = 0

  async def monitor_video(
      input_stream: LiveRequestQueue,
  ) -> AsyncGenerator[str, None]:
    """Simulate an input-streaming tool that tracks invocation count."""
    nonlocal call_count
    call_count += 1
    yield f"started (call {call_count})"
    while True:
      await asyncio.sleep(0.1)
      yield "frame"

  def stop_streaming(function_name: str) -> None:
    """Stop a running streaming tool by name."""
    pass

  root_agent = Agent(
      name="root_agent",
      model=mock_model,
      tools=[monitor_video, stop_streaming],
  )

  runner = _LiveTestRunner(root_agent=root_agent)

  captured_child_context = None
  original_create = root_agent._create_invocation_context

  def capturing_create(*args, **kwargs) -> Any:
    nonlocal captured_child_context
    ctx = original_create(*args, **kwargs)
    captured_child_context = ctx
    return ctx

  root_agent._create_invocation_context = capturing_create

  live_request_queue = LiveRequestQueue()
  live_request_queue.send_realtime(
      blob=types.Blob(data=b"test", mime_type="audio/pcm")
  )

  res_events = runner.run_live(live_request_queue, max_responses=8)

  # monitor_video should appear at least twice in function calls
  # (start + restart). Function response events may add extra
  # occurrences.
  call_names = [
      fc.name for event in res_events for fc in event.get_function_calls()
  ]
  assert (
      call_names.count("monitor_video") >= 2
  ), f"Expected monitor_video called at least twice, got: {call_names}"

  # After re-invocation, stream should be set again (not None).
  assert captured_child_context is not None
  active_tools = captured_child_context.active_streaming_tools or {}
  assert "monitor_video" in active_tools
  assert (
      active_tools["monitor_video"].stream is not None
  ), "Expected .stream to be recreated after stop + re-invocation"


def test_async_gen_with_input_stream_wrong_annotation_gets_no_stream():
  """Test that an async generator with input_stream param but wrong annotation gets no stream."""
  received_input_stream = None

  async def my_tool(input_stream: str) -> AsyncGenerator[str, None]:
    """Simulate an async generator whose input_stream is typed as str."""
    nonlocal received_input_stream
    received_input_stream = input_stream
    yield f"got: {input_stream}"

  active_tools = _run_single_tool_live(
      my_tool, "my_tool", func_args={"input_stream": "some_value"}
  )

  assert (
      "my_tool" in active_tools
  ), "Expected async generator tool to be registered"
  # Stream should be None because annotation is str, not LiveRequestQueue.
  assert active_tools["my_tool"].stream is None, (
      "Expected .stream to be None when input_stream annotation is not"
      " LiveRequestQueue"
  )
  # The tool should have received the model-provided arg value, not a
  # LiveRequestQueue.
  assert (
      received_input_stream == "some_value"
  ), "Expected input_stream to be the model-provided string value"
