import argparse
import gc
import io
import json
import re
import shutil
import subprocess
import sys
from datetime import timedelta
from pathlib import Path

BASE_DIR = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(BASE_DIR))

from ui import ConstData, MyUtils

ConstData.BASE_DIR = BASE_DIR
# 可选：强制工作目录
import os


def _safe_stdio():
    """Avoid Windows GBK encode crashes when third-party libs print emoji."""
    try:
        if hasattr(sys.stdout, "reconfigure"):
            sys.stdout.reconfigure(encoding="utf-8", errors="replace")
        else:
            sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace")
    except Exception:
        pass
    try:
        if hasattr(sys.stderr, "reconfigure"):
            sys.stderr.reconfigure(encoding="utf-8", errors="replace")
        else:
            sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding="utf-8", errors="replace")
    except Exception:
        pass


_safe_stdio()
print(f'BASE_DIR---{BASE_DIR}')
os.chdir(BASE_DIR)


def _cleanup(*objs):
    for obj in objs:
        if obj is None:
            continue
        for fn_name in ("close", "release", "cleanup"):
            fn = getattr(obj, fn_name, None)
            if callable(fn):
                try:
                    fn()
                except Exception:
                    pass
                break
    gc.collect()
    try:
        import torch
        if MyUtils.check_cuda_available():
            torch.cuda.empty_cache()
            if hasattr(torch.cuda, "ipc_collect"):
                torch.cuda.ipc_collect()
    except Exception:
        pass


def _init_text_segments(text):
    """分句 → 仅处理逗号/句号 → 统计纯文本字符长度（核心）"""
    raw_text = re.sub(r"\s+", " ", text.strip().replace("\n", " "))
    if not raw_text:
        return [raw_text]
    # 统一新规则：
    # 1) 逗号/句号/问号/叹号都分句（中英文）
    # 2) 数字中的小数点（如 10.0）不做分句，也不移除
    decimal_dot_token = "__DECIMAL_DOT__"
    protected_text = re.sub(r"(?<=\d)\.(?=\d)", decimal_dot_token, raw_text)
    segments = [
        seg.strip()
        for seg in re.findall(
            r'([^,\uFF0C.\u3002!?\uFF01\uFF1F]+[,\uFF0C.\u3002!?\uFF01\uFF1F]?)',
            protected_text
        )
        if seg.strip()
    ]
    if not segments:
        segments = [protected_text]

    clean_segments = []
    for seg in segments:
        clean_text = re.sub(r"[,\uFF0C.\u3002]", "", seg)
        clean_text = re.sub(r"\s+", " ", clean_text).strip()
        clean_text = clean_text.replace(decimal_dot_token, ".")
        clean_segments.append(clean_text if clean_text else seg.replace(decimal_dot_token, "."))

    if not clean_segments:
        fallback_text = re.sub(r"[,\uFF0C.\u3002]", "", protected_text)
        fallback_text = fallback_text.replace(decimal_dot_token, ".")
        clean_segments = [re.sub(r"\s+", " ", fallback_text).strip()]
    return clean_segments


SAMPLE_RATE = 44100


def fmt(t):
    td = timedelta(seconds=t)
    return f"{td.seconds // 3600:02d}:{(td.seconds % 3600) // 60:02d}:{td.seconds % 60:02d},{int(td.microseconds / 1000):03d}"


