"""Diarization operations mixin for NoteFlowClient."""

from __future__ import annotations

from typing import TYPE_CHECKING

import grpc

from noteflow.grpc._client_mixins.converters import job_status_to_str
from noteflow.grpc._types import DiarizationResult, RenameSpeakerResult
from noteflow.grpc.proto import noteflow_pb2
from noteflow.infrastructure.logging import get_logger

if TYPE_CHECKING:
    from noteflow.grpc._client_mixins.protocols import ClientHost

logger = get_logger(__name__)


class DiarizationClientMixin:
    """Mixin providing speaker diarization operations for NoteFlowClient."""

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

        Requests the server to run offline diarization on the meeting audio
        as a background job and update segment speaker assignments.

        Args:
            meeting_id: Meeting ID.
            num_speakers: Optional known number of speakers (auto-detect if None).

        Returns:
            DiarizationResult with job status or None if request fails.
        """
        if not self.stub:
            return None

        try:
            request = noteflow_pb2.RefineSpeakerDiarizationRequest(
                meeting_id=meeting_id,
                num_speakers=num_speakers or 0,
            )
            response = self.stub.RefineSpeakerDiarization(request)
            speaker_ids = response.speaker_ids
            return DiarizationResult(
                job_id=response.job_id,
                status=job_status_to_str(response.status),
                segments_updated=response.segments_updated,
                speaker_ids=list(speaker_ids),
                error_message=response.error_message,
            )
        except grpc.RpcError as e:
            logger.error("Failed to refine speaker diarization: %s", e)
            return None

    def get_diarization_job_status(
        self: ClientHost,
        job_id: str,
    ) -> DiarizationResult | None:
        """Get status for a diarization background job.

        Args:
            job_id: Job ID from refine_speaker_diarization.

        Returns:
            DiarizationResult with current status or None if request fails.
        """
        if not self.stub:
            return None

        try:
            request = noteflow_pb2.GetDiarizationJobStatusRequest(job_id=job_id)
            response = self.stub.GetDiarizationJobStatus(request)
            speaker_ids = response.speaker_ids
            return DiarizationResult(
                job_id=response.job_id,
                status=job_status_to_str(response.status),
                segments_updated=response.segments_updated,
                speaker_ids=list(speaker_ids),
                error_message=response.error_message,
            )
        except grpc.RpcError as e:
            logger.error("Failed to get diarization job status: %s", e)
            return None

    def rename_speaker(
        self: ClientHost,
        meeting_id: str,
        old_speaker_id: str,
        new_speaker_name: str,
    ) -> RenameSpeakerResult | None:
        """Rename a speaker in all segments of a meeting.

        Args:
            meeting_id: Meeting ID.
            old_speaker_id: Current speaker ID (e.g., "SPEAKER_00").
            new_speaker_name: New speaker name (e.g., "Alice").

        Returns:
            RenameSpeakerResult or None if request fails.
        """
        if not self.stub:
            return None

        try:
            request = noteflow_pb2.RenameSpeakerRequest(
                meeting_id=meeting_id,
                old_speaker_id=old_speaker_id,
                new_speaker_name=new_speaker_name,
            )
            response = self.stub.RenameSpeaker(request)
            return RenameSpeakerResult(
                segments_updated=response.segments_updated,
                success=response.success,
            )
        except grpc.RpcError as e:
            logger.error("Failed to rename speaker: %s", e)
            return None
