#!/usr/bin/env python3
"""
KB Whisper - 本地语音识别工具
基于 Whisper.cpp 的音视频转文字

功能：
- 自动将视频/音频转换为 16kHz WAV 格式
- 支持指定语言（默认中文）
- 输出标准 JSON 格式供 KB Workflow 使用
"""

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


def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(
        description='KB Whisper - 本地语音识别工具',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  python whisper.py --input video.mp4
  python whisper.py --input video.mp4 --model base
  python whisper.py --input video.mp4 --output ./output --background
        """
    )
    parser.add_argument(
        '--input',
        required=True,
        help='输入音视频文件路径'
    )
    parser.add_argument(
        '--model',
        default='base',
        choices=['tiny', 'base', 'small', 'medium', 'large'],
        help='Whisper 模型类型 (默认: base)'
    )
    parser.add_argument(
        '--output',
        help='输出目录 (默认: 输入文件所在目录)'
    )
    parser.add_argument(
        '--background',
        action='store_true',
        help='后台异步执行'
    )
    parser.add_argument(
        '--language',
        default='zh',
        help='指定语言 (默认: zh)，例如: en, zh, ja'
    )

    return parser.parse_args()


def get_whisper_cpp_path():
    """获取 whisper.cpp 路径"""
    # 优先从环境变量获取
    path = os.environ.get('WHISPER_CPP_PATH')
    if path:
        path = os.path.expanduser(path)
        if os.path.exists(path):
            return path

    # 尝试常见位置
    common_paths = [
        os.path.expanduser('~/whisper.cpp'),
        '/usr/local/bin/whisper',
        '/opt/whisper.cpp',
    ]

    for cp in common_paths:
        if os.path.exists(cp):
            return cp

    raise FileNotFoundError(
        "未找到 whisper.cpp。请设置环境变量 WHISPER_CPP_PATH "
        "或安装到 ~/whisper.cpp\n"
        "安装方法: git clone https://github.com/ggerganov/whisper.cpp.git ~/whisper.cpp && cd ~/whisper.cpp && make"
    )


def check_model(whisper_path, model):
    """检查模型是否存在"""
    model_path = os.path.join(whisper_path, 'models', f'ggml-{model}.bin')
    if not os.path.exists(model_path):
        raise FileNotFoundError(
            f"模型文件不存在: {model_path}\n"
            f"请运行以下命令下载模型:\n"
            f"cd {whisper_path} && ./models/download-ggml-model.sh {model}"
        )
    return model_path


def convert_to_wav(input_path, output_dir):
    """将音视频转换为 16kHz WAV 格式

    Whisper 需要 WAV 格式的输入文件。
    """
    input_path = os.path.abspath(input_path)

    # 如果已经是 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:
                print(f"输入文件已是 16kHz WAV 格式，直接使用")
                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}")

    try:
        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
    except subprocess.CalledProcessError as e:
        print(f"音频转换失败: {e}")
        raise RuntimeError(f"无法转换音频格式: {e}")
    except FileNotFoundError:
        raise RuntimeError("未找到 ffmpeg，请安装: sudo apt-get install ffmpeg")


def run_whisper(whisper_path, audio_path, model, output_dir, language='zh'):
    """运行 Whisper 语音识别

    Args:
        whisper_path: whisper.cpp 目录
        audio_path: WAV 音频文件路径
        model: 模型名称
        output_dir: 输出目录
        language: 语言代码
    """
    # 输出文件名（使用原始输入文件名，不是 audio.wav）
    # 从 output_dir 和 audio_path 推断原始文件名
    original_stem = Path(audio_path).stem.replace('_audio', '')
    output_basename = original_stem

    # 构建命令
    cmd = [
        os.path.join(whisper_path, 'whisper-cli'),
        '-m', os.path.join(whisper_path, 'models', f'ggml-{model}.bin'),
        '-f', audio_path,
        '-of', os.path.join(output_dir, output_basename),
        '-l', language,
        '--output-txt',
        '--output-srt',
        '--output-json',
    ]

    print(f"执行命令: {' '.join(cmd)}")

    # 运行
    result = subprocess.run(
        cmd,
        capture_output=True,
        text=True,
        cwd=whisper_path,
        timeout=3600  # 1小时超时
    )

    if result.returncode != 0:
        # 检查是否有警告但仍成功
        if "error" in result.stderr.lower() and not os.path.exists(
            os.path.join(output_dir, f"{output_basename}.txt")
        ):
            raise RuntimeError(
                f"Whisper 执行失败:\n"
                f"标准输出: {result.stdout}\n"
                f"错误输出: {result.stderr}"
            )

    return result


def parse_srt_for_duration(srt_file):
    """从 SRT 文件解析时长"""
    if not os.path.exists(srt_file):
        return 0

    with open(srt_file, 'r', encoding='utf-8') as f:
        content = f.read()

    # 找到最后一个时间戳
    pattern = r'(\d{2}:\d{2}:\d{2}),\d{3} --> \d{2}:\d{2}:\d{2},\d{3}'
    matches = re.findall(pattern, content)

    if not matches:
        return 0

    # 解析最后一个时间戳
    last_time = matches[-1]
    h, m, s = last_time.split(':')
    return int(h) * 3600 + int(m) * 60 + int(s)


def create_workflow_json(input_file, output_dir, model, language, duration):
    """创建标准格式的 JSON 输出，供 KB Workflow 使用"""
    input_path = Path(input_file)
    basename = input_path.stem

    txt_file = os.path.join(output_dir, f"{basename}.txt")
    srt_file = os.path.join(output_dir, f"{basename}.srt")
    json_file = os.path.join(output_dir, f"{basename}.json")

    # 读取文本内容
    text_content = ""
    if os.path.exists(txt_file):
        with open(txt_file, 'r', encoding='utf-8') as f:
            text_content = f.read().strip()

    # 解析 SRT 获取分段
    segments = []
    if os.path.exists(srt_file):
        with open(srt_file, '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 enumerate(matches):
            segments.append({
                "id": idx,
                "start": round(parse_time(start), 2),
                "end": round(parse_time(end), 2),
                "text": text.strip()
            })

    # 构建 JSON
    workflow_json = {
        "success": True,
        "tool": "whisper",
        "model": model,
        "language": language,
        "input_file": str(input_path),
        "output_files": {
            "txt": f"{basename}.txt",
            "srt": f"{basename}.srt",
            "json": f"{basename}.json"
        },
        "duration_seconds": duration,
        "segments": segments,
        "text": text_content
    }

    # 写入 JSON
    with open(json_file, 'w', encoding='utf-8') as f:
        json.dump(workflow_json, f, ensure_ascii=False, indent=2)

    return workflow_json


def main():
    """主函数"""
    args = parse_arguments()

    # 检查输入文件
    input_file = os.path.abspath(args.input)
    if not os.path.exists(input_file):
        print(f"错误: 输入文件不存在: {input_file}", file=sys.stderr)
        sys.exit(1)

    # 获取输出目录
    if args.output:
        output_dir = os.path.abspath(args.output)
    else:
        output_dir = os.path.dirname(input_file)

    os.makedirs(output_dir, exist_ok=True)

    # 获取 whisper.cpp 路径
    try:
        whisper_path = get_whisper_cpp_path()
        print(f"✓ Whisper.cpp 路径: {whisper_path}")
    except FileNotFoundError as e:
        print(f"✗ 错误: {e}", file=sys.stderr)
        sys.exit(1)

    # 检查模型
    try:
        check_model(whisper_path, args.model)
        print(f"✓ 使用模型: {args.model}")
    except FileNotFoundError as e:
        print(f"✗ 错误: {e}", file=sys.stderr)
        sys.exit(1)

    # 后台执行
    if args.background:
        print("启动后台处理...")
        log_file = os.path.join(output_dir, f"{Path(input_file).stem}_whisper.log")

        with open(log_file, 'w') as log:
            subprocess.Popen(
                [sys.executable, __file__,
                 '--input', input_file,
                 '--model', args.model,
                 '--output', output_dir,
                 '--language', args.language],
                stdout=log,
                stderr=log,
                start_new_session=True
            )

        print(f"✓ 后台任务已启动")
        print(f"  日志: {log_file}")
        print(f"  输出: {output_dir}")

        result = {
            "success": True,
            "background": True,
            "status": "running",
            "log_file": log_file,
            "output_dir": output_dir
        }
        print(json.dumps(result, ensure_ascii=False, indent=2))
        return

    # 前台执行
    print(f"\n📁 处理文件: {input_file}")
    print(f"📁 输出目录: {output_dir}")

    try:
        # 1. 转换为 WAV
        print(f"\n🔄 步骤1: 转换音频格式...")
        audio_path = convert_to_wav(input_file, output_dir)

        # 2. 运行 Whisper
        print(f"\n🎙️ 步骤2: 语音识别 (语言: {args.language})...")
        run_whisper(whisper_path, audio_path, args.model, output_dir, args.language)

        # 3. 清理临时 WAV 文件
        if audio_path != input_file and audio_path.endswith('_audio.wav'):
            try:
                os.remove(audio_path)
                print("✓ 清理临时文件")
            except Exception:
                pass

        # 4. 获取时长并创建 JSON
        print(f"\n📝 步骤3: 生成输出文件...")
        basename = Path(input_file).stem
        srt_file = os.path.join(output_dir, f"{basename}.srt")
        duration = parse_srt_for_duration(srt_file)

        workflow_json = create_workflow_json(
            input_file, output_dir, args.model, args.language, duration
        )

        # 输出结果
        print(f"\n✅ 处理完成!")
        print(f"   时长: {duration} 秒")
        print(f"   分段数: {len(workflow_json['segments'])}")
        print(f"\n📄 输出文件:")
        for fmt in ['txt', 'srt', 'json']:
            fpath = os.path.join(output_dir, f"{basename}.{fmt}")
            if os.path.exists(fpath):
                size = os.path.getsize(fpath)
                print(f"   • {fmt.upper()}: {fpath} ({size:,} bytes)")

        print(f"\n📋 Workflow JSON:")
        print(json.dumps({
            "success": workflow_json["success"],
            "tool": "whisper",
            "text_length": len(workflow_json["text"]),
            "segments_count": len(workflow_json["segments"])
        }, ensure_ascii=False, indent=2))

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


if __name__ == '__main__':
    main()