"""自定义集：JSONL manifest，每行一条样本。

格式（audio 路径相对 manifest 文件所在目录解析）：
  {"audio": "clips/a.wav", "text": "帮我把这份 weekly report 整理一下", "lang": "zh-en", "domain": "tech"}

lang 缺省按文本自动猜（含中文且含 ASCII 字母→zh-en；纯 ASCII→en；否则 zh）。
"""

from __future__ import annotations

import json
import os
import re
import sys
from collections.abc import Iterator

from ..manifest import Sample

_CJK = re.compile(r"[一-鿿]")
_ASCII_ALPHA = re.compile(r"[A-Za-z]")


def _guess_lang(text: str) -> str:
    has_cjk = bool(_CJK.search(text))
    has_en = bool(_ASCII_ALPHA.search(text))
    if has_cjk and has_en:
        return "zh-en"
    if has_en and not has_cjk:
        return "en"
    return "zh"


class CustomDataset:
    def __init__(self, cfg: dict):
        self.name = cfg.get("name", "custom")
        self.manifest = cfg["manifest"]
        self.domain_default = cfg.get("domain", "general")
        self._base = os.path.dirname(os.path.abspath(self.manifest))

    def samples(self, limit: int | None = None) -> Iterator[Sample]:
        with open(self.manifest, encoding="utf-8") as f:
            n = 0
            for i, line in enumerate(f):
                line = line.strip()
                if not line or line.startswith("#"):
                    continue
                if limit is not None and n >= limit:
                    break
                row = json.loads(line)
                text = (row.get("text") or "").strip()
                audio = row["audio"]
                if not os.path.isabs(audio):
                    audio = os.path.join(self._base, audio)
                if not os.path.exists(audio):
                    print(f"[custom] 跳过缺失音频：{audio}", file=sys.stderr)
                    continue
                lang = row.get("lang") or _guess_lang(text)
                n += 1
                yield Sample(
                    id=row.get("id", f"{self.name}-{i:06d}"),
                    audio_path=audio,
                    ref_text=text,
                    lang=lang,
                    domain=row.get("domain", self.domain_default),
                    dataset=self.name,
                )
