#
# Copyright (c) 2025, Speechmatics / Cantab Research Ltd
#

from __future__ import annotations

import datetime
import logging
import os
import urllib.request
from typing import Optional
from urllib.parse import urlparse

import numpy as np
import onnxruntime as ort
from pydantic import BaseModel

os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1"
from transformers import WhisperFeatureExtractor  # noqa: E402

logger = logging.getLogger(__name__)

# Base model from HuggingFace
ONNX_HF_URL = os.getenv(
    "SMART_TURN_HF_URL", "https://huggingface.co/pipecat-ai/smart-turn-v3/resolve/main/smart-turn-v3.0.onnx"
)
ONNX_MODEL_PATH = os.getenv("SMART_TURN_MODEL_PATH", ".models/smart-turn-v3.0.onnx")


class SmartTurnPredictionResult(BaseModel):
    """Prediction result from the smart turn detector.

    Attributes:
        prediction: True for complete, False for incomplete
        probability: Probability of completion (sigmoid output)
        processing_time: Time taken to process the audio (in seconds)
        error: Error message if an error occurred
    """

    prediction: bool = False
    probability: float = 0.0
    processing_time: Optional[float] = None
    error: Optional[str] = None


class SmartTurnDetector:
    """Smart Turn Detector.

    Uses Pipecat's opensource acoustic model for determining if an audio sample
    is predicted to be complete or incomplete.

    Further information at https://github.com/pipecat-ai/smart-turn
    """

    def __init__(self, auto_init: bool = True, threshold: float = 0.8):
        """Create the new SmartTurnDetector.

        Args:
            auto_init: Whether to automatically initialise the detector.
            threshold: Probability threshold for turn completion (0.0-1.0).
        """

        # Has initialized
        self._is_initialized: bool = False

        # Threshold
        self._threshold: float = threshold

        # If auto_init is True, setup the detector
        if auto_init:
            self.setup()

    def setup(self) -> None:
        """Setup the detector.

        Initialises the ONNX model and feature extractor.
        """

        try:
            # Check / download the model
            self.download_model()

            # Check the model downloaded
            if not self.model_exists():
                logger.warning("Smart Turn model not found. Please download the model first.")
                return

            # Build the session
            self.session = self.build_session(ONNX_MODEL_PATH)

            # Load the feature extractor
            self.feature_extractor = WhisperFeatureExtractor(chunk_length=8)

            # Set initialized
            self._is_initialized = True

        except Exception as e:
            logger.error(f"Failed to setup SmartTurnDetector: {e}")

    def build_session(self, onnx_path: str) -> ort.InferenceSession:
        """Build the ONNX session and load resources.

        Args:
            onnx_path: Path to the ONNX model.

        Returns:
            ONNX inference session.
        """

        # Build the session
        so = ort.SessionOptions()
        so.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
        so.inter_op_num_threads = 1
        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

        # Return the new session
        return ort.InferenceSession(onnx_path, sess_options=so)

    async def predict(
        self, audio_array: bytes, language: str, sample_rate: int = 16000, sample_width: int = 2
    ) -> SmartTurnPredictionResult:
        """Predict whether an audio segment is complete (turn ended) or incomplete.

        Args:
            audio_array: Numpy array containing audio samples at 16kHz. The function
                will convert the audio into float32 and truncate to 8 seconds (keeping the end)
                or pad to 8 seconds.
            language: Language of the audio.
            sample_rate: Sample rate of the audio.
            sample_width: Sample width of the audio.

        Returns:
            Prediction result containing completion status and probability.
        """

        # Check if initialized
        if not self._is_initialized:
            return SmartTurnPredictionResult(error="SmartTurnDetector is not initialized")

        # Check a valid language
        if not self.valid_language(language):
            logger.warning(f"Invalid language: {language}.")
            return SmartTurnPredictionResult(error=f"Invalid language: {language}")

        # Record start time
        start_time = datetime.datetime.now()

        # Convert into numpy array
        dtype = np.int16 if sample_width == 2 else np.int8
        int16_array: np.ndarray = np.frombuffer(audio_array, dtype=dtype).astype(np.int16)

        # Process audio using Whisper's feature extractor
        inputs = self.feature_extractor(
            int16_array,
            sampling_rate=sample_rate,
            return_tensors="np",
            padding="max_length",
            max_length=8 * sample_rate,
            truncation=True,
            do_normalize=True,
        )

        # Extract features and ensure correct shape for ONNX
        input_features = inputs.input_features.squeeze(0).astype(np.float32)
        input_features = np.expand_dims(input_features, axis=0)

        # Run ONNX inference
        outputs = self.session.run(None, {"input_features": input_features})

        # Extract probability (ONNX model returns sigmoid probabilities)
        probability = outputs[0][0].item()

        # Make prediction (True for Complete, False for Incomplete)
        prediction = probability >= self._threshold

        # Record end time
        end_time = datetime.datetime.now()

        # Return the result
        return SmartTurnPredictionResult(
            prediction=prediction,
            probability=probability,
            processing_time=float((end_time - start_time).total_seconds()),
        )

    @staticmethod
    def truncate_audio_to_last_n_seconds(
        audio_array: np.ndarray, n_seconds: float = 8.0, sample_rate: int = 16000
    ) -> np.ndarray:
        """Truncate audio to last n seconds or pad with zeros to meet n seconds.

        Args:
            audio_array: Numpy array containing audio samples at 16kHz.
            n_seconds: Number of seconds to truncate to.
            sample_rate: Sample rate of the audio.

        Returns:
            Numpy array truncated to last n seconds or padded with zeros.
        """

        # Calculate the max samples we should have
        max_samples = int(n_seconds * sample_rate)

        # Truncate if longer
        if len(audio_array) > max_samples:
            return audio_array[-max_samples:]

        # Pad if shorter
        elif len(audio_array) < max_samples:
            padding = max_samples - len(audio_array)
            return np.pad(audio_array, (padding, 0), mode="constant", constant_values=0)

        # Otherwise return the array
        return audio_array

    @staticmethod
    def download_model() -> None:
        """Download the ONNX model.

        This will check if the model has been downloaded and is available in the
        location specified by the SMART_TURN_MODEL_PATH environment variable.

        If not, it will download the model from HuggingFace.
        """

        # Check if model file exists
        if SmartTurnDetector.model_exists():
            return

        # Check the URL for valid schemes
        parsed_url = urlparse(ONNX_HF_URL)
        if parsed_url.scheme not in ("http", "https"):
            logger.error(f"Invalid URL scheme: {parsed_url.scheme}")
            return

        # Report to the user
        logger.warning("Smart Turn model not found. Downloading from HuggingFace...")

        # Create the directory
        os.makedirs(os.path.dirname(ONNX_MODEL_PATH), exist_ok=True)

        # Download
        urllib.request.urlretrieve(ONNX_HF_URL, ONNX_MODEL_PATH)  # nosec B310

    @staticmethod
    def model_exists() -> bool:
        """Check the model has been downloaded.

        Returns:
            True if the model file exists, False otherwise.
        """
        return os.path.exists(ONNX_MODEL_PATH)

    @staticmethod
    def valid_language(language: str) -> bool:
        """Check if the language is valid.

        Args:
            language: Language code to validate.

        Returns:
            True if the language is supported, False otherwise.
        """
        return language in [
            "ar",
            "bn",
            "zh",
            "da",
            "nl",
            "de",
            "en",
            "fi",
            "fr",
            "hi",
            "id",
            "it",
            "ja",
            "ko",
            "mr",
            "no",
            "pl",
            "pt",
            "ru",
            "es",
            "tr",
            "uk",
            "vi",
        ]
