"""Diarization engine implementation using pyannote.audio and diart.

Provides speaker diarization for both streaming (real-time) and
offline (post-meeting) processing.

Requires optional dependencies: pip install noteflow[diarization]
"""

from __future__ import annotations

import os
import warnings
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Protocol, Self, cast

from noteflow.config.constants import DEFAULT_SAMPLE_RATE, ERR_HF_TOKEN_REQUIRED
from noteflow.infrastructure.diarization.dto import SpeakerTurn
from noteflow.infrastructure.diarization.session import DiarizationSession
from noteflow.infrastructure.logging import get_logger, log_timing

if TYPE_CHECKING:
    import numpy as np
    from diart import SpeakerDiarization
    from diart.models import EmbeddingModel, SegmentationModel
    from numpy.typing import NDArray
    from pyannote.core import Annotation
    from torch import Tensor
    from torch import device as TorchDevice


class _PipelinePretrainedModel(Protocol):
    @classmethod
    def from_pretrained(
        cls, model: str, *, use_auth_token: str | None = ...
    ) -> Self: ...


class _OfflinePipeline(Protocol):
    def to(self, device: TorchDevice) -> None: ...

    def __call__(
        self,
        waveform: Mapping[str, Tensor | int],
        *,
        num_speakers: int | None = ...,
        min_speakers: int | None = ...,
        max_speakers: int | None = ...,
    ) -> Annotation: ...


class _TorchModule(Protocol):
    def from_numpy(self, ndarray: NDArray[np.float32]) -> Tensor: ...

logger = get_logger(__name__)


