"""Streaming diarization processing."""

from __future__ import annotations

import asyncio
from functools import partial
from typing import TYPE_CHECKING

import numpy as np
from numpy.typing import NDArray

from noteflow.infrastructure.diarization import SpeakerTurn
from noteflow.infrastructure.logging import get_logger
from noteflow.infrastructure.persistence.repositories import StreamingTurn

if TYPE_CHECKING:
    from ..protocols import ServicerHost

logger = get_logger(__name__)


class StreamingDiarizationMixin:
    """Mixin providing streaming diarization processing."""

    async def process_streaming_diarization(
        self: ServicerHost,
        meeting_id: str,
        audio: NDArray[np.float32],
    ) -> None:
        """Process an audio chunk for streaming diarization (best-effort).

        Uses per-meeting sessions to enable concurrent diarization without
        race conditions. Each meeting has its own pipeline state.

        Offloads heavy ML inference to thread pool to avoid blocking the event loop.
        """
        if self.diarization_engine is None:
            return
        if meeting_id in self.diarization_streaming_failed:
            return
        if audio.size == 0:
            return

        loop = asyncio.get_running_loop()

        # Get or create per-meeting session under lock
        async with self.diarization_lock:
            session = self.diarization_sessions.get(meeting_id)
            if session is None:
                try:
                    session = await loop.run_in_executor(
                        None,
                        self.diarization_engine.create_streaming_session,
                        meeting_id,
                    )
                    prior_turns = self.diarization_turns.get(meeting_id, [])
                    prior_stream_time = self.diarization_stream_time.get(meeting_id, 0.0)
                    if prior_turns or prior_stream_time:
                        session.restore(prior_turns, stream_time=prior_stream_time)
                    self.diarization_sessions[meeting_id] = session
                except (RuntimeError, ValueError) as exc:
                    logger.warning(
                        "Streaming diarization disabled for meeting %s: %s",
                        meeting_id,
                        exc,
                    )
                    self.diarization_streaming_failed.add(meeting_id)
                    return

        # Process chunk in thread pool (outside lock for parallelism)
        try:
            new_turns = await loop.run_in_executor(
                None,
                partial(
                    session.process_chunk,
                    audio,
                    sample_rate=self.DEFAULT_SAMPLE_RATE,
                ),
            )
        except (RuntimeError, OSError) as exc:
            logger.warning(
                "Streaming diarization failed for meeting %s: %s",
                meeting_id,
                exc,
            )
            self.diarization_streaming_failed.add(meeting_id)
            return

        # Populate diarization_turns for compatibility with maybe_assign_speaker
        if new_turns:
            diarization_turns = self.diarization_turns.setdefault(meeting_id, [])
            diarization_turns.extend(new_turns)

            # Update stream time for legacy compatibility
            self.diarization_stream_time[meeting_id] = session.stream_time

            # Persist turns immediately for crash resilience (DB only)
            await self.persist_streaming_turns(meeting_id, list(new_turns))

    async def persist_streaming_turns(
        self: ServicerHost,
        meeting_id: str,
        new_turns: list[SpeakerTurn],
    ) -> None:
        """Persist streaming turns to database (fire-and-forget)."""
        try:
            async with self.create_repository_provider() as repo:
                if repo.supports_diarization_jobs:
                    repo_turns = [
                        StreamingTurn(
                            speaker=t.speaker,
                            start_time=t.start,
                            end_time=t.end,
                            confidence=t.confidence,
                        )
                        for t in new_turns
                    ]
                    await repo.diarization_jobs.add_streaming_turns(meeting_id, repo_turns)
                    await repo.commit()
        # INTENTIONAL BROAD HANDLER: Fire-and-forget persistence
        # - Turn persistence should not block streaming
        # - Data can be recovered from audio file if needed
        except Exception:
            logger.exception("Failed to persist streaming turns for %s", meeting_id)
