"""Partial transcription handling for streaming transcription."""

from __future__ import annotations

import time
from typing import TYPE_CHECKING

from ...proto import noteflow_pb2

if TYPE_CHECKING:
    from ..protocols import ServicerHost


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

    Uses consolidated MeetingStreamState for efficient state access.

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

    Returns:
        TranscriptUpdate with partial text, or None if not time yet.
    """
    if host.asr_engine is None or not host.asr_engine.is_loaded:
        return None

    # Single lookup for all partial-related state
    state = host.get_stream_state(meeting_id)
    if state is None:
        return None

    now = time.time()

    # Sync from legacy dicts if they were modified directly (test compatibility)
    legacy_time = host.last_partial_time.get(meeting_id, 0)
    if legacy_time < state.last_partial_time:
        state.last_partial_time = legacy_time
    legacy_text = host.last_partial_text.get(meeting_id, "")
    if legacy_text != state.last_partial_text:
        state.last_partial_text = legacy_text

    # Check if enough time has passed since last partial
    if now - state.last_partial_time < host.PARTIAL_CADENCE_SECONDS:
        return None

    # Check if we have enough audio
    if state.partial_buffer.is_empty:
        return None

    # Check minimum audio duration before extracting
    if state.partial_buffer.duration_seconds < host.MIN_PARTIAL_AUDIO_SECONDS:
        return None

    # Get buffered audio (single allocation, no concatenation)
    combined = state.partial_buffer.get_audio()

    # Run inference on buffered audio (async to avoid blocking event loop)
    results = await host.asr_engine.transcribe_async(combined)
    partial_text = " ".join(result.text for result in results)

    # Clear buffer after inference to keep partials incremental and bounded.
    # Pre-allocated buffer resets write pointer (O(1), no deallocation).
    state.partial_buffer.clear()

    # Only emit if text changed (debounce)
    if partial_text and partial_text != state.last_partial_text:
        state.last_partial_time = now
        state.last_partial_text = partial_text
        # Keep legacy dicts in sync
        host.last_partial_time[meeting_id] = now
        host.last_partial_text[meeting_id] = partial_text
        return noteflow_pb2.TranscriptUpdate(
            meeting_id=meeting_id,
            update_type=noteflow_pb2.UPDATE_TYPE_PARTIAL,
            partial_text=partial_text,
            server_timestamp=now,
        )

    # Update time even if no text change (cadence tracking)
    state.last_partial_time = now
    host.last_partial_time[meeting_id] = now
    return None


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

    Uses consolidated state for efficient access.

    Args:
        host: The servicer host.
        meeting_id: Meeting identifier.
    """
    current_time = time.time()

    # Use consolidated state if available
    if state := host.get_stream_state(meeting_id):
        state.clear_partial_state(current_time)

    # Keep legacy dicts in sync
    if meeting_id in host.partial_buffers:
        host.partial_buffers[meeting_id].clear()  # O(1) pointer reset
    if meeting_id in host.last_partial_text:
        host.last_partial_text[meeting_id] = ""
    if meeting_id in host.last_partial_time:
        host.last_partial_time[meeting_id] = current_time
