"""DubbingBridgeModel: zh/en 语音 -> en/es/ja/zh 配音 (音色保持, 6 方向) 的 HF 风格单模型接口。

单进程 (transformers 5.6 + torch, 见 README.md "运行环境"):
    from modeling_dubbing import DubbingBridgeModel
    model = DubbingBridgeModel.from_pretrained('/path/to/dubbing_2b_fulldir_cv3')
    wav, sr = model.dub('中文.wav', lang='en')                     # zh->en
    wav, sr = model.dub('english.wav', lang='zh')                  # en->zh
    model.dub('english.wav', lang='ja', out_wav='dubbed.wav')      # en->ja, 同时存盘

lang 为目标语种 (∈ {en, es, ja, zh}); 源语种由 ST LM 转写原文自动判定
(含 CJK => zh 源, 否则 en 源), 支持 zh→en/es/ja + en→zh/es/ja 共 6 方向。
"""
import json
import os
import sys

import torch


class DubbingBridgeModel:
    """ST LM (EXP-53, 冻结) + Hidden2CVMapper (lang 条件, 6 方向) + CosyVoice3 (0.5B, 冻结)
    整条全方向配音链。"""

    def __init__(self, model_dir, device='cuda', mode=None):
        model_dir = os.path.abspath(model_dir)
        assert os.path.isfile(os.path.join(model_dir, 'config.json')), model_dir
        self.model_dir = model_dir
        self.config = json.load(open(os.path.join(model_dir, 'config.json')))
        os.environ['DUBBING_HOME'] = model_dir  # vendored code 读这个
        self.mode = mode or self.config.get('default_mode', 'h2cv')
        self.langs = self.config.get('langs', ['en', 'es', 'ja', 'zh'])
        for p in (os.path.join(model_dir, 'code/tts'), os.path.join(model_dir, 'code')):
            if p not in sys.path:
                sys.path.insert(0, p)
        from pipeline import DubPipeline  # noqa: E402  (code/tts/pipeline.py)
        self._pipe = DubPipeline(model_dir, device=device)
        self.sample_rate = self.config.get('sample_rate', 24000)

    @classmethod
    def from_pretrained(cls, model_dir, **kw):
        return cls(model_dir, **kw)

    def dub(self, audio_path, lang='en', out_wav=None, mode=None, chunk=False,
            return_info=False):
        """audio_path: zh 或 en 语音; lang: 目标语种 'en'/'es'/'ja'/'zh'
        -> (waveform (1,T) float32, sample_rate=24000)。
        out_wav 给定时同时存盘; mode 仅接受 'h2cv' (hidden2cv af 口径, 本包唯一模式);
        chunk 为长音频编排扩展位 (v4.0.0-fulldir 未启用);
        return_info=True 时追加返回 info dict (含 zh(源文)/tgt_raw/tgt_cv 译文文本,
        ja 的 tgt_cv 为 kata 化文本, src_lang 为判定出的源语种)。"""
        mode = mode or self.mode
        assert mode in ('h2cv',), mode
        assert lang in self.langs, f'lang must be one of {self.langs}, got {lang!r}'
        if out_wav is None:
            import tempfile
            fd, out_wav = tempfile.mkstemp(suffix='.wav', prefix='dub_')
            os.close(fd)
        w, sr, info = self._pipe.dub(audio_path, lang=lang, out_wav=out_wav, chunk=chunk)
        return (w, sr, info) if return_info else (w, sr)