from typing import Iterable, List

from joblib import Parallel, delayed
from pydantic import validate_call
from typing_extensions import Self

from fasr.config import registry
from fasr.data.audio import (
    Audio,
    AudioChannel,
    AudioList,
    AudioSpan,
    AudioSpanList,
    AudioToken,
    AudioTokenList,
)
from fasr.model import ASRModel

from .base import BaseComponent


@registry.components.register("recognizer")
class SpeechRecognizer(BaseComponent):
    name: str = "recognizer"
    description: str = "Recognize speech from audio channels"
    input_tags: List[str] = ["Audio.channels"]
    output_tags: List[str] = ["AudioChannel.tokens"]
    model: ASRModel | None = None
    num_threads: int = 1
    batch_size_s: int | None = None
    hotwords: frozenset[str] | None = None

    def predict(self, audios: AudioList) -> AudioList:
        if self.hotwords is not None:
            for audio in audios:
                audio: Audio
                if audio.hotwords is None:
                    audio.set_hotwords(self.hotwords)
        groups = audios.group_by_hotwords()
        _ = Parallel(
            n_jobs=self.num_threads, prefer="threads", pre_dispatch="1 * n_jobs"
        )(
            delayed(self.predict_step)(batch_segments, list(hotwords))
            for hotwords, group_audios in groups.items()
            for batch_segments in self.batch_audio_segments(audios=group_audios)
        )
        return audios

    def predict_step(
        self, batch_segments: List[AudioSpan | AudioChannel], hotwords: List[str] = []
    ) -> List[AudioSpan | AudioChannel]:
        batch_waveforms = [seg.waveform for seg in batch_segments]
        if len(batch_waveforms) == 0:
            return batch_segments
        sample_rate = batch_waveforms[0].sample_rate  # 一个batch的音频片段采样率相同
        batch_data = [waveform.data for waveform in batch_waveforms]
        batch_tokens = self.model.transcribe(
            batch=batch_data, sample_rate=sample_rate, hotwords=hotwords
        )
        for seg, tokens in zip(batch_segments, batch_tokens):
            if not isinstance(seg, AudioChannel):
                if seg.start_ms is not None and seg.end_ms is not None:
                    for token in tokens:
                        token: AudioToken
                        if (
                            token.start_ms and token.end_ms
                        ):  # not all asr model predict timestamp
                            token.start_ms += seg.start_ms
                            token.end_ms += seg.start_ms
            seg.tokens = AudioTokenList(docs=tokens)
        return batch_segments

    def batch_audio_segments(self, audios: AudioList) -> Iterable[AudioSpanList]:
        """将音频片段组成批次。
        步骤：
        - 1. 将音频片段按照时长排序。
        - 2. 将音频片段按照时长分组，每组时长不超过batch_size_s。
        """
        all_segments = []
        for audio in audios:
            if not audio.is_bad:
                for channel in audio.channels:
                    if channel.segments is None:  # 兼容没有vad模型的情况
                        all_segments.append(channel)
                    else:
                        for seg in channel.segments:
                            all_segments.append(seg)
        return self.batch_segments(all_segments)

    def batch_segments(
        self, segments: Iterable[AudioSpan | AudioChannel]
    ) -> Iterable[List[AudioSpan | AudioChannel]]:
        """将音频片段组成动态批次（按时长贪心组批）。

        组批目标是让同一批次内的样本尽量“时长相近”，从而降低 padding 带来的浪费。

        具体规则：
        - **批次上限**：`batch_size_s` 表示一个批次的“等效总时长上限”（秒），会换算成
          毫秒 `batch_size_ms = batch_size_s * 1000`。
          当 `batch_size_s is None` 时，不按时长上限切分，直接把全部片段组成一个批次。
        - **排序**：先将 `segments` 收集成 list，并按 `duration_ms` 升序排序。
        - **贪心装箱**：遍历排序后的片段，维护当前批次的最大时长 `batch_max_duration_ms`。
          尝试把片段加入当前批次时，计算加入后的等效总时长：

            `current = max(batch_max_duration_ms, seg.duration_ms) * (len(batch) + 1)`

          这里的 `max_duration_ms * batch_size` 等价于“把批次内所有样本 padding 到同一最大时长”的
          计算量/时长近似。
        - **切批**：若 `current > batch_size_ms`，先 `yield` 当前批次，再以当前片段开启新批次；
          否则将片段加入当前批次并更新 `batch_max_duration_ms`。
        - **收尾**：遍历结束后若当前批次非空，会再 `yield` 一次。

        Returns:
            Iterable[List[AudioSpan | AudioChannel]]: 每次产出一个批次（list）。
        """
        segments = [seg for seg in segments]
        sorted_segments = self.sort_segments(segments)
        if self.batch_size_s is None:
            if len(sorted_segments) > 0:
                yield sorted_segments
            return

        batch_size_ms = self.batch_size_s * 1000
        batch: List[AudioSpan | AudioChannel] = []
        batch_max_duration_ms = 0
        for seg in sorted_segments:
            max_duration_ms = max(batch_max_duration_ms, seg.duration_ms)
            current_batch_duration_ms = max_duration_ms * (
                len(batch) + 1
            )  # 如果加入当前片段，是否超过batch_size_s
            if current_batch_duration_ms > batch_size_ms:
                yield batch
                batch = [seg]
                batch_max_duration_ms = seg.duration_ms
            else:
                batch.append(seg)
                batch_max_duration_ms = max_duration_ms
        if len(batch) > 0:
            yield batch

    def sort_segments(
        self, segments: List[AudioSpan | AudioChannel]
    ) -> List[AudioSpan | AudioChannel]:
        return sorted(segments, key=lambda x: x.duration_ms)

    @validate_call
    def setup(
        self,
        model: str,
        batch_size_s: int | None = None,
        num_threads: int = 1,
        hotwords: Iterable[str] | None = None,
        **kwargs,
    ) -> Self:
        self.load_model_plugin(group="fasr.asr_models", model=model)
        model: ASRModel = registry.asr_models.get(model)()
        self.model = model.from_checkpoint(**kwargs)
        self.num_threads = num_threads
        self.batch_size_s = batch_size_s
        self.hotwords = frozenset(hotwords) if hotwords is not None else None
        return self
