"""DiariZen speaker diarization segmentation inference."""

import json
import os
from pathlib import Path
from typing import Optional

import numpy as np


class DiarizenSegmenter:
    """Speaker diarization segmentation using NPU CNN frontend + CPU backend.

    Pipeline:
        1. Audio preprocessing (resample + LayerNorm) on CPU.
        2. CNN feature extraction on AX650 NPU (U16).
        3. Transformer + Conformer + Classifier on CPU (ONNX Runtime).
    """

    def __init__(
        self,
        cnn_model_path: str,
        backend_model_path: str,
        meta_path: Optional[str] = None,
    ):
        self._cnn_path = Path(cnn_model_path)
        self._backend_path = Path(backend_model_path)

        if meta_path is None:
            meta_path = Path(__file__).parent / "model_meta.json"
        with open(meta_path) as f:
            self._meta = json.load(f)

        pp = self._meta["preprocess"]
        self._sample_rate = pp["sample_rate"]
        self._duration_s = pp["duration_seconds"]
        self._eps = pp.get("layer_norm_eps", 1e-5)
        self._num_samples = int(self._sample_rate * self._duration_s)

        self._cnn_session = None
        self._backend_session = None

    def _init_cnn(self):
        """Initialize NPU CNN inference session."""
        try:
            from axengine import InferenceSession
        except ImportError:
            raise RuntimeError(
                "pyaxengine is required for NPU inference. "
                "Install from: https://github.com/AXERA-TECH/pyaxengine"
            )
        self._cnn_session = InferenceSession(str(self._cnn_path))

    def _init_backend(self):
        """Initialize CPU backend ONNX inference session."""
        import onnxruntime as ort
        self._backend_session = ort.InferenceSession(
            str(self._backend_path),
            providers=["CPUExecutionProvider"],
        )

    def __call__(self, audio: np.ndarray, sample_rate: int) -> np.ndarray:
        """Run segmentation inference.

        Args:
            audio: 1-D float32 waveform.
            sample_rate: Original sample rate.

        Returns:
            Log-probabilities of shape (1, frames, 11), float32.
        """
        from .preprocess import preprocess_audio

        waveform_ln = preprocess_audio(
            audio, sample_rate,
            target_sr=self._sample_rate,
            duration_s=self._duration_s,
            eps=self._eps,
        )

        if self._cnn_session is None:
            self._init_cnn()
        cnn_outputs = self._cnn_session.run(
            {self._cnn_session.input_names()[0]: waveform_ln}
        )
        cnn_features = cnn_outputs[0]

        if self._backend_session is None:
            self._init_backend()
        backend_inputs = {
            self._backend_session.get_inputs()[0].name: cnn_features
        }
        log_probs = self._backend_session.run(None, backend_inputs)[0]
        return log_probs

    @property
    def num_frames(self) -> int:
        return 199

    @property
    def num_classes(self) -> int:
        return 11