#!/usr/bin/env python3
"""
KB Whisper Diarization - 语音识别 + 说话者分离
使用 Whisper 进行语音识别，pyannote.audio 进行说话者分离
"""

import argparse
import json
import os
import subprocess
import sys
import tempfile
from datetime import datetime
from pathlib import Path
import re


def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(
        description='KB Whisper Diarization - 语音识别 + 说话者分离',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  python whisper_diarization.py --input video.mp4
  python whisper_diarization.py --input interview.mp4 --diarize
  python whisper_diarization.py --input meeting.mp4 --diarize --num-speakers 2
        """
    )
    parser.add_argument('--input', required=True, help='输入音视频文件路径')
    parser.add_argument('--output', help='输出目录 (默认: 文件所在目录)')
    parser.add_argument('--diarize', action='store_true', help='启用说话者分离')
    parser.add_argument('--num-speakers', type=int, help='估计说话者数量')
    parser.add_argument('--language', help='指定语言 (默认: 自动检测)')
    parser.add_argument('--model', default='base', choices=['tiny', 'base', 'small', 'medium', 'large'],
                        help='Whisper 模型 (默认: base)')
    parser.add_argument('--background', action='store_true', help='后台异步执行')
    return parser.parse_args()


def get_whisper_path():
    """获取 whisper.cpp 路径"""
    path = os.environ.get('WHISPER_CPP_PATH')
    if path:
        path = os.path.expanduser(path)
        if os.path.exists(path):
            return path
    return os.path.expanduser('~/whisper.cpp')


def check_whisper_model(whisper_path, model):
    """检查 Whisper 模型"""
    model_path = os.path.join(whisper_path, 'models', f'ggml-{model}.bin')
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Whisper 模型不存在: {model_path}")
    return model_path


def convert_to_wav(input_path, output_dir):
    """转换为 16kHz WAV"""
    if input_path.lower().endswith('.wav'):
        try:
            result = subprocess.run(
                ['ffprobe', '-v', 'error', '-show_entries', 'stream=sample_rate',
                 '-of', 'default=noprint_wrappers=1:nokey=1', input_path],
                capture_output=True, text=True
            )
            if result.returncode == 0 and int(result.stdout.strip()) == 16000:
                return input_path
        except Exception:
            pass

    output_path = os.path.join(output_dir, f"{Path(input_path).stem}_audio.wav")
    print(f"转换音频: {input_path} -> {output_path}")

    subprocess.run([
        'ffmpeg', '-y', '-i', input_path,
        '-acodec', 'pcm_s16le', '-ar', '16000', '-ac', '1',
        '-loglevel', 'error', output_path
    ], check=True, capture_output=True)

    print("✓ 音频转换完成")
    return output_path


def run_whisper(whisper_path, audio_path, model, output_dir, language=None):
    """运行 Whisper 识别"""
    basename = Path(audio_path).stem.replace('_audio', '')
    output_file = os.path.join(output_dir, basename)

    cmd = [
        os.path.join(whisper_path, 'whisper-cli'),
        '-m', os.path.join(whisper_path, 'models', f'ggml-{model}.bin'),
        '-f', audio_path,
        '-of', output_file,
        '-l', language or 'auto',
        '--output-txt', '--output-srt', '--output-json',
    ]

    print(f"执行 Whisper: {' '.join(cmd[:6])} ...")

    result = subprocess.run(
        cmd, capture_output=True, text=True, cwd=whisper_path, timeout=3600
    )

    if result.returncode != 0 and not os.path.exists(f"{output_file}.txt"):
        raise RuntimeError(f"Whisper 失败: {result.stderr}")

    return basename


def run_diarization(audio_path, num_speakers=None):
    """运行 pyannote.audio 说话者分离"""
    print("\n🎙️ 步骤2: 说话者分离...")

    try:
        import torch
        from pyannote.audio import Pipeline
    except ImportError as e:
        print(f"✗ 缺少依赖: {e}")
        print("请安装: pip install pyannote.audio torch")
        raise

    # 检查 HuggingFace Token
    hf_token = os.environ.get('HF_TOKEN')
    if not hf_token:
        raise RuntimeError("需要设置 HF_TOKEN 环境变量\n"
                          "获取方式: https://huggingface.co/settings/tokens")

    # 新版 pyannote 需要先登录
    from huggingface_hub import login
    print("登录 HuggingFace...")
    login(token=hf_token)

    print("加载 pyannote 流水线...")
    pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization")

    # 设置设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"使用设备: {device}")
    pipeline.to(device)

    # 运行分离
    print("执行说话者分离...")
    diarization = pipeline(
        audio_path,
        num_speakers=num_speakers
    )

    # 解析结果
    speakers = {}
    for turn, _, speaker in diarization.itertracks(yield_label=True):
        if speaker not in speakers:
            speakers[speaker] = []
        speakers[speaker].append({
            "start": round(turn.start, 2),
            "end": round(turn.end, 2)
        })

    print(f"✓ 检测到 {len(speakers)} 个说话者: {list(speakers.keys())}")
    return speakers


def merge_results(whisper_txt, whisper_srt, diarization, output_dir, basename):
    """合并 Whisper 和说话者分离结果"""
    print("\n📝 步骤3: 合并结果...")

    # 读取 Whisper 文本
    text_segments = []
    if os.path.exists(whisper_txt):
        with open(whisper_txt, 'r', encoding='utf-8') as f:
            content = f.read()

        # 简单按行分段
        for i, line in enumerate(content.strip().split('\n')):
            if line.strip():
                text_segments.append({
                    "id": i,
                    "text": line.strip(),
                    "start": i * 3.0,  # 估算
                    "end": (i + 1) * 3.0
                })

    # 读取 SRT 获取精确时间
    srt_segments = []
    if os.path.exists(whisper_srt):
        with open(whisper_srt, 'r', encoding='utf-8') as f:
            content = f.read()

        # 解析 SRT
        pattern = r'(\d+)\n(\d{2}:\d{2}:\d{2},\d{3}) --> (\d{2}:\d{2}:\d{2},\d{3})\n(.+?)(?=\n\n|\Z)'
        matches = re.findall(pattern, content, re.DOTALL)

        def parse_time(t):
            h, m, s_ms = t.split(':')
            s, ms = s_ms.split(',')
            return float(h) * 3600 + float(m) * 60 + float(s) + float(ms) / 1000

        for idx, num, start, end, text in matches:
            srt_segments.append({
                "id": int(idx),
                "start": parse_time(start),
                "end": parse_time(end),
                "text": text.strip()
            })

    # 为每个 Whisper 段落分配说话者
    merged_segments = []
    for seg in srt_segments:
        speaker = assign_speaker(seg['start'], seg['end'], diarization)
        merged_segments.append({
            "id": len(merged_segments),
            "speaker": speaker,
            "start": seg['start'],
            "end": seg['end'],
            "text": seg['text']
        })

    # 生成带说话者的 SRT
    output_srt = os.path.join(output_dir, f"{basename}_diarized.srt")
    with open(output_srt, 'w', encoding='utf-8') as f:
        for seg in merged_segments:
            start = format_time(seg['start'])
            end = format_time(seg['end'])
            f.write(f"{seg['id'] + 1}\n")
            f.write(f"{start} --> {end}\n")
            f.write(f"[{seg['speaker']}] {seg['text']}\n\n")

    # 生成合并的 TXT
    output_txt = os.path.join(output_dir, f"{basename}_diarized.txt")
    with open(output_txt, 'w', encoding='utf-8') as f:
        for seg in merged_segments:
            f.write(f"[{seg['speaker']}] {seg['text']}\n")

    # 生成 JSON
    output_json = os.path.join(output_dir, f"{basename}_diarized.json")
    json_data = {
        "success": True,
        "tool": "whisper-diarization",
        "diarization": True,
        "speakers": list(diarization.keys()),
        "input_file": basename,
        "output_files": {
            "txt": f"{basename}_diarized.txt",
            "srt": f"{basename}_diarized.srt",
            "json": f"{basename}_diarized.json"
        },
        "segments": merged_segments
    }
    with open(output_json, 'w', encoding='utf-8') as f:
        json.dump(json_data, f, ensure_ascii=False, indent=2)

    print(f"✓ 生成文件:")
    print(f"  - {output_srt}")
    print(f"  - {output_txt}")
    print(f"  - {output_json}")

    return json_data


def assign_speaker(start, end, diarization):
    """将时间段分配给说话者"""
    best_speaker = "UNKNOWN"
    max_overlap = 0

    for speaker, segments in diarization.items():
        for seg in segments:
            overlap = min(end, seg['end']) - max(start, seg['start'])
            if overlap > max_overlap:
                max_overlap = overlap
                best_speaker = speaker

    return best_speaker


def format_time(seconds):
    """秒数转换为 SRT 时间格式"""
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    secs = int(seconds % 60)
    ms = int((seconds % 1) * 1000)
    return f"{hours:02d}:{minutes:02d}:{secs:02d},{ms:03d}"


def main():
    args = parse_arguments()

    input_file = os.path.abspath(args.input)
    if not os.path.exists(input_file):
        print(f"✗ 文件不存在: {input_file}")
        sys.exit(1)

    output_dir = os.path.abspath(args.output) if args.output else os.path.dirname(input_file)
    os.makedirs(output_dir, exist_ok=True)

    basename = Path(input_file).stem

    print(f"\n📁 输入: {input_file}")
    print(f"📁 输出: {output_dir}")
    print(f"🎙️ 说话者分离: {'是' if args.diarize else '否'}")

    try:
        # 步骤1: Whisper 识别
        print(f"\n🎙️ 步骤1: 语音识别...")
        whisper_path = get_whisper_path()
        check_whisper_model(whisper_path, args.model)

        audio_path = convert_to_wav(input_file, output_dir)
        run_whisper(whisper_path, audio_path, args.model, output_dir, args.language)

        # 步骤2: 说话者分离
        if args.diarize:
            diarization = run_diarization(audio_path, args.num_speakers)
        else:
            diarization = {}

        # 步骤3: 合并结果
        whisper_txt = os.path.join(output_dir, f"{basename}.txt")
        whisper_srt = os.path.join(output_dir, f"{basename}.srt")
        result = merge_results(whisper_txt, whisper_srt, diarization, output_dir, basename)

        # 清理临时音频
        if audio_path != input_file and audio_path.endswith('_audio.wav'):
            try:
                os.remove(audio_path)
            except Exception:
                pass

        print(f"\n✅ 处理完成!")
        print(f"   说话者数量: {len(diarization) if args.diarize else 1}")
        print(f"   段落数量: {len(result['segments'])}")

    except Exception as e:
        print(f"\n✗ 错误: {e}")
        sys.exit(1)


if __name__ == '__main__':
    main()