"""Per-meeting diarization session for streaming speaker identification.

Each session maintains its own pipeline state, enabling concurrent meetings
without cross-talk. Shared models are loaded once and reused across sessions.
"""

from __future__ import annotations

from collections.abc import Sequence
from dataclasses import dataclass, field
from typing import TYPE_CHECKING

from noteflow.config.constants import DEFAULT_SAMPLE_RATE
from noteflow.infrastructure.diarization.dto import SpeakerTurn
from noteflow.infrastructure.logging import get_logger

if TYPE_CHECKING:
    import numpy as np
    from diart import SpeakerDiarization
    from numpy.typing import NDArray

logger = get_logger(__name__)


@dataclass
class DiarizationSession:
    """Per-meeting streaming diarization session.

    Maintains independent pipeline state for a single meeting, enabling
    concurrent diarization of multiple meetings without interference.

    The session owns its own SpeakerDiarization pipeline instance but
    shares the underlying segmentation and embedding models with other
    sessions for memory efficiency.
    """

    meeting_id: str
    _pipeline: SpeakerDiarization | None
    _sample_rate: int = DEFAULT_SAMPLE_RATE
    _stream_time: float = field(default=0.0, init=False)
    _turns: list[SpeakerTurn] = field(default_factory=list, init=False)
    _closed: bool = field(default=False, init=False)

    def process_chunk(
        self,
        audio: NDArray[np.float32],
        sample_rate: int | None = None,
    ) -> Sequence[SpeakerTurn]:
        """Process an audio chunk and return new speaker turns.

        Args:
            audio: Audio samples as float32 array (mono).
            sample_rate: Audio sample rate (defaults to session's configured rate).

        Returns:
            Sequence of speaker turns detected in this chunk,
            with times adjusted to absolute stream position.

        Raises:
            RuntimeError: If session is closed.
        """
        if self._closed or self._pipeline is None:
            raise RuntimeError(f"Session {self.meeting_id} is closed")

        if audio.size == 0:
            return []

        rate = sample_rate or self._sample_rate
        duration = len(audio) / rate

        # Import here to avoid import errors when diart not installed
        from pyannote.core import SlidingWindow, SlidingWindowFeature

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

        # Create SlidingWindowFeature for diart
        window = SlidingWindow(start=0.0, duration=duration, step=duration)
        waveform = SlidingWindowFeature(audio, window)

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

        # Convert results to turns with absolute time offsets
        new_turns: list[SpeakerTurn] = []
        for annotation, _ in results:
            for track in annotation.itertracks(yield_label=True):
                if len(track) == 3:
                    segment, _, speaker = track
                    turn = SpeakerTurn(
                        speaker=str(speaker),
                        start=segment.start + self._stream_time,
                        end=segment.end + self._stream_time,
                    )
                    new_turns.append(turn)
                    self._turns.append(turn)

        self._stream_time += duration
        return new_turns

    def reset(self) -> None:
        """Reset session state for restarting diarization.

        Clears accumulated turns and resets stream time to zero.
        The underlying pipeline is also reset.
        """
        if self._closed or self._pipeline is None:
            return

        self._pipeline.reset()
        self._stream_time = 0.0
        self._turns.clear()
        logger.debug("Session %s reset", self.meeting_id)

    def restore(
        self,
        turns: Sequence[SpeakerTurn],
        *,
        stream_time: float | None = None,
    ) -> None:
        """Restore session state from prior streaming turns.

        Use during crash recovery to continue diarization timelines.

        Args:
            turns: Previously collected speaker turns.
            stream_time: Optional stream time to restore. If not provided,
                uses the max end time from turns.
        """
        if self._closed:
            return

        self._turns = list(turns)
        if stream_time is None:
            stream_time = max((t.end for t in turns), default=0.0)
        self._stream_time = max(self._stream_time, stream_time)
        logger.debug(
            "Session %s restored stream_time=%.3f with %d turns",
            self.meeting_id,
            self._stream_time,
            len(self._turns),
        )

    def close(self) -> None:
        """Close the session and release resources.

        After closing, the session cannot be used for further processing.
        Explicitly releases the pipeline reference to allow garbage collection
        and GPU memory to be freed.
        """
        if self._closed:
            return

        self._closed = True
        self._turns.clear()
        # Explicitly release pipeline reference to allow GC and GPU memory release
        self._pipeline = None
        logger.debug("Session %s closed", self.meeting_id)

    @property
    def stream_time(self) -> float:
        """Current stream time position in seconds."""
        return self._stream_time

    @property
    def turns(self) -> list[SpeakerTurn]:
        """All accumulated speaker turns for this session."""
        turns = self._turns
        return list(turns)

    @property
    def is_closed(self) -> bool:
        """Check if session is closed."""
        return self._closed
