"""Tests for partial transcription in the gRPC service."""

from __future__ import annotations

import time
from collections.abc import AsyncIterator
from dataclasses import dataclass
from typing import TYPE_CHECKING
from unittest.mock import MagicMock

import numpy as np
import pytest
from numpy.typing import NDArray

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


@dataclass
class MockAsrResult:
    """Mock ASR transcription result."""

    text: str
    start: float = 0.0
    end: float = 1.0
    language: str = "en"
    language_probability: float = 0.99
    avg_logprob: float = -0.5
    no_speech_prob: float = 0.01


def _create_mockasr_engine(transcribe_results: list[str] | None = None) -> MagicMock:
    """Create mock ASR engine with configurable transcription results."""
    engine = MagicMock()
    engine.is_loaded = True
    engine.model_size = "base"

    results = transcribe_results or ["Test transcription"]

    def _transcribe(_audio: NDArray[np.float32]) -> list[MockAsrResult]:
        return [MockAsrResult(text=text) for text in results]

    async def _transcribe_async(
        _audio: NDArray[np.float32],
        _language: str | None = None,
    ) -> list[MockAsrResult]:
        return [MockAsrResult(text=text) for text in results]

    engine.transcribe = _transcribe
    engine.transcribe_async = _transcribe_async
    return engine


class ExposedServicer(NoteFlowServicer):
    """Subclass of NoteFlowServicer with exposed protected methods for testing.

    This subclass provides public wrapper methods for protected methods that
    tests need to call directly, eliminating type errors from protected member access.
    """

    if TYPE_CHECKING:
        # Type stubs for inherited protected methods (used by wrapper methods)
        async def _maybe_emit_partial(
            self,
            meeting_id: str,
        ) -> noteflow_pb2.TranscriptUpdate | None: ...

        def _clear_partial_buffer(self, meeting_id: str) -> None: ...

        async def _process_audio_with_vad(
            self,
            meeting_id: str,
            audio: NDArray[np.float32],
        ) -> AsyncIterator[noteflow_pb2.TranscriptUpdate]:
            yield noteflow_pb2.TranscriptUpdate()

    async def maybe_emit_partial(
        self,
        meeting_id: str,
    ) -> noteflow_pb2.TranscriptUpdate | None:
        """Public wrapper for _maybe_emit_partial for testing."""
        return await self._maybe_emit_partial(meeting_id)

    def clear_partial_buffer(self, meeting_id: str) -> None:
        """Public wrapper for _clear_partial_buffer for testing."""
        self._clear_partial_buffer(meeting_id)

    async def process_audio_with_vad(
        self,
        meeting_id: str,
        audio: NDArray[np.float32],
    ) -> AsyncIterator[noteflow_pb2.TranscriptUpdate]:
        """Public wrapper for _process_audio_with_vad for testing."""
        async for update in self._process_audio_with_vad(meeting_id, audio):
            yield update


class TestPartialTranscriptionState:
    """Tests for partial transcription state initialization."""

    def testinit_streaming_state_creates_partial_buffer(self) -> None:
        """Initialize streaming state should create empty partial buffer."""
        servicer: ExposedServicer = ExposedServicer()

        servicer.init_streaming_state("meeting-123", next_segment_id=0)

        assert "meeting-123" in servicer.partial_buffers, "partial buffer should be created"
        assert servicer.partial_buffers["meeting-123"].is_empty, "Buffer should start empty"

    def testinit_streaming_state_createslast_partial_time(self) -> None:
        """Initialize streaming state should set last partial time to now."""
        servicer: ExposedServicer = ExposedServicer()
        before = time.time()

        servicer.init_streaming_state("meeting-123", next_segment_id=0)

        assert "meeting-123" in servicer.last_partial_time, "last partial time should be set"
        assert servicer.last_partial_time["meeting-123"] >= before, "time should be current"

    def testinit_streaming_state_creates_empty_last_text(self) -> None:
        """Initialize streaming state should set last partial text to empty."""
        servicer: ExposedServicer = ExposedServicer()

        servicer.init_streaming_state("meeting-123", next_segment_id=0)

        assert "meeting-123" in servicer.last_partial_text, "last partial text should exist"
        assert servicer.last_partial_text["meeting-123"] == "", "text should start empty"

    def testcleanup_streaming_state_removes_partial_state(self) -> None:
        """Cleanup streaming state should remove all partial-related state."""
        servicer: ExposedServicer = ExposedServicer()
        servicer.init_streaming_state("meeting-123", next_segment_id=0)

        servicer.cleanup_streaming_state("meeting-123")

        assert "meeting-123" not in servicer.partial_buffers, "buffer should be removed"
        assert "meeting-123" not in servicer.last_partial_time, "time should be removed"
        assert "meeting-123" not in servicer.last_partial_text, "text should be removed"


