cli.py

import click
import configparser
import os
import sys
import mimetypes
import multiprocessing
import threading
from tqdm import tqdm
from typing import Optional, Callable, Any

from director import (
    direct_it,
    PROGRESS_ALIGN_START,
    PROGRESS_ALIGN_END
)
from only_align import align_it, load_config
from just_utils import is_video_file
from main_logger import setup_logging

AUDIO_EXTENSIONS = {'.wav', '.mp3', '.flac', '.m4a', '.ogg', '.aac'}

def key_error_handler(func: Callable) -> Callable:
    def wrapper(*args, **kwargs) -> Any:
        try:
            return func(*args, **kwargs)
        except KeyError as e:
            click.echo(f'报错: {e}')
            while True:
                a = input('疑似配置文件出现错误，是否删除？(y/n)')
                if a.lower() == 'y':
                    os.remove(os.path.abspath('config.ini'))
                    click.echo('已删除配置文件。')
                    break
                elif a.lower() == 'n':
                    click.echo('已跳过。')
                    break
    return wrapper

def is_audio_file(path: str) -> bool:
    """
    判断文件是否为音频文件。
    """
    ext = os.path.splitext(path)[1].lower()
    if ext in AUDIO_EXTENSIONS:
        return True
    mime_type, _ = mimetypes.guess_type(path)
    return mime_type is not None and mime_type.startswith('audio/')

@click.group(invoke_without_command=True)
@click.pass_context
def cli(ctx: click.Context) -> None:
    """Script Director 命令行程序"""
    if ctx.invoked_subcommand is None:
        click.echo(ctx.get_help())

@cli.command(name='help', short_help='显示命令帮助信息。若不指定命令名称，则显示全局帮助。')
@click.argument('command_name', required=False, type=str)
def help_command(command_name: Optional[str] = None) -> None:
    """显示指定命令的详细用法。如果不提供命令名，则显示所有命令的摘要。"""
    if command_name:
        if command_name not in cli.commands:
            click.echo(f"错误：未知命令 '{command_name}'")
            return
        cmd_obj = cli.commands[command_name]
        with click.Context(cmd_obj) as ctx:
            click.echo(cmd_obj.get_help(ctx))
    else:
        click.echo(cli.get_help(click.Context(cli)))

@cli.command(name='init', short_help='初始化配置文件（交互式）')
def init_config() -> None:
    """
    初始化配置文件 config.ini。

    \b
    你需要依次输入以下配置项：
    - Faster Whisper 本地模型路径（例如：D:/models/whisper-medium）
    - 台本与音频所使用的语言代码（例如：zh, en, ja）
    - 设备类型（cuda 或 cpu）
    - 计算类型（float16 或 int8）
    - 可选的高级配置

    如果配置文件已存在，会询问是否覆盖。
    """
    config_path = os.path.abspath('config.ini')
    if os.path.exists(config_path):
        overwrite = input('配置文件已存在，是否覆盖？(y/n): ').lower()
        if overwrite != 'y':
            click.echo('已取消初始化。')
            return
    click.echo('开始配置...')
    conf = configparser.ConfigParser()

    # 获取 common 节参数
    while True:
        model = input('请输入Faster Whisper本地模型路径: ').strip()
        lang = input('请输入台本与音频所使用的语言代码: ').strip()
        device = input('请输入设备类型(cuda/cpu): ').strip()
        compute = input('请输入计算类型(float16/int8): ').strip()
        if model and lang and device and compute:
            break
        click.echo('配置文件不能有空项，请重新输入。')
    conf['common'] = {
        'model': model,
        'lang': lang,
        'device': device,
        'compute': compute
    }

    # 获取 advanced 节参数（可选）
    click.echo("可选高级参数设置（直接回车使用默认值）:")
    gap_penalty = input('gap_penalty (对齐惩罚值，默认 -10): ').strip()
    similarity_offset = input('similarity_offset (相似度偏移，默认 50): ').strip()
    default_duration = input('default_duration (默认字幕时长/秒，默认 5.0): ').strip()
    max_combine = input('max_combine (最大合并片段数，默认 5): ').strip()
    beam_size = input('beam_size (束搜索宽度，默认 5): ').strip()
    vad_filter = input('vad_filter (启用语音活动检测，默认 False): ').strip()
    vad_parameters = input('vad_parameters (VAD参数 JSON，默认 {}): ').strip()

    conf['advanced'] = {}
    if gap_penalty:
        conf['advanced']['gap_penalty'] = gap_penalty
    if similarity_offset:
        conf['advanced']['similarity_offset'] = similarity_offset
    if default_duration:
        conf['advanced']['default_duration'] = default_duration
    if max_combine:
        conf['advanced']['max_combine'] = max_combine
    if beam_size:
        conf['advanced']['beam_size'] = beam_size
    if vad_filter:
        conf['advanced']['vad_filter'] = str(vad_filter.lower() in ('true', '1', 'yes'))
    if vad_parameters:
        conf['advanced']['vad_parameters'] = vad_parameters

    with open(config_path, 'w', encoding='utf-8') as configfile:
        conf.write(configfile)
    click.echo('配置文件已保存。')

