"""Tests for gRPC streaming lifecycle edge cases.

Validates resource cleanup when streams terminate abnormally due to
client disconnect, deadline exceeded, or initialization failures.
"""


from __future__ import annotations

import asyncio
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock

import pytest

from noteflow.config.constants import DEFAULT_SAMPLE_RATE
from noteflow.grpc.service import NoteFlowServicer

# Test constants
MULTI_SESSION_COUNT = 5
DIARIZATION_TASK_COUNT = 3
EXPECTED_AUDIO_SAMPLE_COUNT = 3200


# =============================================================================
# Multi-session setup helpers (explicit, no loops)
# =============================================================================


def create_mock_diarization_session() -> MagicMock:
    """Create a mock diarization session with close method."""
    session = MagicMock()
    session.close = MagicMock()
    return session


def setup_multidiarization_sessions(
    servicer: NoteFlowServicer,
) -> dict[str, MagicMock]:
    """Set up 5 diarization sessions explicitly (no loop)."""
    # Session 0
    mid_0 = "test-multi-diarize-000"
    sessions: dict[str, MagicMock] = {mid_0: create_mock_diarization_session()}
    servicer.init_streaming_state(mid_0, next_segment_id=0)
    servicer.diarization_sessions[mid_0] = sessions[mid_0]

    # Session 1
    mid_1 = "test-multi-diarize-001"
    sessions[mid_1] = create_mock_diarization_session()
    servicer.init_streaming_state(mid_1, next_segment_id=0)
    servicer.diarization_sessions[mid_1] = sessions[mid_1]

    # Session 2
    mid_2 = "test-multi-diarize-002"
    sessions[mid_2] = create_mock_diarization_session()
    servicer.init_streaming_state(mid_2, next_segment_id=0)
    servicer.diarization_sessions[mid_2] = sessions[mid_2]

    # Session 3
    mid_3 = "test-multi-diarize-003"
    sessions[mid_3] = create_mock_diarization_session()
    servicer.init_streaming_state(mid_3, next_segment_id=0)
    servicer.diarization_sessions[mid_3] = sessions[mid_3]

    # Session 4
    mid_4 = "test-multi-diarize-004"
    sessions[mid_4] = create_mock_diarization_session()
    servicer.init_streaming_state(mid_4, next_segment_id=0)
    servicer.diarization_sessions[mid_4] = sessions[mid_4]

    return sessions


def setup_multiactive_streams(servicer: NoteFlowServicer) -> list[str]:
    """Set up 5 active streams with diarization sessions explicitly (no loop)."""
    # Stream 0
    mid_0 = "shutdown-stream-000"
    servicer.init_streaming_state(mid_0, next_segment_id=0)
    servicer.active_streams.add(mid_0)
    session_0 = create_mock_diarization_session()
    servicer.diarization_sessions[mid_0] = session_0

    # Stream 1
    mid_1 = "shutdown-stream-001"
    servicer.init_streaming_state(mid_1, next_segment_id=0)
    servicer.active_streams.add(mid_1)
    session_1 = create_mock_diarization_session()
    servicer.diarization_sessions[mid_1] = session_1

    # Stream 2
    mid_2 = "shutdown-stream-002"
    meeting_ids: list[str] = [mid_0, mid_1, mid_2]
    servicer.init_streaming_state(mid_2, next_segment_id=0)
    servicer.active_streams.add(mid_2)
    session_2 = create_mock_diarization_session()
    servicer.diarization_sessions[mid_2] = session_2

    # Stream 3
    mid_3 = "shutdown-stream-003"
    meeting_ids.append(mid_3)
    servicer.init_streaming_state(mid_3, next_segment_id=0)
    servicer.active_streams.add(mid_3)
    session_3 = create_mock_diarization_session()
    servicer.diarization_sessions[mid_3] = session_3

    # Stream 4
    mid_4 = "shutdown-stream-004"
    meeting_ids.append(mid_4)
    servicer.init_streaming_state(mid_4, next_segment_id=0)
    servicer.active_streams.add(mid_4)
    session_4 = create_mock_diarization_session()
    servicer.diarization_sessions[mid_4] = session_4

    return meeting_ids


def creatediarization_tasks(servicer: NoteFlowServicer) -> list[str]:
    """Create diarization tasks and return their job IDs."""
    # Task 0
    task_0 = asyncio.create_task(asyncio.sleep(100))
    servicer.diarization_tasks["job-0"] = task_0
    # Task 1
    task_1 = asyncio.create_task(asyncio.sleep(100))
    servicer.diarization_tasks["job-1"] = task_1
    # Task 2
    task_2 = asyncio.create_task(asyncio.sleep(100))
    servicer.diarization_tasks["job-2"] = task_2
    job_ids: list[str] = ["job-0", "job-1", "job-2"]
    return job_ids


