import gc
import sys
from pathlib import Path

import torch
from faster_whisper import WhisperModel

BASE_DIR = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(BASE_DIR))
# 可选：强制工作目录
import os
print(f'BASE_DIR---{BASE_DIR}')
os.chdir(BASE_DIR)


def convert_seconds_to_srt_time(seconds):
    """你的 SRT 时间转换函数（复制过来）"""
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    secs = int(seconds % 60)
    ms = int((seconds - int(seconds)) * 1000)
    return f"{hours:02d}:{minutes:02d}:{secs:02d},{ms:03d}"


def get_packed_file_path(relative_path):
    # 拼接完整路径
    return os.path.join(BASE_DIR, relative_path)


def process_audio(audio_path, subtitle_path):
    """独立的转录+生成SRT函数"""
    success = False
    try:
        # 加载模型（独立进程中，无需考虑主进程资源）
        print("cuda" if os.environ.get("CUDA_VISIBLE_DEVICES") else "cpu")
        model_path = get_packed_file_path("_internal/models/guillaumekln__faster-whisper-medium")
        print(f'model_path---{model_path}')
        model = WhisperModel(
            model_path,
            device="cuda" if os.environ.get("CUDA_VISIBLE_DEVICES") else "cpu",
            compute_type="float16" if os.environ.get("CUDA_VISIBLE_DEVICES") else "int8"
        )

        print(f'whisper 开始处理，{subtitle_path}')
        # 转录音频
        segments, info = model.transcribe(
            audio=audio_path,
            beam_size=5,
            vad_parameters={"min_silence_duration_ms": 500},
            vad_filter=True,
            chunk_length=30,
            condition_on_previous_text=False,
            suppress_tokens=[-1]
        )
        print(
            f"识别语言={info.language}, "
            f"语言概率={info.language_probability}, "
            f"音频时长={info.duration}, "
            f"VAD后时长={info.duration_after_vad}"
        )
        # 生成 SRT
        with open(subtitle_path, 'w', encoding='utf-8') as f:
            idx = 1
            for segment in segments:
                srt_start = convert_seconds_to_srt_time(segment.start)
                srt_end = convert_seconds_to_srt_time(segment.end)
                f.write(f"{idx}\n{srt_start} --> {srt_end}\n{segment.text.strip()}\n\n")
                idx += 1
        print('whisper处理完成，完成生成字幕')
        success = True
    except Exception as e:
        import traceback
        print(f"子进程处理失败：{traceback.format_exc()}")  # 输出到标准错误流
    finally:
        model = None
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    return success


if __name__ == "__main__":
    # 从命令行参数获取 audio_path 和 subtitle_path
    if len(sys.argv) != 3:
        print("参数错误：需要传入 audio_path 和 subtitle_path", file=sys.stderr)
        sys.exit(1)
    audio_path = sys.argv[1]
    subtitle_path = sys.argv[2]
    # 执行处理
    success = process_audio(audio_path, subtitle_path)
    sys.exit(0 if success else 1)  # 成功返回0，失败返回1