@cli.command(name='config', short_help='快速修改配置项，格式：key=value')
@click.argument('key_value', type=str)
def modify_config(key_value: str) -> None:
    """
    直接通过命令行修改配置文件中的某项配置。

    \b
    用法示例：
        python cli.py config model=/new/path/to/model
        python cli.py config lang=en

    可修改的配置项有common与advanced节，详见已经初始化的config.ini。
    """
    config_path = os.path.abspath('config.ini')
    if not os.path.exists(config_path):
        click.echo('错误：配置文件不存在，请先运行 init 命令。')
        sys.exit(1)
    if '=' not in key_value:
        click.echo('错误：参数格式应为 key=value')
        sys.exit(1)
    key, value = key_value.split('=', 1)
    key = key.strip()
    value = value.strip()
    if not key or not value:
        click.echo('错误：键和值不能为空')
        sys.exit(1)

    conf = configparser.ConfigParser()
    conf.read(config_path, encoding='utf-8')

    # 判断属于哪个节（默认 common）
    advanced_keys = ['gap_penalty', 'similarity_offset', 'default_duration', 'max_combine', 'beam_size', 'vad_filter', 'vad_parameters']
    section = 'advanced' if key in advanced_keys else 'common'
    if section not in conf:
        conf[section] = {}
    conf[section][key] = value

    with open(config_path, 'w', encoding='utf-8') as configfile:
        conf.write(configfile)
    click.echo(f'已更新配置项 {section}.{key} = {value}')

@key_error_handler
@cli.command(name='process', short_help='通过台本文件与音频文件生成字幕文件')
@click.argument('input_str', type=str)
@click.option('-t', '--type', type=click.Choice(['srt', 'lrc'], case_sensitive=False),
              default='srt', help='输出字幕格式，支持 srt 或 lrc，默认为 srt')
@click.option('-n', '--name', type=str, default=None, help='指定输出文件的名称，该项不包含扩展名')
@click.option('-p', '--preprocess', is_flag=True, default=False,
              help='预处理台本，删除空行和方括号内容（如角色标识）')
@click.option('-s', '--shorter', is_flag=True, default=False,
              help='启用短句模式（按标点分割长句，生成更精确的字幕）')
