"""Audio chunk processing and VAD handling for streaming transcription."""

from __future__ import annotations

import time
from collections import deque
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, decode_audio_chunk, validate_stream_format
from .._types import GrpcContext
from ..converters import create_ack_update, create_congestion_info, create_vad_update
from ..errors import abort_invalid_argument
from ._asr import process_audio_segment
from ._partials import clear_partial_buffer, maybe_emit_partial

if TYPE_CHECKING:
    from ..protocols import ServicerHost

logger = get_logger(__name__)

# Congestion thresholds (public for testability)
PROCESSING_DELAY_THRESHOLD_MS = 1000  # 1 second delay triggers throttle
QUEUE_DEPTH_THRESHOLD = 20  # 20 pending chunks triggers throttle


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

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

    Yields:
        Transcript updates from processing.
    """
    # Track chunk sequence for acknowledgment (default 0 for backwards compat)
    chunk_sequence = max(chunk.chunk_sequence, 0)
    ack_update = track_chunk_sequence(host, meeting_id, chunk_sequence)
    if ack_update is not None:
        yield ack_update

    try:
        sample_rate, channels = normalize_stream_format(
            host,
            meeting_id,
            chunk.sample_rate,
            chunk.channels,
        )
    except ValueError as e:
        await abort_invalid_argument(context, str(e))
        raise  # Unreachable but helps type checker

    audio = decode_audio_chunk(chunk.audio_data)
    if audio is None:
        return

    try:
        audio = _convert_audio_format(host, audio, sample_rate, channels)
    except ValueError as e:
        await abort_invalid_argument(context, str(e))

    # Write to encrypted audio file
    write_audio_chunk_safe(host, meeting_id, audio)

    # VAD-driven segmentation
    async for update in process_audio_with_vad(host, meeting_id, audio):
        yield update


def normalize_stream_format(
    host: ServicerHost,
    meeting_id: str,
    sample_rate: int,
    channels: int,
) -> tuple[int, int]:
    """Validate and persist stream audio format for a meeting."""
    existing = host.stream_formats.get(meeting_id)
    result = validate_stream_format(
        sample_rate,
        channels,
        host.DEFAULT_SAMPLE_RATE,
        frozenset(host.SUPPORTED_SAMPLE_RATES),
        existing,
    )
    host.stream_formats.setdefault(meeting_id, result)
    return result


# Emit ack every N chunks (~500ms at 100ms per chunk)
ACK_CHUNK_INTERVAL = 5

# Maximum receipt timestamps to track for processing delay calculation
_RECEIPT_TIMES_WINDOW = 10


def track_chunk_sequence(
    host: ServicerHost,
    meeting_id: str,
    chunk_sequence: int,
) -> noteflow_pb2.TranscriptUpdate | None:
    """Track chunk sequence and emit ack with congestion info every N chunks.

    Args:
        host: The servicer host.
        meeting_id: Meeting identifier.
        chunk_sequence: Sequence number from the chunk (0 if not provided).

    Returns:
        TranscriptUpdate with ack_sequence and congestion info if ack interval reached.
    """
    receipt_time = time.monotonic()

    # Initialize receipt times tracking if needed
    if not hasattr(host, "_chunk_receipt_times"):
        host.chunk_receipt_times = {}
    if meeting_id not in host.chunk_receipt_times:
        host.chunk_receipt_times[meeting_id] = deque(maxlen=_RECEIPT_TIMES_WINDOW)

    # Track receipt timestamp for processing delay calculation
    host.chunk_receipt_times[meeting_id].append(receipt_time)

    # Initialize pending chunks counter if needed
    if not hasattr(host, "_pending_chunks"):
        host.pending_chunks = {}
    host.pending_chunks[meeting_id] = host.pending_chunks.get(meeting_id, 0) + 1

    # Track highest received sequence (only if client provides sequences)
    if chunk_sequence > 0:
        prev_seq = host.chunk_sequences.get(meeting_id, 0)
        if chunk_sequence > prev_seq + 1:
            # Gap detected - log for debugging (client may retry)
            logger.warning(
                "Chunk sequence gap for meeting %s: expected %d, got %d",
                meeting_id,
                prev_seq + 1,
                chunk_sequence,
            )
        host.chunk_sequences[meeting_id] = max(prev_seq, chunk_sequence)

    # Increment chunk count and check if we should emit ack
    count = host.chunk_counts.get(meeting_id, 0) + 1
    host.chunk_counts[meeting_id] = count

    if count >= ACK_CHUNK_INTERVAL:
        host.chunk_counts[meeting_id] = 0
        ack_seq = host.chunk_sequences.get(meeting_id, 0)
        # Only emit ack if client is sending sequences
        if ack_seq > 0:
            # Calculate congestion info
            congestion = calculate_congestion_info(host, meeting_id, receipt_time)
            return create_ack_update(meeting_id, ack_seq, congestion)

    return None


def calculate_congestion_info(
    host: ServicerHost,
    meeting_id: str,
    current_time: float,
) -> noteflow_pb2.CongestionInfo:
    """Calculate congestion info for backpressure signaling.

    Args:
        host: The servicer host.
        meeting_id: Meeting identifier.
        current_time: Current monotonic timestamp.

    Returns:
        CongestionInfo with processing delay, queue depth, and throttle recommendation.
    """
    if receipt_times := host.chunk_receipt_times.get(meeting_id, deque()):
        oldest_receipt = receipt_times[0]
        processing_delay_ms = int((current_time - oldest_receipt) * 1000)
    else:
        processing_delay_ms = 0

    # Get queue depth (pending chunks not yet processed through ASR)
    queue_depth = host.pending_chunks.get(meeting_id, 0)

    # Determine if throttle is recommended
    throttle_recommended = (
        processing_delay_ms > PROCESSING_DELAY_THRESHOLD_MS
        or queue_depth > QUEUE_DEPTH_THRESHOLD
    )

    return create_congestion_info(
        processing_delay_ms=processing_delay_ms,
        queue_depth=queue_depth,
        throttle_recommended=throttle_recommended,
    )


def decrement_pending_chunks(host: ServicerHost, meeting_id: str) -> None:
    """Decrement pending chunks counter after processing.

    Call this after ASR processing completes for a segment.
    """
    if hasattr(host, "_pending_chunks") and meeting_id in host.pending_chunks:
        # Decrement by ACK_CHUNK_INTERVAL since we process in batches
        host.pending_chunks[meeting_id] = max(
            0, host.pending_chunks[meeting_id] - ACK_CHUNK_INTERVAL
        )
        if receipt_times := host.chunk_receipt_times.get(meeting_id):
            # Remove timestamps corresponding to processed chunks
            for _ in range(min(ACK_CHUNK_INTERVAL, len(receipt_times))):
                if receipt_times:
                    receipt_times.popleft()


def _convert_audio_format(
    host: 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, host.DEFAULT_SAMPLE_RATE)


def write_audio_chunk_safe(
    host: ServicerHost,
    meeting_id: str,
    audio: NDArray[np.float32],
) -> None:
    """Write audio chunk to encrypted file, logging errors without raising."""
    if meeting_id not in host.audio_writers:
        return
    if meeting_id in host.audio_write_failed:
        return  # Already failed, skip to avoid log spam
    try:
        host.audio_writers[meeting_id].write_chunk(audio)
    except (OSError, ValueError) as e:
        logger.error(
            "Audio write failed for meeting %s: %s. Recording may be incomplete.",
            meeting_id,
            e,
        )
        host.audio_write_failed.add(meeting_id)


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

    Uses consolidated MeetingStreamState for O(1) lookup instead of 13+ dict accesses.
    """
    # Single dict lookup replaces 6+ separate lookups per audio chunk
    state = host.get_stream_state(meeting_id)
    if state is None:
        return

    # Get VAD decision using consolidated state
    is_speech = state.vad.process_chunk(audio)

    # Streaming diarization (optional)
    await host.process_streaming_diarization(meeting_id, audio)

    # Emit VAD state change events using consolidated state
    if is_speech and not state.was_speaking:
        # Speech started
        yield create_vad_update(meeting_id, noteflow_pb2.UPDATE_TYPE_VAD_START)
        state.was_speaking = True
        host.was_speaking[meeting_id] = True  # Keep legacy dict in sync
    elif not is_speech and state.was_speaking:
        # Speech ended
        yield create_vad_update(meeting_id, noteflow_pb2.UPDATE_TYPE_VAD_END)
        state.was_speaking = False
        host.was_speaking[meeting_id] = False  # Keep legacy dict in sync

    # Buffer audio for partial transcription (pre-allocated buffer handles copy)
    if is_speech:
        state.partial_buffer.append(audio)

        # Check if we should emit a partial
        partial_update = await maybe_emit_partial(host, meeting_id)
        if partial_update is not None:
            yield partial_update

    # Process through segmenter using consolidated state
    for audio_segment in state.segmenter.process_audio(audio, is_speech):
        # Clear partial buffer when we get a final segment
        clear_partial_buffer(host, meeting_id)
        async for update in process_audio_segment(
            host,
            meeting_id,
            audio_segment.audio,
            audio_segment.start_time,
        ):
            yield update


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

    Args:
        host: The servicer host.
        meeting_id: Meeting identifier.

    Yields:
        TranscriptUpdates for final segment.
    """
    segmenter = host.segmenters.get(meeting_id)
    if segmenter is None:
        return

    # Clear partial buffer since we're flushing to final
    clear_partial_buffer(host, meeting_id)

    final_segment = segmenter.flush()
    if final_segment is not None:
        async for update in process_audio_segment(
            host,
            meeting_id,
            final_segment.audio,
            final_segment.start_time,
        ):
            yield update