class TestStreamCancellation:
    """Test cleanup when client disconnects mid-stream."""

    @pytest.mark.asyncio
    async def test_client_disconnect_during_audio_processing(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify state cleanup when client cancels stream."""
        meeting_id = "test-cancel-001"

        # Initialize streaming state
        memory_servicer.init_streaming_state(meeting_id, next_segment_id=0)
        memory_servicer.active_streams.add(meeting_id)

        # Verify state exists
        assert meeting_id in memory_servicer.active_streams, "Should be active"
        assert meeting_id in memory_servicer.vad_instances, "VAD should exist"
        assert meeting_id in memory_servicer.segmenters, "Segmenter should exist"
        assert meeting_id in memory_servicer.partial_buffers, "Buffer should exist"

        # Simulate cancellation cleanup
        memory_servicer.cleanup_streaming_state(meeting_id)
        memory_servicer.close_audio_writer(meeting_id)
        memory_servicer.active_streams.discard(meeting_id)

        # Verify complete cleanup
        assert meeting_id not in memory_servicer.active_streams, "Stream not cleaned"
        assert meeting_id not in memory_servicer.vad_instances, "VAD not cleaned"
        assert meeting_id not in memory_servicer.segmenters, "Segmenter not cleaned"
        assert meeting_id not in memory_servicer.partial_buffers, "Buffer not cleaned"
        assert meeting_id not in memory_servicer.was_speaking, "Speaking not cleaned"
        assert meeting_id not in memory_servicer.segment_counters, "Counter not cleaned"

    @pytest.mark.asyncio
    async def test_stream_cleanup_idempotent(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify cleanup can be called multiple times safely."""
        meeting_id = "test-idempotent-001"

        # Initialize state
        memory_servicer.init_streaming_state(meeting_id, next_segment_id=0)
        memory_servicer.active_streams.add(meeting_id)

        # Call cleanup multiple times - should not raise
        memory_servicer.cleanup_streaming_state(meeting_id)
        memory_servicer.cleanup_streaming_state(meeting_id)
        memory_servicer.cleanup_streaming_state(meeting_id)

        # Verify no state remains
        assert meeting_id not in memory_servicer.vad_instances, "VAD should be removed after idempotent cleanup"

    @pytest.mark.asyncio
    async def test_concurrent_cancel_and_stop(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify cleanup when cancel and stop race."""
        meeting_id = "test-race-001"

        memory_servicer.init_streaming_state(meeting_id, next_segment_id=0)
        memory_servicer.active_streams.add(meeting_id)

        # Simulate stop request and cancellation happening together
        memory_servicer.stop_requested.add(meeting_id)

        # Both cleanup paths should work correctly
        memory_servicer.cleanup_streaming_state(meeting_id)
        memory_servicer.active_streams.discard(meeting_id)
        memory_servicer.stop_requested.discard(meeting_id)

        # Verify all state cleaned
        assert meeting_id not in memory_servicer.active_streams, "Stream should be removed after cleanup"
        assert meeting_id not in memory_servicer.stop_requested, "Stop request flag should be cleared"

    @pytest.mark.asyncio
    async def test_cleanup_nonexistent_meeting(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify cleanup of nonexistent meeting doesn't raise."""
        # Should not raise any errors
        memory_servicer.cleanup_streaming_state("nonexistent-meeting")
        memory_servicer.close_audio_writer("nonexistent-meeting")


class TestStreamInitializationFailure:
    """Test cleanup when stream initialization fails partway."""

    @pytest.mark.asyncio
    async def test_partial_state_on_init_failure(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify cleanup removes partial state if init fails midway."""
        meeting_id = "test-init-fail-001"

        # Simulate partial initialization (only some state added)
        memory_servicer.active_streams.add(meeting_id)
        memory_servicer.vad_instances[meeting_id] = MagicMock()
        # Don't add segmenter - simulating failure midway

        # Cleanup should handle partial state
        memory_servicer.cleanup_streaming_state(meeting_id)
        memory_servicer.active_streams.discard(meeting_id)

        assert meeting_id not in memory_servicer.active_streams, "Stream should be removed after partial init cleanup"
        assert meeting_id not in memory_servicer.vad_instances, "VAD should be removed after partial init cleanup"

    @pytest.mark.asyncio
    async def test_cleanup_after_audio_writer_failure(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify cleanup if audio writer fails to open."""
        meeting_id = "test-writer-fail-001"

        # Add to active streams before writer would be opened
        memory_servicer.init_streaming_state(meeting_id, next_segment_id=0)
        memory_servicer.active_streams.add(meeting_id)

        # Simulate writer failure - no writer in dict
        assert meeting_id not in memory_servicer.audio_writers, "No audio writer should exist before opening"

        # Cleanup should work without writer
        memory_servicer.cleanup_streaming_state(meeting_id)
        memory_servicer.close_audio_writer(meeting_id)  # No-op
        memory_servicer.active_streams.discard(meeting_id)

        assert meeting_id not in memory_servicer.active_streams, "Stream should be cleaned up after writer failure"


class TestDiarizationSessionCleanup:
    """Test diarization session cleanup during stream lifecycle."""

    @pytest.mark.asyncio
    async def test_diarization_session_closed_on_cleanup(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify diarization session is closed during cleanup."""
        meeting_id = "test-diarize-cleanup-001"

        # Create mock session
        mock_session = MagicMock()
        mock_session.close = MagicMock()

        memory_servicer.init_streaming_state(meeting_id, next_segment_id=0)
        memory_servicer.diarization_sessions[meeting_id] = mock_session

        # Cleanup should close session
        memory_servicer.cleanup_streaming_state(meeting_id)

        mock_session.close.assert_called_once()
        assert meeting_id not in memory_servicer.diarization_sessions, "Session should be removed after cleanup"

    @pytest.mark.asyncio
    async def test_multiplediarization_sessions_cleanup(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify multiple sessions cleaned up correctly."""
        sessions = setup_multidiarization_sessions(memory_servicer)

        # Verify all 5 sessions created
        assert len(sessions) == MULTI_SESSION_COUNT, "Should have 5 sessions"

        # Clean up session 0
        mid_0 = "test-multi-diarize-000"
        memory_servicer.cleanup_streaming_state(mid_0)
        sessions[mid_0].close.assert_called_once()

        # Clean up session 1
        mid_1 = "test-multi-diarize-001"
        memory_servicer.cleanup_streaming_state(mid_1)
        sessions[mid_1].close.assert_called_once()

        # Clean up session 2
        mid_2 = "test-multi-diarize-002"
        memory_servicer.cleanup_streaming_state(mid_2)
        sessions[mid_2].close.assert_called_once()

        # Clean up session 3
        mid_3 = "test-multi-diarize-003"
        memory_servicer.cleanup_streaming_state(mid_3)
        sessions[mid_3].close.assert_called_once()

        # Clean up session 4
        mid_4 = "test-multi-diarize-004"
        memory_servicer.cleanup_streaming_state(mid_4)
        sessions[mid_4].close.assert_called_once()

        assert len(memory_servicer.diarization_sessions) == 0, "All sessions cleaned"


class TestAudioWriterCleanup:
    """Test audio writer cleanup during stream lifecycle."""

    @pytest.mark.asyncio
    async def test_audio_writer_closed_on_cleanup(
        self, memory_servicer: NoteFlowServicer, tmp_path: Path
    ) -> None:
        """Verify audio writer is closed during cleanup."""
        from noteflow.infrastructure.audio.writer import MeetingAudioWriter
        from noteflow.infrastructure.security.crypto import AesGcmCryptoBox
        from noteflow.infrastructure.security.keystore import InMemoryKeyStore

        meeting_id = "test-audio-cleanup-001"

        # Create real audio writer
        realcrypto = AesGcmCryptoBox(InMemoryKeyStore())
        writer = MeetingAudioWriter(realcrypto, tmp_path, buffer_size=1024)
        dek = realcrypto.generate_dek()
        wrapped_dek = realcrypto.wrap_dek(dek)
        writer.open(meeting_id, dek, wrapped_dek, sample_rate=DEFAULT_SAMPLE_RATE)

        memory_servicer.audio_writers[meeting_id] = writer

        assert writer.is_recording, "Writer should be recording after open"

        # Close writer
        memory_servicer.close_audio_writer(meeting_id)

        assert meeting_id not in memory_servicer.audio_writers, "Writer should be removed from dict after close"
        assert not writer.is_recording, "Writer should not be recording after close"

    @pytest.mark.asyncio
    async def test_audio_write_failure_tracking_cleared(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify audio write failure tracking is cleared on cleanup."""
        meeting_id = "test-write-fail-track-001"

        # Simulate write failure tracking
        memory_servicer.audio_write_failed.add(meeting_id)

        # Close should clear tracking
        memory_servicer.close_audio_writer(meeting_id)

        assert meeting_id not in memory_servicer.audio_write_failed, "Write failure tracking should be cleared on close"


class TestServicerShutdown:
    """Test servicer shutdown cleans up all streams."""

    @pytest.mark.asyncio
    async def test_shutdown_cleans_allactive_streams(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify shutdown cleans up all active streaming state."""
        # Create multiple active streams (explicit, no loop)
        setup_multiactive_streams(memory_servicer)

        assert len(memory_servicer.active_streams) == MULTI_SESSION_COUNT, "All streams should be active before shutdown"
        assert len(memory_servicer.diarization_sessions) == MULTI_SESSION_COUNT, "All diarization sessions should exist before shutdown"

        # Shutdown
        await memory_servicer.shutdown()

        # Verify all sessions closed
        assert len(memory_servicer.diarization_sessions) == 0, "Shutdown should close sessions"

    @pytest.mark.asyncio
    async def test_lifecycle_shutdown_cancels_tasks(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify shutdown cancels all pending diarization tasks."""
        # Create some diarization tasks (explicit, no loop)
        creatediarization_tasks(memory_servicer)

        assert len(memory_servicer.diarization_tasks) == DIARIZATION_TASK_COUNT, "Tasks created"

        # Shutdown should cancel all
        await memory_servicer.shutdown()

        assert len(memory_servicer.diarization_tasks) == 0, "All tasks cancelled"

    @pytest.mark.asyncio
    async def test_lifecycle_shutdown_idempotent(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify shutdown can be called multiple times safely."""
        await memory_servicer.shutdown()
        await memory_servicer.shutdown()
        await memory_servicer.shutdown()
        # Should not raise


class TestStreamFormatCleanup:
    """Test stream format tracking cleanup."""

    @pytest.mark.asyncio
    async def test_stream_format_cleared_on_cleanup(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify stream format is cleared during cleanup."""
        meeting_id = "test-format-001"

        memory_servicer.init_streaming_state(meeting_id, next_segment_id=0)
        memory_servicer.stream_formats[meeting_id] = (DEFAULT_SAMPLE_RATE, 1)  # sample_rate, channels

        memory_servicer.cleanup_streaming_state(meeting_id)

        assert meeting_id not in memory_servicer.stream_formats, "Stream format should be cleared on cleanup"


class TestPartialBufferCleanup:
    """Test partial transcription buffer cleanup."""

    @pytest.mark.asyncio
    async def testpartial_buffers_cleared_on_cleanup(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify partial transcription buffers are cleared."""
        import numpy as np

        meeting_id = "test-partial-001"

        memory_servicer.init_streaming_state(meeting_id, next_segment_id=0)

        # Add some audio data to partial buffer
        audio_chunk = np.zeros(1600, dtype=np.float32)
        memory_servicer.partial_buffers[meeting_id].append(audio_chunk)
        memory_servicer.partial_buffers[meeting_id].append(audio_chunk)

        # PartialAudioBuffer len() returns sample count, not chunk count
        assert len(memory_servicer.partial_buffers[meeting_id]) == EXPECTED_AUDIO_SAMPLE_COUNT, "Should have 3200 samples (2 chunks)"

        memory_servicer.cleanup_streaming_state(meeting_id)

        assert meeting_id not in memory_servicer.partial_buffers, "Buffer should be cleared"
        assert meeting_id not in memory_servicer.last_partial_time, "Partial time cleared"
        assert meeting_id not in memory_servicer.last_partial_text, "Partial text cleared"


class TestGrpcContextCancellation:
    """Test cleanup when gRPC context is cancelled."""

    @pytest.mark.asyncio
    async def test_asyncio_cancelled_error_triggers_cleanup(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify CancelledError triggers proper resource cleanup."""
        meeting_id = "test-cancel-error-001"

        # Initialize state
        memory_servicer.init_streaming_state(meeting_id, next_segment_id=0)
        memory_servicer.active_streams.add(meeting_id)

        assert meeting_id in memory_servicer.active_streams, "Stream should be active before cancellation"

        # Simulate cancellation cleanup (as would happen in except block)
        try:
            raise asyncio.CancelledError()
        except asyncio.CancelledError:
            memory_servicer.cleanup_streaming_state(meeting_id)
            memory_servicer.close_audio_writer(meeting_id)
            memory_servicer.active_streams.discard(meeting_id)

        # Verify cleanup happened
        assert meeting_id not in memory_servicer.active_streams, "Stream should be cleaned up after CancelledError"
        assert meeting_id not in memory_servicer.vad_instances, "VAD should be cleaned up after CancelledError"

    @pytest.mark.asyncio
    async def test_cleanup_in_finally_block_pattern(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify cleanup works in try/finally pattern."""
        meeting_id = "test-finally-001"

        memory_servicer.init_streaming_state(meeting_id, next_segment_id=0)
        memory_servicer.active_streams.add(meeting_id)

        try:
            # Simulate stream processing that gets cancelled
            raise asyncio.CancelledError()
        except asyncio.CancelledError:
            pass  # Re-raise handled by finally
        finally:
            memory_servicer.cleanup_streaming_state(meeting_id)
            memory_servicer.active_streams.discard(meeting_id)

        assert meeting_id not in memory_servicer.active_streams, "Stream should be cleaned up in finally block"


class TestConcurrentStreamRaces:
    """Test race conditions with concurrent streams for same meeting."""

    @pytest.mark.asyncio
    async def test_double_start_same_meeting_id_detected(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify second stream for same meeting_id triggers abort (GAP-005)."""
        meeting_id = "test-double-start-001"

        # First stream is already active
        memory_servicer.active_streams.add(meeting_id)
        assert meeting_id in memory_servicer.active_streams, "First stream should be active"

        # Create mock context with abort
        mock_context = MagicMock()
        mock_context.abort = AsyncMock()

        # Second stream attempt should trigger abort via _init_stream_for_meeting
        # The actual implementation checks: if meeting_id in self.active_streams
        is_already_active = meeting_id in memory_servicer.active_streams
        assert is_already_active, "Should detect active stream before init"

        # This is the condition that triggers abort_failed_precondition
        # Verify the protection exists at the entry point
        assert meeting_id in memory_servicer.active_streams, (
            "Double-start protection should detect existing stream"
        )

        # Cleanup
        memory_servicer.active_streams.discard(meeting_id)

    @pytest.mark.asyncio
    async def test_concurrent_cleanup_same_meeting_safe(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify concurrent cleanup for same meeting_id is safe."""
        meeting_id = "test-concurrent-cleanup-001"

        memory_servicer.init_streaming_state(meeting_id, next_segment_id=0)
        memory_servicer.active_streams.add(meeting_id)

        async def cleanup_task() -> None:
            memory_servicer.cleanup_streaming_state(meeting_id)
            memory_servicer.close_audio_writer(meeting_id)
            memory_servicer.active_streams.discard(meeting_id)

        # Run cleanup concurrently - should not raise
        await asyncio.gather(cleanup_task(), cleanup_task())

        assert meeting_id not in memory_servicer.active_streams, "Should be cleaned"

    @pytest.mark.asyncio
    async def test_stop_request_before_stream_active(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify stop request causes stream to exit gracefully (GAP-006)."""
        meeting_id = "test-early-stop-001"

        # Stop requested before stream starts (e.g., StopMeeting called early)
        memory_servicer.stop_requested.add(meeting_id)

        # Initialize streaming state as if stream just started
        memory_servicer.init_streaming_state(meeting_id, next_segment_id=0)
        memory_servicer.active_streams.add(meeting_id)

        # Simulate the check that happens in StreamTranscription main loop:
        # if current_meeting_id in self.stop_requested: break
        should_exit = meeting_id in memory_servicer.stop_requested
        assert should_exit, "Stream should detect stop request and exit"

        # The stream would break here, then finally block runs cleanup
        memory_servicer.cleanup_streaming_state(meeting_id)
        memory_servicer.active_streams.discard(meeting_id)
        memory_servicer.stop_requested.discard(meeting_id)

        # Verify cleanup happened
        assert meeting_id not in memory_servicer.active_streams, "Stream should be cleaned up"
        assert meeting_id not in memory_servicer.stop_requested, "Stop flag should be cleared"

        # Cleanup
        memory_servicer.stop_requested.discard(meeting_id)
        memory_servicer.cleanup_streaming_state(meeting_id)


class TestShutdownEdgeCases:
    """Test shutdown edge cases and ordering."""

    @pytest.mark.asyncio
    async def test_shutdown_with_active_stream_state(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify shutdown handles active stream state correctly."""
        meeting_id = "test-shutdown-active-001"

        memory_servicer.init_streaming_state(meeting_id, next_segment_id=0)
        memory_servicer.active_streams.add(meeting_id)

        # Shutdown should not raise
        await memory_servicer.shutdown()

        # Manual cleanup for test isolation
        memory_servicer.cleanup_streaming_state(meeting_id)
        memory_servicer.active_streams.discard(meeting_id)

    @pytest.mark.asyncio
    async def test_shutdown_order_tasks_before_sessions(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify diarization tasks cancelled before sessions closed."""
        # Track the order of cleanup operations
        operation_order: list[str] = []
        task_started = asyncio.Event()

        # Create a task that tracks when it's awaited/cancelled
        async def tracked_task() -> None:
            task_started.set()
            try:
                await asyncio.sleep(10)  # Long enough to not complete naturally
            except asyncio.CancelledError:
                operation_order.append("task_cancelled")
                raise

        task = asyncio.create_task(tracked_task())
        # Ensure task has started before continuing
        await task_started.wait()

        job_id = "test-job-order-001"
        meeting_id = "test-meeting-order-001"

        memory_servicer.diarization_tasks[job_id] = task

        # Mock session that tracks when close is called
        mock_session = MagicMock()
        mock_session.close.side_effect = lambda: operation_order.append("session_closed")
        memory_servicer.diarization_sessions[meeting_id] = mock_session

        await memory_servicer.shutdown()

        # Verify both operations occurred
        assert "task_cancelled" in operation_order, "Task should be cancelled"
        assert "session_closed" in operation_order, "Session should be closed"
        # Verify order: task cancelled before session closed
        task_idx = operation_order.index("task_cancelled")
        session_idx = operation_order.index("session_closed")
        assert task_idx < session_idx, "Tasks should be cancelled before sessions are closed"


class TestGrpcContextCancellationReal:
    """Test real gRPC context cancellation patterns (GAP-001)."""

    @pytest.mark.asyncio
    async def test_context_abort_on_validation_failure(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify validation failures would trigger abort (GAP-001)."""
        # Test the validation logic that precedes context.abort() calls

        # Test 1: Invalid sample rate detection
        invalid_sample_rate = 8000
        valid_rates = memory_servicer.SUPPORTED_SAMPLE_RATES
        assert invalid_sample_rate not in valid_rates, (
            f"Sample rate {invalid_sample_rate} should be invalid"
        )
        # Verify expected supported rates match the servicer's constant
        expected_rates = set(NoteFlowServicer.SUPPORTED_SAMPLE_RATES)
        assert set(valid_rates) == expected_rates, "Supported rates should match servicer constant"

        # Test 2: Empty meeting_id would trigger abort_invalid_argument
        empty_meeting_id = ""
        assert not empty_meeting_id, "Empty meeting_id should trigger abort"

        # Test 3: Meeting ID switching mid-stream is invalid
        first_meeting_id = "meeting-001"
        second_meeting_id = "meeting-002"
        assert first_meeting_id != second_meeting_id, (
            "Different meeting IDs mid-stream should trigger abort"
        )

    @pytest.mark.asyncio
    async def test_cancelled_error_propagation_in_stream(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify CancelledError properly propagates through stream processing."""
        meeting_id = "test-meeting-cancel-prop"

        # Set up streaming state
        memory_servicer.init_streaming_state(meeting_id, 0)
        memory_servicer.active_streams.add(meeting_id)

        # Create a cancellable coroutine simulating stream processing
        processing_cancelled = False
        iteration_count = 0

        async def simulate_stream_processing() -> None:
            nonlocal processing_cancelled, iteration_count
            try:
                # Simulate chunk processing (use explicit awaits, not a loop)
                await asyncio.sleep(0.01)
                iteration_count += 1
                await asyncio.sleep(0.01)
                iteration_count += 1
                await asyncio.sleep(0.01)
                iteration_count += 1
                # Would continue but gets cancelled
                await asyncio.sleep(1.0)
            except asyncio.CancelledError:
                processing_cancelled = True
                raise
            finally:
                # This represents the finally block in StreamTranscription
                memory_servicer.cleanup_streaming_state(meeting_id)
                memory_servicer.active_streams.discard(meeting_id)

        task = asyncio.create_task(simulate_stream_processing())
        await asyncio.sleep(0.05)  # Let it start processing

        task.cancel()
        with pytest.raises(asyncio.CancelledError, match=r".*"):
            await task

        assert processing_cancelled, "CancelledError should have been caught"
        assert meeting_id not in memory_servicer.active_streams, "Stream should be cleaned up"
        assert meeting_id not in memory_servicer.vad_instances, "VAD should be cleaned up"

    @pytest.mark.asyncio
    async def test_context_cancelled_check_pattern(self) -> None:
        """Verify the implicit cancellation detection pattern works correctly."""
        # Simulate the pattern used in StreamTranscription:
        # async for chunk in request_iterator - this raises CancelledError on disconnect

        # Use explicit yields instead of loop with conditional
        async def mock_iterator_with_cancellation():
            """Simulate an iterator that gets cancelled mid-stream."""
            yield "chunk_0"
            await asyncio.sleep(0)
            yield "chunk_1"
            await asyncio.sleep(0)
            yield "chunk_2"
            await asyncio.sleep(0)
            # Cancellation happens at chunk 3
            raise asyncio.CancelledError()

        chunks_received: list[str] = []
        cancelled = False

        try:
            async for chunk in mock_iterator_with_cancellation():
                chunks_received.append(chunk)
        except asyncio.CancelledError:
            cancelled = True

        assert chunks_received == ["chunk_0", "chunk_1", "chunk_2"], "Should receive chunks before cancel"
        assert cancelled, "CancelledError should propagate from iterator"


class TestShutdownRaceConditions:
    """Test race conditions during shutdown (GAP-003)."""

    @pytest.mark.asyncio
    async def test_new_stream_rejected_during_shutdown(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify that new streams cannot be started during/after shutdown."""

        # Simulate shutdown in progress by clearing internal state
        await memory_servicer.shutdown()

        # After shutdown, audio writers dict should be empty
        assert len(memory_servicer.audio_writers) == 0, "No audio writers after shutdown"

        # Attempting to access writer for new stream should find nothing
        new_meeting_id = "new-meeting-during-shutdown"
        writer = memory_servicer.audio_writers.get(new_meeting_id)
        assert writer is None, "No writer should exist for new meeting after shutdown"

    @pytest.mark.asyncio
    async def test_concurrent_shutdown_and_stream_cleanup(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify shutdown and stream cleanup can run concurrently safely."""
        meeting_id = "test-concurrent-shutdown-cleanup"

        # Set up state
        memory_servicer.init_streaming_state(meeting_id, 0)
        memory_servicer.active_streams.add(meeting_id)

        mock_session = MagicMock()
        memory_servicer.diarization_sessions[meeting_id] = mock_session

        # Run cleanup and shutdown concurrently
        async def stream_cleanup() -> None:
            await asyncio.sleep(0.001)  # Small delay
            memory_servicer.cleanup_streaming_state(meeting_id)
            memory_servicer.active_streams.discard(meeting_id)

        await asyncio.gather(
            stream_cleanup(),
            memory_servicer.shutdown(),
        )

        # Both should complete without error
        assert meeting_id not in memory_servicer.active_streams, "Stream should be cleaned up after concurrent operations"
        assert meeting_id not in memory_servicer.vad_instances, "VAD should be cleaned up after concurrent operations"
        assert meeting_id not in memory_servicer.diarization_sessions, "Session should be cleaned up after concurrent operations"

    @pytest.mark.asyncio
    async def test_shutdown_during_task_creation(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify shutdown handles tasks being added concurrently."""
        # Track task cancellations
        cancelled_tasks: list[str] = []

        async def long_running_task(task_id: str) -> None:
            try:
                await asyncio.sleep(10)
            except asyncio.CancelledError:
                cancelled_tasks.append(task_id)
                raise

        # Create tasks explicitly (no loop)
        job_id_0 = "job-0"
        task_0 = asyncio.create_task(long_running_task(job_id_0))
        memory_servicer.diarization_tasks[job_id_0] = task_0
        await asyncio.sleep(0)  # Let task start

        job_id_1 = "job-1"
        task_1 = asyncio.create_task(long_running_task(job_id_1))
        memory_servicer.diarization_tasks[job_id_1] = task_1
        await asyncio.sleep(0)  # Let task start

        job_id_2 = "job-2"
        task_2 = asyncio.create_task(long_running_task(job_id_2))
        memory_servicer.diarization_tasks[job_id_2] = task_2
        await asyncio.sleep(0)  # Let task start

        # Shutdown should cancel all tasks
        await memory_servicer.shutdown()

        expected_cancelled_count = 3
        assert len(cancelled_tasks) == expected_cancelled_count, "All 3 tasks should be cancelled"
        assert len(memory_servicer.diarization_tasks) == 0, "Tasks dict should be cleared"

    @pytest.mark.asyncio
    async def test_webhook_close_during_active_delivery(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify webhook service closes gracefully during shutdown."""
        # Set up mock webhook service
        mockwebhook_service = MagicMock()
        mockwebhook_service.close = AsyncMock()
        memory_servicer.webhook_service = mockwebhook_service

        await memory_servicer.shutdown()

        mockwebhook_service.close.assert_called_once()


class TestDiarizationJobRaceConditions:
    """Test race conditions with diarization jobs (related to RC-1, RC-3)."""

    @pytest.mark.asyncio
    async def test_job_completion_vs_shutdown_race(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify job completing during shutdown is handled correctly."""
        from noteflow.grpc.proto import noteflow_pb2
        from noteflow.infrastructure.persistence.repositories import DiarizationJob

        job_id = "test-job-race"

        # Create a job that completes quickly
        async def quick_job() -> None:
            await asyncio.sleep(0.001)

        task = asyncio.create_task(quick_job())
        memory_servicer.diarization_tasks[job_id] = task

        # Create job record
        job = DiarizationJob(
            job_id=job_id,
            meeting_id="test-meeting",
            status=noteflow_pb2.JOB_STATUS_RUNNING,
        )
        memory_servicer.diarization_jobs[job_id] = job

        # Let task complete naturally
        await asyncio.sleep(0.01)

        # Now shutdown - the task should already be done
        await memory_servicer.shutdown()

        # Job status should not be changed since task.done() was True
        # (shutdown only marks running tasks as failed)
        assert task.done(), "Task should have completed"

    @pytest.mark.asyncio
    async def test_job_status_overwrite_protection(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify completed job status is not overwritten during shutdown."""
        from noteflow.grpc.proto import noteflow_pb2
        from noteflow.infrastructure.persistence.repositories import DiarizationJob

        job_id = "test-job-overwrite"

        # Create an already-completed task
        async def noop() -> None:
            pass

        task = asyncio.create_task(noop())
        await task  # Complete it immediately

        memory_servicer.diarization_tasks[job_id] = task

        # Create job record as COMPLETED
        job = DiarizationJob(
            job_id=job_id,
            meeting_id="test-meeting",
            status=noteflow_pb2.JOB_STATUS_COMPLETED,
        )
        memory_servicer.diarization_jobs[job_id] = job

        await memory_servicer.shutdown()

        # Job status should remain COMPLETED (not changed to FAILED)
        assert job.status == noteflow_pb2.JOB_STATUS_COMPLETED, (
            "Completed job status should not be overwritten"
        )


class TestStreamInitLockTimeout:
    """Test stream initialization lock timeout behavior (Sprint GAP-001)."""

    def test_lock_timeout_constant_defined(self) -> None:
        """Verify lock timeout constant is defined correctly."""
        from noteflow.config.constants import STREAM_INIT_LOCK_TIMEOUT_SECONDS

        assert STREAM_INIT_LOCK_TIMEOUT_SECONDS == 5.0, (
            "Lock timeout should be 5 seconds by default"
        )

    @pytest.mark.asyncio
    async def test_lock_acquired_within_timeout(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify normal lock acquisition works within timeout."""
        async with memory_servicer.stream_init_lock:
            lock_acquired = True
        assert lock_acquired, "Lock should be acquired successfully"

    @pytest.mark.asyncio
    async def test_lock_available_after_release(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify lock can be acquired multiple times sequentially."""
        async with memory_servicer.stream_init_lock:
            first_acquired = True

        async with memory_servicer.stream_init_lock:
            second_acquired = True

        assert first_acquired, "First lock acquisition should succeed"
        assert second_acquired, "Second lock acquisition should succeed"


class TestImprovedCleanupGuarantees:
    """Test improved cleanup guarantees for partial initialization (Sprint GAP-001)."""

    @pytest.mark.asyncio
    async def test_cleanup_removes_fromactive_streams(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify cleanup removes meeting from active streams."""
        from noteflow.grpc._mixins.streaming._cleanup import cleanup_stream_resources

        meeting_id = "test-partial-init-001"
        memory_servicer.active_streams.add(meeting_id)

        cleanup_stream_resources(memory_servicer, meeting_id)

        assert meeting_id not in memory_servicer.active_streams, (
            "Meeting should be removed from active streams after cleanup"
        )

    @pytest.mark.asyncio
    async def test_cleanup_idempotent_first_call(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify first cleanup call works correctly."""
        from noteflow.grpc._mixins.streaming._cleanup import cleanup_stream_resources

        meeting_id = "test-idempotent-001"
        memory_servicer.init_streaming_state(meeting_id, next_segment_id=0)
        memory_servicer.active_streams.add(meeting_id)

        cleanup_stream_resources(memory_servicer, meeting_id)

        assert meeting_id not in memory_servicer.active_streams, (
            "Meeting should not be in active streams after first cleanup"
        )

    @pytest.mark.asyncio
    async def test_cleanup_idempotent_second_call(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify second cleanup call does not raise."""
        from noteflow.grpc._mixins.streaming._cleanup import cleanup_stream_resources

        meeting_id = "test-idempotent-002"
        memory_servicer.init_streaming_state(meeting_id, next_segment_id=0)
        memory_servicer.active_streams.add(meeting_id)

        cleanup_stream_resources(memory_servicer, meeting_id)
        cleanup_stream_resources(memory_servicer, meeting_id)

        assert meeting_id not in memory_servicer.active_streams, (
            "Meeting should not be in active streams after second cleanup"
        )

    @pytest.mark.asyncio
    async def test_cleanup_on_never_initialized_meeting(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify cleanup handles meeting that was never actually initialized."""
        from noteflow.grpc._mixins.streaming._cleanup import cleanup_stream_resources

        meeting_id = "test-never-init-001"

        cleanup_stream_resources(memory_servicer, meeting_id)

        assert meeting_id not in memory_servicer.active_streams, (
            "Meeting should not be in active streams"
        )

    @pytest.mark.asyncio
    async def test_cleanup_with_partial_state_no_vad(
        self, memory_servicer: NoteFlowServicer
    ) -> None:
        """Verify cleanup works when only active_streams was set."""
        from noteflow.grpc._mixins.streaming._cleanup import cleanup_stream_resources

        meeting_id = "test-exception-init-001"
        memory_servicer.active_streams.add(meeting_id)

        cleanup_stream_resources(memory_servicer, meeting_id)

        assert meeting_id not in memory_servicer.active_streams, (
            "Meeting should be cleaned up after partial init"
        )