def process_command(input_str: str, type: str, name: str, preprocess: bool, shorter: bool) -> None:
    """
    处理音频和台本，生成字幕文件。

    \b
    INPUT_STR 必须包含两个文件路径，用英文逗号分隔，例如：
        python cli.py process "audio.wav,script.txt"
        或
        python cli.py process "script.txt,subtitles.srt"

    \b
    程序会自动识别文件类型：
    - 扩展名为 .txt 的视为台本文件
    - 扩展名为 .srt 或 .lrc 的视为已有字幕文件（启用只对齐模式）
    - 其他音频扩展名或常见视频扩展名视为音频/视频文件（视频将自动提取音频轨）

    \b
    生成的字幕文件与音频文件同名，扩展名为 .srt 或 .lrc，保存在同一目录。
    如果输入的是字幕文件（只对齐模式），则输出文件默认与字幕文件同名。

    运行此命令前必须通过 init 命令创建配置文件 config.ini。

    \b
    如果指定 -p 或 --preprocess，则会对台本进行清洗：
    - 删除空行
    - 删除方括号内的内容（例如 [人名]、[动作说明]）
    - 去除多余空格

    \b
    如果指定 -s 或 --shorter，则启用短句模式，按标点分割长句，生成更精确的字幕。
    注意：只对齐模式下短句模式无效，程序会发出警告并忽略该选项。
    """
    setup_logging(console=True, file=True)

    files = [f.strip() for f in input_str.split(',')]
    if len(files) != 2:
        raise click.UsageError('输入参数必须包含两个文件路径，用逗号分隔。')

    # 识别文件类型（同时转为绝对路径）
    script_path = None
    audio_path = None
    subtitle_path = None
    for f in files:
        f = os.path.abspath(f)
        if not os.path.isfile(f):
            raise click.FileError(f, f'文件不存在：{f}')
        ext = os.path.splitext(f)[1].lower()
        if ext == '.txt':
            script_path = f
        elif ext in ('.srt', '.lrc'):
            subtitle_path = f
        else:
            if is_audio_file(f) or is_video_file(f):
                audio_path = f
            else:
                raise click.UsageError(f'无法识别文件类型: {f}')

    if not script_path:
        raise click.UsageError('未找到台本文件（.txt）')

    config_path = os.path.abspath('config.ini')

    # 只对齐模式：提供了字幕文件
    if subtitle_path:
        click.echo("检测到已有字幕文件，启用只对齐模式（不进行语音识别）")
        if shorter:
            click.echo("警告：只对齐模式下短句模式无效，将忽略 -s/--shorter 选项。")
        output_dir = os.path.dirname(subtitle_path)
        if name:
            base = name
        else:
            base = os.path.splitext(os.path.basename(subtitle_path))[0]
        output_path = os.path.join(output_dir, f"{base}.{type}")
        align_it(
            script_path=script_path,
            subtitle_path=subtitle_path,
            output_path=output_path,
            output_format=type,
            preprocess=preprocess,
            short_sentences=shorter,
            config_path=config_path
        )
        click.echo(f'字幕已生成：{output_path}')
        return

    # 否则走原有流程（音频+台本）
    if not audio_path:
        raise click.UsageError('未找到音频文件或已有字幕文件')

    if not os.path.exists(config_path):
        raise click.ClickException('配置文件不存在，请先运行 init 命令。')

    settings = load_config(config_path)
    model = settings['model']
    lang = settings['lang']
    device = settings['device']
    compute = settings['compute']
    if not all([model, lang, device, compute]):
        raise ValueError('配置文件不完整，请重新运行 init 命令或检查 config.ini')

    audio_dir = os.path.dirname(audio_path)
    if name:
        audio_basename = name
    else:
        audio_basename = os.path.splitext(os.path.basename(audio_path))[0]
    output_filename = f"{audio_basename}.{type}"
    output_path = os.path.join(audio_dir, output_filename)

    progress_queue = multiprocessing.Queue()
    progress_bar = None
    stop_event = threading.Event()

    def progress_updater():
        nonlocal progress_bar
        while not stop_event.is_set():
            try:
                p = progress_queue.get(timeout=0.5)
                if p is None:
                    break
                # 仅当进入对齐阶段（p > PROGRESS_ALIGN_START）才创建进度条
                if p > PROGRESS_ALIGN_START:
                    if progress_bar is None:
                        progress_bar = tqdm(total=100, desc="Aligning", unit="%")
                        progress_bar.n = 0
                        progress_bar.refresh()
                    # 将底层进度值映射到 0~100%
                    # 底层范围: PROGRESS_ALIGN_START (80) -> PROGRESS_ALIGN_END (99)
                    # 映射公式: percent = (p - 80) / (99 - 80) * 100
                    percent = int((p - PROGRESS_ALIGN_START) / (PROGRESS_ALIGN_END - PROGRESS_ALIGN_START) * 100)
                    percent = max(0, min(100, percent))  # 限制在 0~100
                    if progress_bar.n < percent:
                        progress_bar.n = percent
                        progress_bar.refresh()
                    # 对齐完成后强制进度条达到 100% 并关闭
                    if p >= PROGRESS_ALIGN_END:
                        if progress_bar and progress_bar.n < 100:
                            progress_bar.n = 100
                            progress_bar.refresh()
                        if progress_bar:
                            progress_bar.close()
                        break
            except:
                continue

    progress_thread = threading.Thread(target=progress_updater, daemon=True)
    progress_thread.start()

    try:
        direct_it(
            audio_path=audio_path,
            script_path=script_path,
            output_path=output_path,
            local_model_path=model,
            language=lang,
            device=device,
            compute_type=compute,
            preprocess=preprocess,
            short_sentences=shorter,
            progress_queue=progress_queue,
            verbose=False
        )
    finally:
        stop_event.set()
        progress_queue.put(None)
        progress_thread.join(timeout=2)
        if progress_bar is not None and not progress_bar.disable:
            progress_bar.close()

    click.echo(f'字幕已生成：{output_path}')

if __name__ == '__main__':
    import multiprocessing
    multiprocessing.freeze_support()
    cli()