from typing import List

from joblib import Parallel, delayed
from loguru import logger
from pydantic import Field
from typing_extensions import Self

from fasr.components.base import BaseComponent
from fasr.config import registry
from fasr.data import AudioChannel, AudioList
from fasr.model import VADModel


@registry.components.register("detector")
class VoiceDetector(BaseComponent):
    name: str = "detector"
    description: str = "Detect voice segments from audio channels"
    input_tags: List[str] = ["Audio.channels"]
    output_tags: List[str] = ["AudioChannel.segments"]

    model: VADModel | None = Field(
        None, description="The vad model to use for voice detection"
    )
    threshold: float | None = Field(
        None,
        description="The threshold for audio duration, only audios with duration > threshold will be detected",
    )
    num_threads: int | None = Field(
        None, description="The number of threads to use for parallel processing"
    )

    def predict(self, audios: AudioList) -> AudioList:
        _audios = self.check_audios_threshold(audios)
        all_channels = [channel for audio in _audios for channel in audio.channels]
        num_threads = self.num_threads or 2
        batch_size = max(1, len(all_channels) // num_threads)
        _ = Parallel(n_jobs=num_threads, prefer="threads", batch_size=batch_size)(
            delayed(self.predict_channel)(channel) for channel in all_channels
        )
        return audios

    def predict_channel(self, channel: AudioChannel) -> AudioChannel:
        segments = self.model.detect(waveform=channel.waveform)
        channel.segments = segments
        return channel

    def check_audios_threshold(self, audios: AudioList) -> AudioList:
        # 大于threshold的音频才进行检测
        if self.threshold is not None:
            new_audios = AudioList()
            for audio in audios:
                if audio.duration > self.threshold:
                    new_audios.append(audio)
                else:
                    logger.warning(
                        f"Audio {audio.id} duration {audio.duration} < threshold {self.threshold}, skip detection"
                    )
            return new_audios
        return audios

    def setup(
        self,
        model: str,
        num_threads: int | None = None,
        threshold: float = None,
        **kwargs,
    ) -> Self:
        self.load_model_plugin(group="fasr.vad_models", model=model)
        self.model: VADModel = registry.vad_models.get(model)()
        self.model = self.model.from_checkpoint(**kwargs)
        self.num_threads = num_threads
        self.threshold = threshold
        return self
