# Copyright 2025 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 AsyncMock
from unittest.mock import Mock

from google.adk.flows.llm_flows.transcription_manager import TranscriptionManager
from google.genai import types
import pytest

from ... import testing_utils


class TestTranscriptionManager:
  """Test the TranscriptionManager class."""

  def setup_method(self):
    """Set up test fixtures."""
    self.manager = TranscriptionManager()

  @pytest.mark.asyncio
  async def test_handle_input_transcription(self):
    """Test handling user input transcription events."""
    invocation_context = await testing_utils.create_invocation_context(
        testing_utils.create_test_agent()
    )

    # Set up mock session service
    mock_session_service = AsyncMock()
    invocation_context.session_service = mock_session_service

    # Create test transcription
    transcription = types.Transcription(text='Hello from user')

    # Handle transcription
    await self.manager.handle_input_transcription(
        invocation_context, transcription
    )

    # Verify session service was called
    mock_session_service.append_event.assert_called_once()

    # Check the event that was created
    call_args = mock_session_service.append_event.call_args
    event = call_args[0][1]  # Second argument is the event

    assert event.author == 'user'
    assert event.input_transcription == transcription
    assert event.output_transcription is None
    assert event.invocation_id == invocation_context.invocation_id
    assert isinstance(event.timestamp, float)

  @pytest.mark.asyncio
  async def test_handle_output_transcription(self):
    """Test handling model output transcription events."""
    agent = testing_utils.create_test_agent()
    invocation_context = await testing_utils.create_invocation_context(agent)

    # Set up mock session service
    mock_session_service = AsyncMock()
    invocation_context.session_service = mock_session_service

    # Create test transcription
    transcription = types.Transcription(text='Hello from model')

    # Handle transcription
    await self.manager.handle_output_transcription(
        invocation_context, transcription
    )

    # Verify session service was called
    mock_session_service.append_event.assert_called_once()

    # Check the event that was created
    call_args = mock_session_service.append_event.call_args
    event = call_args[0][1]  # Second argument is the event

    assert event.author == agent.name
    assert event.input_transcription is None
    assert event.output_transcription == transcription
    assert event.invocation_id == invocation_context.invocation_id
    assert isinstance(event.timestamp, float)

  @pytest.mark.asyncio
  async def test_handle_multiple_transcriptions(self):
    """Test handling multiple transcription events."""
    invocation_context = await testing_utils.create_invocation_context(
        testing_utils.create_test_agent()
    )

    # Set up mock session service
    mock_session_service = AsyncMock()
    invocation_context.session_service = mock_session_service

    # Handle multiple input transcriptions
    for i in range(3):
      transcription = types.Transcription(text=f'User message {i}')
      await self.manager.handle_input_transcription(
          invocation_context, transcription
      )

    # Handle multiple output transcriptions
    for i in range(2):
      transcription = types.Transcription(text=f'Model response {i}')
      await self.manager.handle_output_transcription(
          invocation_context, transcription
      )

    # Verify session service was called for each transcription
    assert mock_session_service.append_event.call_count == 5

  @pytest.mark.asyncio
  async def test_error_handling_input_transcription(self):
    """Test error handling during input transcription processing."""
    invocation_context = await testing_utils.create_invocation_context(
        testing_utils.create_test_agent()
    )

    # Set up mock session service that raises an error
    mock_session_service = AsyncMock()
    mock_session_service.append_event.side_effect = Exception(
        'Session service error'
    )
    invocation_context.session_service = mock_session_service

    # Create test transcription
    transcription = types.Transcription(text='Test transcription')

    # Handle transcription should raise the exception
    with pytest.raises(Exception, match='Session service error'):
      await self.manager.handle_input_transcription(
          invocation_context, transcription
      )

  @pytest.mark.asyncio
  async def test_error_handling_output_transcription(self):
    """Test error handling during output transcription processing."""
    invocation_context = await testing_utils.create_invocation_context(
        testing_utils.create_test_agent()
    )

    # Set up mock session service that raises an error
    mock_session_service = AsyncMock()
    mock_session_service.append_event.side_effect = Exception(
        'Session service error'
    )
    invocation_context.session_service = mock_session_service

    # Create test transcription
    transcription = types.Transcription(text='Test transcription')

    # Handle transcription should raise the exception
    with pytest.raises(Exception, match='Session service error'):
      await self.manager.handle_output_transcription(
          invocation_context, transcription
      )

  def test_get_transcription_stats_empty_session(self):
    """Test getting transcription statistics for empty session."""
    invocation_context = Mock()
    invocation_context.session.events = []

    stats = self.manager.get_transcription_stats(invocation_context)

    expected = {
        'input_transcriptions': 0,
        'output_transcriptions': 0,
        'total_transcriptions': 0,
    }
    assert stats == expected

  def test_get_transcription_stats_with_events(self):
    """Test getting transcription statistics for session with events."""
    invocation_context = Mock()

    # Create mock events
    input_event1 = Mock()
    input_event1.input_transcription = types.Transcription(text='User 1')
    input_event1.output_transcription = None

    input_event2 = Mock()
    input_event2.input_transcription = types.Transcription(text='User 2')
    input_event2.output_transcription = None

    output_event = Mock()
    output_event.input_transcription = None
    output_event.output_transcription = types.Transcription(
        text='Model response'
    )

    regular_event = Mock()
    regular_event.input_transcription = None
    regular_event.output_transcription = None

    invocation_context.session.events = [
        input_event1,
        output_event,
        input_event2,
        regular_event,
    ]

    stats = self.manager.get_transcription_stats(invocation_context)

    expected = {
        'input_transcriptions': 2,
        'output_transcriptions': 1,
        'total_transcriptions': 3,
    }
    assert stats == expected

  def test_get_transcription_stats_missing_attributes(self):
    """Test getting transcription statistics when events don't have transcription attributes."""
    invocation_context = Mock()

    # Create mock events and explicitly set transcription attributes to None
    event1 = Mock()
    event1.input_transcription = None
    event1.output_transcription = None

    event2 = Mock()
    event2.input_transcription = None
    event2.output_transcription = None

    invocation_context.session.events = [event1, event2]

    stats = self.manager.get_transcription_stats(invocation_context)

    expected = {
        'input_transcriptions': 0,
        'output_transcriptions': 0,
        'total_transcriptions': 0,
    }
    assert stats == expected

  @pytest.mark.asyncio
  async def test_transcription_event_fields(self):
    """Test that transcription events have correct field values."""
    invocation_context = await testing_utils.create_invocation_context(
        testing_utils.create_test_agent()
    )

    # Set up mock session service
    mock_session_service = AsyncMock()
    invocation_context.session_service = mock_session_service

    # Create test transcription with specific content
    transcription = types.Transcription(
        text='Test transcription content', finished=True
    )

    # Handle input transcription
    await self.manager.handle_input_transcription(
        invocation_context, transcription
    )

    # Verify the event structure
    call_args = mock_session_service.append_event.call_args
    event = call_args[0][1]

    # Check all required fields are present
    assert hasattr(event, 'id')
    assert hasattr(event, 'invocation_id')
    assert hasattr(event, 'author')
    assert hasattr(event, 'input_transcription')
    assert hasattr(event, 'output_transcription')
    assert hasattr(event, 'timestamp')

    # Check values
    assert event.id is not None
    assert event.invocation_id == invocation_context.invocation_id
    assert event.author == 'user'
    assert event.input_transcription == transcription
    assert event.output_transcription is None

  @pytest.mark.asyncio
  async def test_transcription_with_different_data_types(self):
    """Test handling transcriptions with different data types."""
    invocation_context = await testing_utils.create_invocation_context(
        testing_utils.create_test_agent()
    )

    # Set up mock session service
    mock_session_service = AsyncMock()
    invocation_context.session_service = mock_session_service

    # Test with transcription that has basic fields only
    transcription = types.Transcription(
        text='Advanced transcription', finished=True
    )

    # Handle transcription
    await self.manager.handle_input_transcription(
        invocation_context, transcription
    )

    # Verify the transcription object is preserved as-is
    call_args = mock_session_service.append_event.call_args
    event = call_args[0][1]

    assert event.input_transcription == transcription
    assert event.input_transcription.text == 'Advanced transcription'