"""数据集适配器基类与音频物料化工具。"""

from __future__ import annotations

import os
from collections.abc import Iterator

import numpy as np
import soundfile as sf

from ..manifest import Sample

TARGET_SR = 16000


class Dataset:
    name: str = "base"

    def samples(self, limit: int | None = None) -> Iterator[Sample]:
        raise NotImplementedError


def write_wav_16k(array: np.ndarray, sr: int, out_path: str) -> None:
    """把（可能任意采样率的）单/多声道数组写成 16k/mono/16bit wav。"""
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    data = np.asarray(array, dtype=np.float32)
    if data.ndim > 1:
        data = data.mean(axis=1)
    if sr != TARGET_SR:
        import soxr

        data = soxr.resample(data, sr, TARGET_SR)
    data = np.clip(data, -1.0, 1.0)
    sf.write(out_path, (data * 32767.0).astype(np.int16), TARGET_SR, subtype="PCM_16")