class TestClearPartialBuffer:
    """Tests for _clear_partial_buffer method."""

    def test_clear_partial_buffer_empties_buffer(self) -> None:
        """Clear partial buffer should empty the audio buffer."""
        from noteflow.infrastructure.audio.partial_buffer import PartialAudioBuffer

        servicer: ExposedServicer = ExposedServicer()
        buffer = PartialAudioBuffer(sample_rate=DEFAULT_SAMPLE_RATE)
        buffer.append(np.zeros(1600, dtype=np.float32))
        servicer.partial_buffers["meeting-123"] = buffer

        servicer.clear_partial_buffer("meeting-123")

        assert servicer.partial_buffers["meeting-123"].is_empty, "Buffer should be empty after clear"

    def test_clear_partial_buffer_resets_last_text(self) -> None:
        """Clear partial buffer should reset last partial text."""
        servicer: ExposedServicer = ExposedServicer()
        servicer.last_partial_text["meeting-123"] = "Previous partial"

        servicer.clear_partial_buffer("meeting-123")

        assert servicer.last_partial_text["meeting-123"] == "", "text should be reset"

    def test_clear_partial_buffer_updates_time(self) -> None:
        """Clear partial buffer should update last partial time."""
        servicer: ExposedServicer = ExposedServicer()
        servicer.last_partial_time["meeting-123"] = 0.0
        before = time.time()

        servicer.clear_partial_buffer("meeting-123")

        assert servicer.last_partial_time["meeting-123"] >= before, "time should be updated"

    def test_clear_partial_buffer_handles_missing_meeting(self) -> None:
        """Clear partial buffer should handle missing meeting gracefully."""
        servicer: ExposedServicer = ExposedServicer()

        servicer.clear_partial_buffer("nonexistent")  # Should not raise