def get_media_duration(media_path: str) -> float:
    """通用函数：获取视频/音频时长"""
    try:
        cmd = (
            f'ffprobe -v error -show_entries format=duration '
            f'-of default=noprint_wrappers=1:nokey=1 "{media_path}"'
        )
        result = subprocess.run(
            cmd, shell=True, capture_output=True, text=True, timeout=30
        )
        duration = float(result.stdout.strip()) if result.stdout.strip() else 0.0
        return duration
    except Exception as e:
        print(f"获取媒体时长失败（{media_path}）：{e}")
        return 0.0


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--payload", required=True)
    parser.add_argument("--result", required=True)
    args = parser.parse_args()

    result = {"ok": False, "duration": 0, "error": ""}
    voxcpm_model = None
    audio_generator = None
    vox_cpm_audio = None
    temp_dir = os.path.join(BASE_DIR, "temp_audio_t")
    os.makedirs(temp_dir, exist_ok=True)
    audio_files = []
    try:
        with open(args.payload, "r", encoding="utf-8") as f:
            payload = json.load(f)

        from app import voxcpm
        from ui import ConstData

        voxcpm_model_path = payload.get("voxcpm_model_path") or os.path.join(BASE_DIR,
                                                                             "_internal/models/openbmb__VoxCPM1.5"
                                                                             )
        zipenhancer_model_path = payload.get("zipenhancer_model_path") or os.path.join(BASE_DIR,
                                                                                       "_internal/models/iic/speech_zipenhancer_ans_multiloss_16k_base"
                                                                                       )
        print(f'voxcpm_model_path---{voxcpm_model_path}')
        voxcpm_model = voxcpm.VoxCPM(
            voxcpm_model_path=voxcpm_model_path,
            zipenhancer_model_path=zipenhancer_model_path,
        )
        import soundfile as sf
        import numpy as np
        import librosa
        segments = _init_text_segments(payload["text"])
        speed = payload["voice_rate"]
        pitch = payload["voice_pitch"]
        volume = payload["voice_volume"]
        current_time = 0.0
        srt_lines = []
        audio_path = payload["audio_path"]
        # full_wav = []
        from ui import MyUtils
        prompt_wav_path, prompt_text = MyUtils.prepare_for_voxcpm(payload["prompt_wav_path"], payload["prompt_text"])
        for i, text in enumerate(segments, 1):
            full_audio = voxcpm_model.generate(
                text=text,
                prompt_wav_path=prompt_wav_path,
                prompt_text=prompt_text,
                cfg_value=2.0,
                inference_timesteps=10,
                normalize=True,
                denoise=True,
                retry_badcase=True,
                retry_badcase_max_times=3,
                retry_badcase_ratio_threshold=6.0,
            )
            if isinstance(full_audio, np.ndarray) and full_audio.size > 0:
                # ✅ 核心修复2：统一音频为一维数组（适配soundfile写入，避免维度报错）
                if full_audio.ndim > 1:
                    full_audio = full_audio.squeeze()  # 降维：二维→一维
                    print(f"✅ 音频数组降维处理：{full_audio.ndim} 维")

                # -------------------------- 1. 语速处理（直接使用0.5-2.0） --------------------------
                if not (0.5 <= speed <= 2.0):
                    raise ValueError(f"语速参数超出范围！请输入0.5-2.0，当前值：{speed}")
                if speed != 1.0:
                    full_audio = librosa.effects.time_stretch(full_audio, rate=speed)
                    print(f"✅ 语速调整完成 | 倍率：{speed}x")

                # -------------------------- 2. 语调处理（0.5-1.5 → 半音偏移） --------------------------
                if not (0.5 <= pitch <= 1.5):
                    raise ValueError(f"语调参数超出范围！请输入0.5-1.5，当前值：{pitch}")
                    # 映射公式：半音偏移量 = (pitch - 1.0) * 20 → 1.0=0半音，0.5=-10半音，1.5=+10半音
                    # 可根据需求调整乘数（如*12则0.5=-6半音，1.5=+6半音）
                pitch_shift = (pitch - 1.0) * 20
                if pitch_shift != 0:
                    full_audio = librosa.effects.pitch_shift(
                        y=full_audio,
                        sr=SAMPLE_RATE,
                        n_steps=pitch_shift
                    )
                    print(f"✅ 语调调整完成 | 相对值：{pitch} → 半音偏移：{pitch_shift:.1f}")

                    # -------------------------- 3. 音量处理（0-1 → 分贝增益） --------------------------
                if not (0 <= volume <= 1):
                    raise ValueError(f"音量参数超出范围！请输入0-1，当前值：{volume}")
                    # 映射公式：分贝增益 = (volume - 0.5) * 20 → 0.5=0dB，0=-10dB，1=+10dB
                volume_gain = (volume - 0.5) * 20
                if volume_gain != 0:
                    volume_multiplier = 10 ** (volume_gain / 20)
                    full_audio = full_audio * volume_multiplier
                    full_audio = np.clip(full_audio, -1.0, 1.0)  # 防止削波失真
                    print(f"✅ 音量调整完成 | 相对值：{volume} → 分贝增益：{volume_gain:.1f}dB")

                duration = len(full_audio) / SAMPLE_RATE
                start = fmt(current_time)
                end = fmt(current_time + duration)
                # 字幕
                srt_lines.append(f"{i}\n{start} --> {end}\n{text}\n\n")
                current_time += duration
                # 拼接音频
                # full_wav.extend(full_audio)
                # 生成分段文件名
                ext = "wav"
                if '.mp3' in audio_path:
                    ext = 'mp3'
                segment_path = os.path.join(temp_dir, f"seg_{i}.{ext}")

                # 写入单段音频
                sf.write(segment_path, full_audio, SAMPLE_RATE)

                audio_files.append(segment_path)
        concat_file = os.path.join(temp_dir, "list.txt")

        replace_data = '\\'
        with open(concat_file, "w", encoding="utf-8") as f:
            for path in audio_files:
                f.write(f"file '{path.replace(replace_data, '/')}'\n")
        cmd = [
            "ffmpeg",
            "-y",
            "-f", "concat",
            "-safe", "0",
            "-i", concat_file,
            "-c", "copy",
            payload["audio_path"]
        ]
        subprocess.run(cmd, check=True)

        # sf.write(payload["audio_path"], full_wav, SAMPLE_RATE)
        total_duration = get_media_duration(payload["audio_path"])
        print(f"✅ 音频同步生成完成 | 总时长：{total_duration:.2f}秒 | 字幕段数：{len(srt_lines)}")
        with open(payload["srt_path"], "w", encoding="utf-8") as f:
            f.writelines(srt_lines)
        result["ok"] = total_duration > 0
        result["duration"] = total_duration
        if total_duration <= 0:
            result["error"] = "duration <= 0"
    except Exception as e:
        import traceback

        err_detail = traceback.format_exc()
        result["ok"] = False
        result["duration"] = 0
        result["error"] = err_detail
    finally:
        try:
            if audio_generator and hasattr(audio_generator, "close"):
                audio_generator.close()
        except Exception:
            pass
        try:
            print(f'temp_dir---{temp_dir}---{os.path.exists(temp_dir)}')
            if os.path.exists(temp_dir):
                shutil.rmtree(temp_dir, ignore_errors=True)
        except Exception:
            pass
        _cleanup(vox_cpm_audio, voxcpm_model, audio_generator)

    os.makedirs(os.path.dirname(args.result), exist_ok=True)
    with open(args.result, "w", encoding="utf-8") as f:
        json.dump(result, f, ensure_ascii=False)


if __name__ == "__main__":
    main()
