"""Main DiarizationMixin class with RPC entry points."""

from __future__ import annotations

import asyncio
from typing import TYPE_CHECKING

from noteflow.infrastructure.logging import get_logger

from ...proto import noteflow_pb2
from .._types import GrpcStatusContext
from ._jobs import JobsMixin
from ._refinement import RefinementMixin
from ._speaker import SpeakerMixin
from ._streaming import StreamingDiarizationMixin

if TYPE_CHECKING:
    from ..protocols import ServicerHost

logger = get_logger(__name__)


class DiarizationMixin(
    StreamingDiarizationMixin,
    RefinementMixin,
    JobsMixin,
    SpeakerMixin,
):
    """Mixin providing speaker diarization functionality.

    Requires host to implement ServicerHost protocol.

    This mixin composes functionality from:
    - StreamingDiarizationMixin: Real-time audio processing
    - RefinementMixin: Offline diarization inference
    - JobsMixin: Background job management
    - SpeakerMixin: Speaker assignment and renaming
    """

    async def RefineSpeakerDiarization(
        self: ServicerHost,
        request: noteflow_pb2.RefineSpeakerDiarizationRequest,
        context: GrpcStatusContext,
    ) -> noteflow_pb2.RefineSpeakerDiarizationResponse:
        """Run post-meeting speaker diarization refinement.

        Load the full meeting audio, run offline diarization, and update
        segment speaker assignments. Job state is persisted when DB available.
        """
        await self.prune_diarization_jobs()
        return await self.start_diarization_job(request, context)

    async def refine_speaker_diarization(
        self: ServicerHost,
        meeting_id: str,
        num_speakers: int | None = None,
    ) -> int:
        """Run post-meeting speaker diarization refinement.

        Load the full meeting audio, run offline diarization, and update
        segment speaker assignments. This provides higher quality speaker
        labels than streaming diarization.

        Args:
            meeting_id: Meeting UUID string.
            num_speakers: Known number of speakers (None for auto-detect).

        Returns:
            Number of segments updated with speaker labels.

        Raises:
            RuntimeError: If diarization engine not available or meeting not found.
        """
        async with self.diarization_lock:
            turns = await asyncio.to_thread(
                self.run_diarization_inference,
                meeting_id,
                num_speakers,
            )

        updated_count = await self.apply_diarization_turns(meeting_id, turns)

        logger.info(
            "Updated %d segments with speaker labels for meeting %s",
            updated_count,
            meeting_id,
        )

        return updated_count