class TestMaybeEmitPartial:
    """Tests for _maybe_emit_partial method."""

    @pytest.mark.asyncio
    async def test_returns_none_when_asr_not_loaded(self) -> None:
        """Return None when ASR engine is not loaded."""
        servicer: ExposedServicer = ExposedServicer()
        servicer.init_streaming_state("meeting-123", next_segment_id=0)

        result: noteflow_pb2.TranscriptUpdate | None = await servicer.maybe_emit_partial("meeting-123")

        assert result is None, "should return None when ASR not loaded"

    @pytest.mark.asyncio
    async def test_returns_none_when_cadence_not_reached(self) -> None:
        """Return None when not enough time has passed since last partial."""
        engine = _create_mockasr_engine(["Test"])
        servicer: ExposedServicer = ExposedServicer(asr_engine=engine)
        servicer.init_streaming_state("meeting-123", next_segment_id=0)
        # Set last time to now (cadence not reached)
        servicer.last_partial_time["meeting-123"] = time.time()
        # Add some audio
        audio: NDArray[np.float32] = np.ones(DEFAULT_SAMPLE_RATE, dtype=np.float32) * 0.1  # 1 second of audio
        servicer.partial_buffers["meeting-123"].append(audio)

        result: noteflow_pb2.TranscriptUpdate | None = await servicer.maybe_emit_partial("meeting-123")

        assert result is None, "should return None when cadence not reached"

    @pytest.mark.asyncio
    async def test_returns_none_when_buffer_empty(self) -> None:
        """Return None when partial buffer is empty."""
        engine = _create_mockasr_engine(["Test"])
        servicer: ExposedServicer = ExposedServicer(asr_engine=engine)
        servicer.init_streaming_state("meeting-123", next_segment_id=0)
        # Set last time to past (cadence reached)
        servicer.last_partial_time["meeting-123"] = time.time() - 10.0

        result: noteflow_pb2.TranscriptUpdate | None = await servicer.maybe_emit_partial("meeting-123")

        assert result is None, "should return None when buffer empty"

    @pytest.mark.asyncio
    async def test_returns_none_when_audio_too_short(self) -> None:
        """Return None when buffered audio is less than minimum."""
        engine = _create_mockasr_engine(["Test"])
        servicer: ExposedServicer = ExposedServicer(asr_engine=engine)
        servicer.init_streaming_state("meeting-123", next_segment_id=0)
        servicer.last_partial_time["meeting-123"] = time.time() - 10.0
        # Add only 0.1 seconds of audio (minimum is 0.5s)
        audio: NDArray[np.float32] = np.ones(1600, dtype=np.float32) * 0.1  # 0.1 second
        servicer.partial_buffers["meeting-123"].append(audio)

        result: noteflow_pb2.TranscriptUpdate | None = await servicer.maybe_emit_partial("meeting-123")

        assert result is None, "should return None when audio too short"

    @pytest.mark.asyncio
    async def test_emits_partial_when_conditions_met(self) -> None:
        """Emit partial when cadence reached and sufficient audio buffered."""
        engine = _create_mockasr_engine(["Hello world"])
        servicer: ExposedServicer = ExposedServicer(asr_engine=engine)
        servicer.init_streaming_state("meeting-123", next_segment_id=0)
        servicer.last_partial_time["meeting-123"] = time.time() - 10.0
        # Add 1 second of audio (above minimum of 0.5s)
        audio: NDArray[np.float32] = np.ones(DEFAULT_SAMPLE_RATE, dtype=np.float32) * 0.1
        servicer.partial_buffers["meeting-123"].append(audio)

        result: noteflow_pb2.TranscriptUpdate | None = await servicer.maybe_emit_partial("meeting-123")

        assert result is not None, "Expected partial result when conditions met"
        assert result.update_type == 1, "Expected UPDATE_TYPE_PARTIAL"
        assert result.partial_text == "Hello world", "Expected transcribed text"
        assert result.meeting_id == "meeting-123", "Expected correct meeting_id"

    @pytest.mark.asyncio
    async def test_debounces_duplicate_text(self) -> None:
        """Return None when text is same as last partial (debounce)."""
        engine = _create_mockasr_engine(["Same text"])
        servicer: ExposedServicer = ExposedServicer(asr_engine=engine)
        servicer.init_streaming_state("meeting-123", next_segment_id=0)
        servicer.last_partial_time["meeting-123"] = time.time() - 10.0
        servicer.last_partial_text["meeting-123"] = "Same text"  # Same as transcription
        audio: NDArray[np.float32] = np.ones(DEFAULT_SAMPLE_RATE, dtype=np.float32) * 0.1
        servicer.partial_buffers["meeting-123"].append(audio)

        result: noteflow_pb2.TranscriptUpdate | None = await servicer.maybe_emit_partial("meeting-123")

        assert result is None, "should return None for duplicate text"

    @pytest.mark.asyncio
    async def test_updates_last_partial_state(self) -> None:
        """Emitting partial should update last text and time."""
        engine = _create_mockasr_engine(["New text"])
        servicer: ExposedServicer = ExposedServicer(asr_engine=engine)
        servicer.init_streaming_state("meeting-123", next_segment_id=0)
        servicer.last_partial_time["meeting-123"] = time.time() - 10.0
        audio: NDArray[np.float32] = np.ones(DEFAULT_SAMPLE_RATE, dtype=np.float32) * 0.1
        servicer.partial_buffers["meeting-123"].append(audio)
        before = time.time()

        await servicer.maybe_emit_partial("meeting-123")

        assert servicer.last_partial_text["meeting-123"] == "New text", "text should update"
        assert servicer.last_partial_time["meeting-123"] >= before, "time should update"