class DiarizationEngine:
    """Speaker diarization engine using pyannote.audio and diart.

    Supports both streaming (real-time via diart) and offline
    (post-meeting via pyannote.audio) diarization modes.
    """

    def __init__(
        self,
        device: str = "auto",
        hf_token: str | None = None,
        streaming_latency: float = 0.5,
        min_speakers: int = 1,
        max_speakers: int = 10,
    ) -> None:
        """Initialize the diarization engine.

        Args:
            device: Device to use ("auto", "cpu", "cuda", "mps").
                    "auto" selects CUDA > MPS > CPU based on availability.
            hf_token: HuggingFace token for pyannote model access.
            streaming_latency: Latency for streaming diarization in seconds.
            min_speakers: Minimum expected speakers for offline diarization.
            max_speakers: Maximum expected speakers for offline diarization.
        """
        self._device_preference = device
        self._device: str | None = None
        self._hf_token = hf_token
        self._streaming_latency = streaming_latency

        # Set HF_TOKEN for huggingface_hub library to use for authenticated requests
        if hf_token:
            os.environ["HF_TOKEN"] = hf_token
        self._min_speakers = min_speakers
        self._max_speakers = max_speakers

        # Lazy-loaded models
        self._streaming_pipeline: SpeakerDiarization | None = None
        self._offline_pipeline: _OfflinePipeline | None = None

        # Shared models for per-session pipelines (loaded once, reused)
        self._segmentation_model: SegmentationModel | None = None
        self._embedding_model: EmbeddingModel | None = None

    def _resolve_device(self) -> str:
        """Resolve the actual device to use based on availability.

        Returns:
            Device string ("cuda", "mps", or "cpu").
        """
        if self._device is not None:
            return self._device

        import torch

        if self._device_preference == "auto":
            if torch.cuda.is_available():
                self._device = "cuda"
            elif torch.backends.mps.is_available():
                self._device = "mps"
            else:
                self._device = "cpu"
        else:
            self._device = self._device_preference

        logger.info("Diarization device resolved to: %s", self._device)
        return self._device

    def load_streaming_model(self) -> None:
        """Load the streaming diarization model (diart).

        Raises:
            RuntimeError: If model loading fails.
            ValueError: If HuggingFace token is not provided.
        """
        if self._streaming_pipeline is not None:
            logger.debug("Streaming model already loaded")
            return

        if not self._hf_token:
            raise ValueError(ERR_HF_TOKEN_REQUIRED)

        device = self._resolve_device()

        logger.info(
            "Loading streaming diarization model on %s with latency %.2fs...",
            device,
            self._streaming_latency,
        )

        try:
            import torch
            from diart import SpeakerDiarization, SpeakerDiarizationConfig
            from diart.models import EmbeddingModel, SegmentationModel

            segmentation = SegmentationModel.from_pretrained(
                "pyannote/segmentation-3.0",
                use_hf_token=self._hf_token,
            )
            embedding = EmbeddingModel.from_pretrained(
                "pyannote/wespeaker-voxceleb-resnet34-LM",
                use_hf_token=self._hf_token,
            )

            torch_device = torch.device(device)
            config = SpeakerDiarizationConfig(
                segmentation=segmentation,
                embedding=embedding,
                step=self._streaming_latency,
                latency=self._streaming_latency,
                device=torch_device,
            )

            self._streaming_pipeline = SpeakerDiarization(config)
            logger.info("Streaming diarization model loaded successfully")

        except (RuntimeError, OSError) as e:
            raise RuntimeError(f"Failed to load streaming diarization model: {e}") from e

    def _ensure_streaming_models_loaded(self) -> None:
        """Ensure shared streaming models are loaded.

        Loads the segmentation and embedding models that are shared across
        all streaming sessions. These models are stateless and can be safely
        reused by multiple concurrent sessions.

        Raises:
            RuntimeError: If model loading fails.
            ValueError: If HuggingFace token is not provided.
        """
        if self._segmentation_model is not None and self._embedding_model is not None:
            return

        if not self._hf_token:
            raise ValueError(ERR_HF_TOKEN_REQUIRED)

        device = self._resolve_device()
        logger.info("Loading shared streaming diarization models on %s...", device)

        try:
            from diart.models import EmbeddingModel, SegmentationModel

            self._segmentation_model = SegmentationModel.from_pretrained(
                "pyannote/segmentation-3.0",
                use_hf_token=self._hf_token,
            )
            self._embedding_model = EmbeddingModel.from_pretrained(
                "pyannote/wespeaker-voxceleb-resnet34-LM",
                use_hf_token=self._hf_token,
            )
            logger.info("Shared streaming models loaded successfully")

        except (RuntimeError, OSError) as e:
            raise RuntimeError(f"Failed to load streaming models: {e}") from e

    def create_streaming_session(self, meeting_id: str) -> DiarizationSession:
        """Create a new per-meeting streaming diarization session.

        Each session maintains its own pipeline state, enabling concurrent
        diarization of multiple meetings without interference. The underlying
        models are shared across sessions for memory efficiency.

        Args:
            meeting_id: Unique identifier for the meeting.

        Returns:
            New DiarizationSession for the meeting.

        Raises:
            RuntimeError: If model loading fails.
            ValueError: If HuggingFace token is not provided.
        """
        self._ensure_streaming_models_loaded()

        import torch
        from diart import SpeakerDiarization, SpeakerDiarizationConfig

        config = SpeakerDiarizationConfig(
            segmentation=self._segmentation_model,
            embedding=self._embedding_model,
            step=self._streaming_latency,
            latency=self._streaming_latency,
            device=torch.device(self._resolve_device()),
        )

        pipeline = SpeakerDiarization(config)
        logger.info("Created streaming session for meeting %s", meeting_id)

        return DiarizationSession(
            meeting_id=meeting_id,
            _pipeline=pipeline,
            _sample_rate=DEFAULT_SAMPLE_RATE,
        )

    def load_offline_model(self) -> None:
        """Load the offline diarization model (pyannote.audio).

        Raises:
            RuntimeError: If model loading fails.
            ValueError: If HuggingFace token is not provided.
        """
        if self._offline_pipeline is not None:
            logger.debug("Offline model already loaded")
            return

        if not self._hf_token:
            raise ValueError(ERR_HF_TOKEN_REQUIRED)

        device = self._resolve_device()

        with log_timing("diarization_offline_model_load", device=device):
            try:
                import torch
                from pyannote.audio import Pipeline

                pipeline_class = cast(type[_PipelinePretrainedModel], Pipeline)
                pipeline = cast(
                    _OfflinePipeline,
                    pipeline_class.from_pretrained(
                        "pyannote/speaker-diarization-3.1",
                        use_auth_token=self._hf_token,
                    ),
                )

                torch_device = torch.device(device)
                pipeline.to(torch_device)
                self._offline_pipeline = pipeline

            except (RuntimeError, OSError) as e:
                raise RuntimeError(f"Failed to load offline diarization model: {e}") from e

        logger.info("diarization_offline_model_loaded", device=device)

    def process_chunk(
        self,
        audio: NDArray[np.float32],
        sample_rate: int = DEFAULT_SAMPLE_RATE,
    ) -> Sequence[SpeakerTurn]:
        """Process an audio chunk for streaming diarization.

        Args:
            audio: Audio samples as float32 array (mono).
            sample_rate: Audio sample rate in Hz.

        Returns:
            Sequence of speaker turns detected in this chunk.

        Raises:
            RuntimeError: If streaming model not loaded.
        """
        if self._streaming_pipeline is None:
            raise RuntimeError("Streaming model not loaded. Call load_streaming_model() first.")

        from pyannote.core import SlidingWindowFeature

        # Reshape audio for diart: (samples,) -> (1, samples)
        if audio.ndim == 1:
            audio = audio.reshape(1, -1)

        # Create SlidingWindowFeature for diart
        from pyannote.core import SlidingWindow

        duration = audio.shape[1] / sample_rate
        window = SlidingWindow(start=0.0, duration=duration, step=duration)
        waveform = SlidingWindowFeature(audio, window)

        # Process through pipeline
        results = self._streaming_pipeline([waveform])

        turns: list[SpeakerTurn] = []
        for annotation, _ in results:
            turns.extend(self._annotation_to_turns(annotation))

        return turns

    def diarize_full(
        self,
        audio: NDArray[np.float32],
        sample_rate: int = DEFAULT_SAMPLE_RATE,
        num_speakers: int | None = None,
    ) -> Sequence[SpeakerTurn]:
        """Diarize a complete audio recording.

        Args:
            audio: Audio samples as float32 array (mono).
            sample_rate: Audio sample rate in Hz.
            num_speakers: Known number of speakers (None for auto-detect).

        Returns:
            Sequence of speaker turns for the full recording.

        Raises:
            RuntimeError: If offline model not loaded.
        """
        if self._offline_pipeline is None:
            raise RuntimeError("Offline model not loaded. Call load_offline_model() first.")

        import torch

        # Prepare audio tensor: (samples,) -> (channels, samples)
        torch_typed = cast(_TorchModule, torch)
        if audio.ndim == 1:
            audio_tensor = torch_typed.from_numpy(audio).unsqueeze(0)
        else:
            audio_tensor = torch_typed.from_numpy(audio)

        audio_duration_seconds = audio_tensor.shape[1] / sample_rate

        # Create waveform dict for pyannote
        waveform: dict[str, Tensor | int] = {"waveform": audio_tensor, "sample_rate": sample_rate}

        with log_timing(
            "diarization_full_audio",
            audio_duration_seconds=round(audio_duration_seconds, 2),
            num_speakers=num_speakers,
        ):
            # Run diarization with speaker hints
            if num_speakers is not None:
                annotation = self._offline_pipeline(waveform, num_speakers=num_speakers)
            else:
                annotation = self._offline_pipeline(
                    waveform,
                    min_speakers=self._min_speakers,
                    max_speakers=self._max_speakers,
                )

        turns = self._annotation_to_turns(annotation)
        logger.info(
            "diarization_full_completed",
            audio_duration_seconds=round(audio_duration_seconds, 2),
            speaker_turn_count=len(turns),
        )
        return turns

    def _annotation_to_turns(self, annotation: Annotation) -> list[SpeakerTurn]:
        """Convert pyannote Annotation to SpeakerTurn list.

        Args:
            annotation: Pyannote diarization annotation.

        Returns:
            List of SpeakerTurn objects.
        """
        turns: list[SpeakerTurn] = []

        # itertracks(yield_label=True) returns 3-tuples: (segment, track, label)
        for track in annotation.itertracks(yield_label=True):
            # Unpack with len check for type safety with pyannote's union return
            if len(track) == 3:
                segment, _, speaker = track
                turns.append(
                    SpeakerTurn(
                        speaker=str(speaker),
                        start=segment.start,
                        end=segment.end,
                    )
                )

        return turns

    def reset_streaming(self) -> None:
        """Reset streaming pipeline state for a new recording.

        .. deprecated::
            Use create_streaming_session() for per-meeting sessions instead.
            This method resets global state and will cause race conditions
            with concurrent meetings.
        """
        warnings.warn(
            "reset_streaming() is deprecated. Use create_streaming_session() for "
            "per-meeting sessions to avoid race conditions with concurrent meetings.",
            DeprecationWarning,
            stacklevel=2,
        )
        if self._streaming_pipeline is not None:
            self._streaming_pipeline.reset()
            logger.debug("Streaming pipeline state reset")

    def unload(self) -> None:
        """Unload all models to free memory."""
        self._streaming_pipeline = None
        self._offline_pipeline = None
        self._segmentation_model = None
        self._embedding_model = None
        self._device = None
        logger.info("Diarization models unloaded")

    @property
    def is_streaming_loaded(self) -> bool:
        """Return True if streaming model is loaded."""
        return self._streaming_pipeline is not None

    @property
    def is_offline_loaded(self) -> bool:
        """Return True if offline model is loaded."""
        return self._offline_pipeline is not None

    @property
    def device(self) -> str | None:
        """Return the resolved device, or None if not yet resolved."""
        return self._device
