src/omnisub/cli.py

"""CLI cho OmniSub.

Ví dụ:
  # Kiểm tra nhanh Bước 1 (không nạp model): parse + gom cảnh
  python -m omnisub.cli prepare phim.srt

  # Toàn bộ pipeline trên Colab (có video + Qwen3-Omni)
  python -m omnisub.cli run phim.srt --video phim.mp4 --config config.yaml
"""

from __future__ import annotations

import argparse
import logging
import sys
from pathlib import Path

from .config import Config
from .pipeline import prepare_subtitles, run_pipeline


def _force_utf8() -> None:
    """Ép stdout/stderr sang UTF-8 (console Windows mặc định cp1252 không in được tiếng Việt)."""
    for stream in (sys.stdout, sys.stderr):
        reconfigure = getattr(stream, "reconfigure", None)
        if reconfigure is not None:
            try:
                reconfigure(encoding="utf-8")
            except (ValueError, OSError):
                pass


def _setup_logging(verbose: bool) -> None:
    _force_utf8()
    logging.basicConfig(
        level=logging.DEBUG if verbose else logging.INFO,
        format="%(asctime)s [%(levelname)s] %(message)s",
        datefmt="%H:%M:%S",
    )


def _build_backend(config: Config, args):
    """Khởi tạo backend Qwen3-Omni (lazy import để Bước 1 không cần GPU).

    `--backend none` chỉ chạy Bước 1 (chuẩn bị SRT) để kiểm tra nhanh, không nạp model.
    """
    if args.backend == "none":
        return None
    if args.backend == "qwen":
        from .backends.transformers_qwen import TransformersQwenOmniBackend

        return TransformersQwenOmniBackend(
            model_name=config.models["omni"],
            quant=config.models["quant"],
            cache_dir=args.cache_dir,
        )
    raise SystemExit(f"Backend không hỗ trợ: {args.backend}")


def cmd_prepare(args) -> int:
    config = Config.load(args.config)
    cues, scenes = prepare_subtitles(args.srt, config)
    print(f"Cue: {len(cues)} | Cảnh: {len(scenes)}")
    for sc in scenes[: args.preview]:
        print(f"  Cảnh #{sc.scene_id} [{sc.start:.1f}-{sc.end:.1f}s] {len(sc.cues)} cue")
        for c in sc.cues:
            print(f"    [{c.index}] {c.text[:60]}")
    if args.out:
        from .srt import write_srt

        write_srt(cues, args.out)
        print(f"Đã ghi cue đã gom cảnh ra: {args.out}")
    return 0


def cmd_run(args) -> int:
    config = Config.load(args.config)
    backend = _build_backend(config, args)
    result = run_pipeline(
        args.srt,
        video=args.video,
        config=config,
        backend=backend,
        work_dir=args.work_dir,
        hf_token=args.hf_token,
        do_diarize=not args.no_diarize,
        do_profile=not args.no_profile,
        do_translate=not args.no_translate,
    )
    print(f"Xong. Output: {result.output_srt}")
    print(f"Report: {result.report_path}")
    return 0


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="omnisub",
        description="Dịch phụ đề SRT ZH→VI bằng Qwen3-Omni (đa phương thức).",
    )
    p.add_argument("-v", "--verbose", action="store_true", help="log chi tiết")
    p.add_argument("--config", default="config.yaml", help="đường dẫn config.yaml")
    sub = p.add_subparsers(dest="command", required=True)

    pp = sub.add_parser("prepare", help="Bước 1: parse + gom cảnh (không cần model)")
    pp.add_argument("srt", help="file SRT nguồn (ZH)")
    pp.add_argument("--out", help="ghi cue đã gom cảnh ra file SRT (tùy chọn)")
    pp.add_argument("--preview", type=int, default=5, help="số cảnh in thử")
    pp.set_defaults(func=cmd_prepare)

    pr = sub.add_parser("run", help="Chạy pipeline đầy đủ (Bước 1→5)")
    pr.add_argument("srt", help="file SRT nguồn (ZH)")
    pr.add_argument("--video", help="file video tương ứng (mp4...)")
    pr.add_argument(
        "--backend", choices=["none", "qwen"], default="qwen",
        help="backend model (none = chỉ chạy Bước 1)",
    )
    pr.add_argument("--cache-dir", help="thư mục cache model (Drive/Cache)")
    pr.add_argument("--work-dir", help="thư mục làm việc tạm (frame/audio)")
    pr.add_argument("--hf-token", help="HuggingFace token cho pyannote")
    pr.add_argument("--no-diarize", action="store_true", help="bỏ Bước 2")
    pr.add_argument("--no-profile", action="store_true", help="bỏ Bước 3")
    pr.add_argument("--no-translate", action="store_true", help="bỏ Bước 4")
    pr.set_defaults(func=cmd_run)

    return p


def main(argv=None) -> int:
    _force_utf8()
    parser = build_parser()
    args = parser.parse_args(argv)
    _setup_logging(getattr(args, "verbose", False))
    return args.func(args)


if __name__ == "__main__":
    sys.exit(main())