class TestPartialCadence:
    """Tests for partial transcription cadence constants."""

    def test_partial_cadence_is_positive(self) -> None:
        """Partial cadence should be a positive duration."""
        cadence: float = NoteFlowServicer.PARTIAL_CADENCE_SECONDS
        assert cadence > 0, "PARTIAL_CADENCE_SECONDS must be positive"

    def test_min_partial_audio_is_positive(self) -> None:
        """Minimum partial audio should be a positive duration."""
        min_audio: float = NoteFlowServicer.MIN_PARTIAL_AUDIO_SECONDS
        assert min_audio > 0, "MIN_PARTIAL_AUDIO_SECONDS must be positive"

    def test_min_partial_less_than_cadence(self) -> None:
        """Minimum partial audio should be less than cadence interval."""
        min_audio: float = NoteFlowServicer.MIN_PARTIAL_AUDIO_SECONDS
        cadence: float = NoteFlowServicer.PARTIAL_CADENCE_SECONDS
        assert min_audio < cadence, "MIN_PARTIAL_AUDIO_SECONDS must be less than cadence"


class TestPartialBufferAccumulation:
    """Tests for audio buffer accumulation during speech."""

    @pytest.mark.asyncio
    async def test_speech_audio_added_to_buffer(self) -> None:
        """Speech audio should be accumulated in partial buffer."""
        engine = _create_mockasr_engine()
        servicer: ExposedServicer = ExposedServicer(asr_engine=engine)
        servicer.init_streaming_state("meeting-123", next_segment_id=0)

        # Simulate speech detection by processing audio
        audio: NDArray[np.float32] = np.ones(1600, dtype=np.float32) * 0.1

        # Mock VAD to return True (is_speech)
        servicer.vad_instances["meeting-123"].process_chunk = MagicMock(return_value=True)

        updates: list[noteflow_pb2.TranscriptUpdate] = []
        async for update in servicer.process_audio_with_vad("meeting-123", audio):
            updates.append(update)

        # Buffer should have audio added
        assert len(servicer.partial_buffers["meeting-123"]) >= 1, "buffer should have audio"

    @pytest.mark.asyncio
    async def test_silence_does_not_add_to_buffer(self) -> None:
        """Silent audio should not be added to partial buffer."""
        engine = _create_mockasr_engine()
        servicer: ExposedServicer = ExposedServicer(asr_engine=engine)
        servicer.init_streaming_state("meeting-123", next_segment_id=0)

        audio: NDArray[np.float32] = np.zeros(1600, dtype=np.float32)  # Silence

        # Mock VAD to return False (is_silence)
        servicer.vad_instances["meeting-123"].process_chunk = MagicMock(return_value=False)

        updates: list[noteflow_pb2.TranscriptUpdate] = []
        async for update in servicer.process_audio_with_vad("meeting-123", audio):
            updates.append(update)

        # Buffer should still be empty
        assert servicer.partial_buffers["meeting-123"].is_empty, "Buffer should be empty"


class TestPartialIntegrationWithFinal:
    """Tests for partial buffer clearing when final segment emitted."""

    @pytest.mark.asyncio
    async def test_buffer_cleared_on_final_segment(self) -> None:
        """Partial buffer should be cleared when a final segment is produced."""
        servicer: ExposedServicer = ExposedServicer()
        servicer.init_streaming_state("meeting-123", next_segment_id=0)

        # Add some audio to buffer
        audio: NDArray[np.float32] = np.ones(DEFAULT_SAMPLE_RATE, dtype=np.float32) * 0.1
        servicer.partial_buffers["meeting-123"].append(audio)
        servicer.last_partial_text["meeting-123"] = "Some partial"

        # Clear buffer (simulates final segment emission)
        servicer.clear_partial_buffer("meeting-123")

        assert servicer.partial_buffers["meeting-123"].is_empty, "Buffer should be empty"
        assert servicer.last_partial_text["meeting-123"] == "", "text should be cleared"
