"""音视频转写：faster-whisper 本地模型 + ffmpeg 抽音轨。
输出与 .srt 同格式的「去时间戳、按语义间隙合并」文本，直接进入既有 RAG 分块管线。
"""
from __future__ import annotations

import os
import re
import subprocess
import threading
from pathlib import Path
from .parse import clean_text

MEDIA_SUFFIXES = {".mp4", ".mov", ".m4a", ".mp3", ".wav", ".flv", ".mkv", ".avi", ".webm", ".aac", ".ogg", ".flac"}

# 模型目录：优先环境变量，其次 modelscope 缓存里的 gpustack/faster-whisper-small
_MODEL_DIR_CANDIDATES = [
    os.environ.get("WHISPER_MODEL_DIR", ""),
    str(Path.home() / ".cache/modelscope/models/gpustack--faster-whisper-small/snapshots/master"),
]
_model = None
_model_lock = threading.Lock()


def _get_model():
    global _model
    if _model is not None:
        return _model
    with _model_lock:
        if _model is not None:
            return _model
        from faster_whisper import WhisperModel
        for cand in _MODEL_DIR_CANDIDATES:
            if cand and Path(cand).exists():
                _model = WhisperModel(cand, device="cpu", compute_type="int8", local_files_only=True)
                return _model
        _model = WhisperModel("small", device="cpu", compute_type="int8")
        return _model


def is_media(path: Path) -> bool:
    return path.suffix.lower() in MEDIA_SUFFIXES


def _probe_duration(path: Path) -> float:
    try:
        r = subprocess.run(
            ["ffprobe", "-v", "quiet", "-show_entries", "format=duration",
             "-of", "csv=p=0", str(path)],
            capture_output=True, text=True, timeout=30)
        return float(r.stdout.strip() or 0)
    except Exception:
        return 0.0


def transcribe_media(path: Path, progress_cb=None) -> list[dict]:
    """音视频 → 语义段落列表 [{text, meta:{heading}}]。
    超长音先压缩码率降体积；VAD 过滤静音；按标点+间隙合并成语义段。"""
    src = Path(path)
    work = src
    # 大文件先转 16k mono wav（whisper 输入要求，同时把视频降为纯音轨）
    if src.suffix.lower() not in (".wav",) or _probe_duration(src) > 0:
        work = src.with_suffix(".wav.wav")
        cmd = ["ffmpeg", "-y", "-i", str(src), "-vn", "-ar", "16000", "-ac", "1",
               "-b:a", "64k", str(work)]
        r = subprocess.run(cmd, capture_output=True, text=True, timeout=3600)
        if r.returncode != 0 or not work.exists() or work.stat().st_size == 0:
            raise ValueError(f"ffmpeg 抽音轨失败: {r.stderr[-200:]}")
    try:
        model = _get_model()
        segments, _info = model.transcribe(str(work), language="zh", vad_filter=True,
                                           vad_parameters={"min_silence_duration_ms": 700})
        lines: list[str] = []
        for seg in segments:
            t = (seg.text or "").strip()
            if t:
                lines.append(t)
            if progress_cb and len(lines) % 20 == 0:
                progress_cb(len(lines))
        if not lines:
            return []
        # 合并成语义段：约 15 句或 1200 字一段（与既有 2200 字分块目标匹配）
        merged, buf, size = [], [], 0
        for t in lines:
            buf.append(t)
            size += len(t)
            if size >= 1200 or t.endswith(("。", "！", "？", "!", "?")) and size >= 600:
                merged.append(" ".join(buf))
                buf, size = [], 0
        if buf:
            merged.append(" ".join(buf))
        # 转写结果也经过同一套结构清洗：去掉字幕式序号/时间戳、不可见字符和句首独立口头填充词，
        # 保留正文中的日期、金额、型号等数字。
        return [{"text": clean_text(m), "meta": {"heading": src.stem}}
                for m in merged if clean_text(m)]
    finally:
        if work != src:
            try:
                work.unlink(missing_ok=True)
            except OSError:
                pass