"""通用 HuggingFace 数据集适配器（流式）。

config 条目示例：
  - name: ascend
    type: hf
    hf_id: CAiRE/ASCEND
    split: test
    text_field: transcription
    audio_field: audio          # 默认 audio
    lang: zh-en
    domain: general
    trust_remote_code: false

流式（streaming=True）避免为跑小样而下载整个数据集；每条把音频物料化为 16k wav 缓存。
"""

from __future__ import annotations

import io
import os
from collections.abc import Iterator

import soundfile as sf

from ..manifest import Sample
from .base import Dataset, write_wav_16k


class HFDataset(Dataset):
    def __init__(self, cfg: dict, audio_cache: str):
        self.name = cfg["name"]
        self.hf_id = cfg["hf_id"]
        self.hf_config = cfg.get("hf_config")
        self.split = cfg.get("split", "test")
        self.text_field = cfg["text_field"]
        self.audio_field = cfg.get("audio_field", "audio")
        self.lang = cfg.get("lang", "zh")
        self.domain = cfg.get("domain", "general")
        self.trust_remote_code = cfg.get("trust_remote_code", False)
        self.cache_dir = os.path.join(audio_cache, self.name)

    def samples(self, limit: int | None = None) -> Iterator[Sample]:
        from datasets import Audio, load_dataset

        ds = load_dataset(
            self.hf_id,
            self.hf_config,
            split=self.split,
            streaming=True,
            trust_remote_code=self.trust_remote_code,
        )
        # 关闭自动解码（新版 datasets 需 torchcodec）：拿原始 bytes/path 自己用 soundfile 读，避开重依赖。
        ds = ds.cast_column(self.audio_field, Audio(decode=False))
        n = 0
        for i, row in enumerate(ds):
            if limit is not None and n >= limit:
                break
            text = (row.get(self.text_field) or "").strip()
            if not text:
                continue
            audio = row.get(self.audio_field)
            if not audio:
                continue
            wav_path = os.path.join(self.cache_dir, f"{i:06d}.wav")
            if not os.path.exists(wav_path):
                array, sr = self._decode(audio)
                if array is None:
                    continue
                write_wav_16k(array, sr, wav_path)
            n += 1
            yield Sample(
                id=f"{self.name}-{i:06d}",
                audio_path=wav_path,
                ref_text=text,
                lang=self.lang,
                domain=self.domain,
                dataset=self.name,
            )

    @staticmethod
    def _decode(audio: dict):
        """从未解码的 HF 音频列读出 (array, sr)：优先 bytes，其次本地 path。"""
        if audio.get("bytes"):
            data, sr = sf.read(io.BytesIO(audio["bytes"]), dtype="float32", always_2d=False)
            return data, sr
        if audio.get("path") and os.path.exists(audio["path"]):
            data, sr = sf.read(audio["path"], dtype="float32", always_2d=False)
            return data, sr
        return None, None