"""Offline diarization refinement logic."""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

from noteflow.infrastructure.audio.reader import MeetingAudioReader
from noteflow.infrastructure.diarization import SpeakerTurn
from noteflow.infrastructure.logging import get_logger

from ..converters import parse_meeting_id_or_none
from ._speaker import apply_speaker_to_segment

if TYPE_CHECKING:
    from ..protocols import ServicerHost

logger = get_logger(__name__)


class RefinementMixin:
    """Mixin providing offline diarization refinement functionality."""

    def run_diarization_inference(
        self: ServicerHost,
        meeting_id: str,
        num_speakers: int | None,
    ) -> list[SpeakerTurn]:
        """Run offline diarization and return speaker turns (blocking)."""
        if self.diarization_engine is None:
            raise RuntimeError("Diarization engine not configured")

        if not self.diarization_engine.is_offline_loaded:
            logger.info("Loading offline diarization model for refinement...")
            self.diarization_engine.load_offline_model()

        audio_reader = MeetingAudioReader(self.crypto, self.meetings_dir)
        if not audio_reader.audio_exists(meeting_id):
            raise RuntimeError("No audio file found for meeting")

        logger.info("Loading audio for meeting %s", meeting_id)
        try:
            audio_chunks = audio_reader.load_meeting_audio(meeting_id)
        except (FileNotFoundError, ValueError) as exc:
            raise RuntimeError(f"Failed to load audio: {exc}") from exc

        if not audio_chunks:
            raise RuntimeError("No audio chunks loaded for meeting")

        sample_rate = audio_reader.sample_rate
        all_audio = np.concatenate([chunk.frames for chunk in audio_chunks])

        logger.info(
            "Running offline diarization on %.2f seconds of audio",
            len(all_audio) / sample_rate,
        )

        turns = self.diarization_engine.diarize_full(
            all_audio,
            sample_rate=sample_rate,
            num_speakers=num_speakers,
        )

        logger.info("Diarization found %d speaker turns", len(turns))
        return list(turns)

    async def apply_diarization_turns(
        self: ServicerHost,
        meeting_id: str,
        turns: list[SpeakerTurn],
    ) -> int:
        """Apply diarization turns to segments and return updated count."""
        updated_count = 0
        parsed_meeting_id = parse_meeting_id_or_none(meeting_id)
        if parsed_meeting_id is None:
            logger.warning("Invalid meeting_id %s while applying diarization turns", meeting_id)
            return 0

        async with self.create_repository_provider() as repo:
            segments = await repo.segments.get_by_meeting(parsed_meeting_id)
            for segment in segments:
                if apply_speaker_to_segment(segment, turns):
                    # For DB segments with db_id, use update_speaker
                    if segment.db_id is not None:
                        await repo.segments.update_speaker(
                            segment.db_id,
                            segment.speaker_id,
                            segment.speaker_confidence,
                        )
                    updated_count += 1
            await repo.commit()

        return updated_count
