from typing import List

from joblib import Parallel, delayed
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 LIDModel


@registry.components.register("identifier")
class LanguageIdentifier(BaseComponent):
    name: str = "identifier"
    description: str = "Identify spoken language of audio channels"
    input_tags: List[str] = ["Audio.channels"]
    output_tags: List[str] = ["AudioChannel.language"]

    model: LIDModel | None = Field(
        None, description="The LID model to use for language identification"
    )
    num_threads: int | None = Field(
        None, description="The number of threads to use for parallel processing"
    )

    def predict(self, audios: AudioList) -> AudioList:
        all_channels = [channel for audio in audios for channel in audio.channels]
        batch_size = max(1, len(all_channels) // self.num_threads)
        _ = Parallel(n_jobs=self.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:
        language = self.model.identify(waveform=channel.waveform)
        channel.language = language
        return channel

    def setup(
        self,
        model: str,
        num_threads: int = 1,
        **kwargs,
    ) -> Self:
        self.load_model_plugin(group="fasr.lid_models", model=model)
        model: LIDModel = registry.lid_models.get(model)()
        self.model = model.from_checkpoint(**kwargs)
        self.num_threads = num_threads
        return self
