"""Streaming audio transcription mixin for gRPC service."""

from __future__ import annotations

from collections.abc import AsyncIterator
from typing import TYPE_CHECKING

import numpy as np
from numpy.typing import NDArray

from noteflow.infrastructure.logging import get_logger

from ...proto import noteflow_pb2
from .._audio_helpers import convert_audio_format
from .._types import GrpcContext
from ..errors import abort_failed_precondition, abort_invalid_argument
from ._asr import process_audio_segment
from ._cleanup import cleanup_stream_resources
from ._partials import clear_partial_buffer, maybe_emit_partial
from ._processing import (
    flush_segmenter,
    normalize_stream_format,
    process_audio_with_vad,
    process_stream_chunk,
    write_audio_chunk_safe,
)
from ._session import StreamSessionManager
from ._types import StreamSessionInit

if TYPE_CHECKING:
    from ..protocols import ServicerHost

logger = get_logger(__name__)


class StreamingMixin:
    """Mixin providing streaming transcription functionality.

    Requires host to implement ServicerHost protocol.
    """

    async def StreamTranscription(
        self: ServicerHost,
        request_iterator: AsyncIterator[noteflow_pb2.AudioChunk],
        context: GrpcContext,
    ) -> AsyncIterator[noteflow_pb2.TranscriptUpdate]:
        """Handle bidirectional audio streaming with persistence.

        Receive audio chunks from client, process through ASR,
        persist segments, and yield transcript updates.
        Works with both database and memory backends via RepositoryProvider.
        """
        if self.asr_engine is None or not self.asr_engine.is_loaded:
            await abort_failed_precondition(context, "ASR engine not loaded")

        current_meeting_id: str | None = None
        # Track separately to guarantee cleanup even if exception occurs
        # between _init_stream_for_meeting returning and current_meeting_id assignment
        initialized_meeting_id: str | None = None

        try:
            async for chunk in request_iterator:
                meeting_id = chunk.meeting_id
                if not meeting_id:
                    await abort_invalid_argument(context, "meeting_id required")

                # Initialize stream on first chunk
                if current_meeting_id is None:
                    # Track meeting_id BEFORE init to guarantee cleanup on any exception
                    # (cleanup_stream_resources is idempotent, safe to call even if init aborts)
                    initialized_meeting_id = meeting_id
                    init_result = await self.init_stream_for_meeting(meeting_id, context)
                    if init_result is None:
                        return  # Error already sent via context.abort
                    current_meeting_id = meeting_id
                elif meeting_id != current_meeting_id:
                    await abort_invalid_argument(
                        context, "Stream may only contain a single meeting_id"
                    )

                # Check for stop request (graceful shutdown from StopMeeting)
                if current_meeting_id in self.stop_requested:
                    logger.info(
                        "Stop requested for meeting %s, exiting stream gracefully",
                        current_meeting_id,
                    )
                    break

                # Process audio chunk
                async for update in self.process_stream_chunk(
                    current_meeting_id, chunk, context
                ):
                    yield update

            # Flush any remaining audio from segmenter
            if current_meeting_id and current_meeting_id in self.segmenters:
                async for update in self.flush_segmenter(current_meeting_id):
                    yield update
        finally:
            if cleanup_meeting := current_meeting_id or initialized_meeting_id:
                cleanup_stream_resources(self, cleanup_meeting)

    async def init_stream_for_meeting(
        self: ServicerHost,
        meeting_id: str,
        context: GrpcContext,
    ) -> StreamSessionInit | None:
        """Initialize streaming for a meeting.

        Uses unified repository provider for both DB and memory backends.
        Protected by lock to prevent race conditions on concurrent init attempts.

        Args:
            meeting_id: Meeting ID string.
            context: gRPC context for error handling.

        Returns:
            Initialization result, or None if error was sent.
        """
        return await StreamSessionManager.init_stream_for_meeting(self, meeting_id, context)

    def process_stream_chunk(
        self: ServicerHost,
        meeting_id: str,
        chunk: noteflow_pb2.AudioChunk,
        context: GrpcContext,
    ) -> AsyncIterator[noteflow_pb2.TranscriptUpdate]:
        """Process a single audio chunk from the stream.

        Args:
            meeting_id: Meeting ID string.
            chunk: Audio chunk from client.
            context: gRPC context for error handling.

        Returns:
            Async iterator of transcript updates from processing.
        """
        return process_stream_chunk(self, meeting_id, chunk, context)

    def _normalize_stream_format(
        self: ServicerHost,
        meeting_id: str,
        sample_rate: int,
        channels: int,
    ) -> tuple[int, int]:
        """Validate and persist stream audio format for a meeting."""
        return normalize_stream_format(self, meeting_id, sample_rate, channels)

    def _convert_audio_format(
        self: ServicerHost,
        audio: NDArray[np.float32],
        sample_rate: int,
        channels: int,
    ) -> NDArray[np.float32]:
        """Downmix/resample audio to the server's expected format."""
        return convert_audio_format(audio, sample_rate, channels, self.DEFAULT_SAMPLE_RATE)

    def _write_audio_chunk_safe(
        self: ServicerHost,
        meeting_id: str,
        audio: NDArray[np.float32],
    ) -> None:
        """Write audio chunk to encrypted file, logging errors without raising."""
        write_audio_chunk_safe(self, meeting_id, audio)

    async def _process_audio_with_vad(
        self: ServicerHost,
        meeting_id: str,
        audio: NDArray[np.float32],
    ) -> AsyncIterator[noteflow_pb2.TranscriptUpdate]:
        """Process audio through VAD and Segmenter, yielding transcript updates."""
        async for update in process_audio_with_vad(self, meeting_id, audio):
            yield update

    async def _maybe_emit_partial(
        self: ServicerHost,
        meeting_id: str,
    ) -> noteflow_pb2.TranscriptUpdate | None:
        """Check if it's time to emit a partial and generate if so."""
        return await maybe_emit_partial(self, meeting_id)

    def _clear_partial_buffer(self: ServicerHost, meeting_id: str) -> None:
        """Clear the partial buffer and reset state after a final is emitted."""
        clear_partial_buffer(self, meeting_id)

    def flush_segmenter(
        self: ServicerHost,
        meeting_id: str,
    ) -> AsyncIterator[noteflow_pb2.TranscriptUpdate]:
        """Flush remaining audio from segmenter at stream end."""
        return flush_segmenter(self, meeting_id)

    async def _process_audio_segment(
        self: ServicerHost,
        meeting_id: str,
        audio: NDArray[np.float32],
        segment_start_time: float,
    ) -> AsyncIterator[noteflow_pb2.TranscriptUpdate]:
        """Process a complete audio segment through ASR."""
        async for update in process_audio_segment(
            self, meeting_id, audio, segment_start_time
        ):
            yield update
