"""Transcribe the first audio stream of one local video with faster-whisper."""

from __future__ import annotations

import argparse
import json
import math
import os
import shutil
import subprocess
import sys
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Iterable, Sequence


MODEL_NAME = "large-v3"
DEVICE = "cuda"
COMPUTE_TYPE = "float16"
SHORT_AUDIO_LIMIT_SECONDS = 1200.0
CHUNK_CORE_SECONDS = 1200.0
CHUNK_OVERLAP_SECONDS = 5.0


class MediaTranscriptionError(RuntimeError):
    """A user-facing error that should stop the command without a traceback."""


@dataclass(frozen=True)
class OutputPaths:
    directory: Path
    audio: Path
    text: Path
    srt: Path
    json: Path

    @property
    def files(self) -> tuple[Path, Path, Path, Path]:
        return (self.audio, self.text, self.srt, self.json)


@dataclass(frozen=True)
class RuntimeInfo:
    cuda_device_count: int
    cuda_compute_types: tuple[str, ...]


@dataclass(frozen=True)
class TranscriptSegment:
    start: float
    end: float
    text: str


@dataclass(frozen=True)
class TranscriptResult:
    language: str | None
    language_probability: float | None
    segments: tuple[TranscriptSegment, ...]


@dataclass(frozen=True)
class AudioChunk:
    index: int
    total: int
    core_start: float
    core_end: float
    audio_start: float
    audio_end: float

    @property
    def audio_duration(self) -> float:
        return self.audio_end - self.audio_start


ProcessRunner = Callable[[Sequence[str]], subprocess.CompletedProcess[str]]
ModelLoader = Callable[[], tuple[Any, RuntimeInfo]]


def _run_process(command: Sequence[str]) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        list(command),
        check=False,
        capture_output=True,
        text=True,
        encoding="utf-8",
        errors="replace",
    )


def _error_detail(result: subprocess.CompletedProcess[str]) -> str:
    detail = (result.stderr or result.stdout or "未返回错误详情").strip()
    return detail[-1200:]


def _require_executable(name: str) -> str:
    executable = shutil.which(name)
    if executable is None:
        raise MediaTranscriptionError(
            f"未找到 {name}，请先安装并确保它位于当前 PATH。"
        )
    return executable


def _resolve_video(video_path: str | Path) -> Path:
    video = Path(video_path).expanduser().resolve()
    if not video.exists():
        raise MediaTranscriptionError(f"视频文件不存在：{video}")
    if not video.is_file():
        raise MediaTranscriptionError(f"输入路径不是文件：{video}")
    return video


def _build_output_paths(video: Path, output_dir: str | Path | None) -> OutputPaths:
    if output_dir is None:
        directory = video.with_name(f"{video.stem}.transcript")
    else:
        directory = Path(output_dir).expanduser().resolve()

    stem = video.stem
    return OutputPaths(
        directory=directory,
        audio=directory / f"{stem}.audio.flac",
        text=directory / f"{stem}.txt",
        srt=directory / f"{stem}.srt",
        json=directory / f"{stem}.json",
    )


def _ensure_output_available(outputs: OutputPaths) -> None:
    if outputs.directory.exists() and not outputs.directory.is_dir():
        raise MediaTranscriptionError(
            f"输出路径已存在且不是目录：{outputs.directory}"
        )

    conflicts = [path for path in outputs.files if path.exists()]
    if conflicts:
        formatted = "\n".join(f"- {path}" for path in conflicts)
        raise MediaTranscriptionError(
            "输出文件已存在；默认不会覆盖。请改用新的输出目录或先人工处理已有结果：\n"
            f"{formatted}"
        )


def _probe_first_audio_stream(
    video: Path,
    ffprobe: str,
    runner: ProcessRunner,
) -> None:
    command = [
        ffprobe,
        "-v",
        "error",
        "-select_streams",
        "a:0",
        "-show_entries",
        "stream=index",
        "-of",
        "csv=p=0",
        str(video),
    ]
    result = runner(command)
    if result.returncode != 0:
        raise MediaTranscriptionError(
            f"FFprobe 无法检查视频音轨：{_error_detail(result)}"
        )
    if not result.stdout.strip():
        raise MediaTranscriptionError(f"视频不包含可用音轨：{video}")


def _extract_first_audio_stream(
    video: Path,
    audio_path: Path,
    ffmpeg: str,
    runner: ProcessRunner,
) -> None:
    command = [
        ffmpeg,
        "-hide_banner",
        "-loglevel",
        "error",
        "-nostdin",
        "-n",
        "-i",
        str(video),
        "-map",
        "0:a:0",
        "-vn",
        "-ac",
        "1",
        "-ar",
        "16000",
        "-c:a",
        "flac",
        str(audio_path),
    ]
    result = runner(command)
    if result.returncode != 0:
        raise MediaTranscriptionError(
            f"FFmpeg 提取第一条音轨失败：{_error_detail(result)}"
        )
    if not audio_path.is_file() or audio_path.stat().st_size == 0:
        raise MediaTranscriptionError("FFmpeg 未生成有效的 FLAC 音频文件。")


def _probe_audio_duration(
    audio_path: Path,
    ffprobe: str,
    runner: ProcessRunner,
) -> float:
    command = [
        ffprobe,
        "-v",
        "error",
        "-show_entries",
        "format=duration",
        "-of",
        "default=noprint_wrappers=1:nokey=1",
        str(audio_path),
    ]
    result = runner(command)
    if result.returncode != 0:
        raise MediaTranscriptionError(
            f"FFprobe 无法读取 FLAC 时长：{_error_detail(result)}"
        )
    try:
        duration = float(result.stdout.strip())
    except ValueError as exc:
        raise MediaTranscriptionError(
            f"FFprobe 返回了无效的 FLAC 时长：{result.stdout.strip() or '空'}"
        ) from exc
    if not math.isfinite(duration) or duration <= 0:
        raise MediaTranscriptionError(f"FLAC 时长必须为有限正数，实际为：{duration}")
    return duration


def _plan_audio_chunks(duration: float) -> tuple[AudioChunk, ...]:
    if not math.isfinite(duration) or duration <= SHORT_AUDIO_LIMIT_SECONDS:
        return ()

    total = int(math.ceil(duration / CHUNK_CORE_SECONDS))
    chunks: list[AudioChunk] = []
    for offset in range(total):
        core_start = offset * CHUNK_CORE_SECONDS
        core_end = min(duration, core_start + CHUNK_CORE_SECONDS)
        chunks.append(
            AudioChunk(
                index=offset + 1,
                total=total,
                core_start=core_start,
                core_end=core_end,
                audio_start=max(0.0, core_start - CHUNK_OVERLAP_SECONDS),
                audio_end=min(duration, core_end + CHUNK_OVERLAP_SECONDS),
            )
        )
    return tuple(chunks)


def _extract_audio_chunk(
    full_audio: Path,
    chunk_path: Path,
    chunk: AudioChunk,
    ffmpeg: str,
    runner: ProcessRunner,
) -> None:
    if chunk_path.exists():
        raise MediaTranscriptionError(f"临时音频块已存在，拒绝覆盖：{chunk_path}")
    command = [
        ffmpeg,
        "-hide_banner",
        "-loglevel",
        "error",
        "-nostdin",
        "-n",
        "-ss",
        f"{chunk.audio_start:.6f}",
        "-i",
        str(full_audio),
        "-t",
        f"{chunk.audio_duration:.6f}",
        "-ac",
        "1",
        "-ar",
        "16000",
        "-c:a",
        "flac",
        str(chunk_path),
    ]
    result = runner(command)
    if result.returncode != 0:
        raise MediaTranscriptionError(
            f"FFmpeg 提取转写分块 {chunk.index}/{chunk.total} 失败："
            f"{_error_detail(result)}"
        )
    if not chunk_path.is_file() or chunk_path.stat().st_size == 0:
        raise MediaTranscriptionError(
            f"FFmpeg 未生成有效的转写分块 {chunk.index}/{chunk.total}。"
        )


def _load_cuda_model() -> tuple[Any, RuntimeInfo]:
    try:
        import ctranslate2
        from faster_whisper import WhisperModel
    except ImportError as exc:
        raise MediaTranscriptionError(
            "缺少 faster-whisper/CTranslate2，请在项目隔离 Conda 环境中运行。"
        ) from exc

    try:
        device_count = int(ctranslate2.get_cuda_device_count())
        compute_types = tuple(
            sorted(str(item) for item in ctranslate2.get_supported_compute_types(DEVICE))
        )
    except Exception as exc:
        raise MediaTranscriptionError(
            f"CUDA GPU 自检失败；不会降级到 CPU。原始错误：{exc}"
        ) from exc

    if device_count < 1:
        raise MediaTranscriptionError("未检测到可用 CUDA GPU；不会降级到 CPU。")
    if COMPUTE_TYPE not in compute_types:
        raise MediaTranscriptionError(
            "当前 CUDA GPU 不支持 float16；不会降级到 CPU。"
            f"检测到的计算类型：{', '.join(compute_types) or '无'}"
        )

    try:
        model = WhisperModel(
            MODEL_NAME,
            device=DEVICE,
            compute_type=COMPUTE_TYPE,
        )
    except Exception as exc:
        raise MediaTranscriptionError(
            "large-v3 模型加载失败；请检查模型缓存/网络以及 CUDA 12、cuBLAS、"
            f"cuDNN 9。不会降级到 CPU。原始错误：{exc}"
        ) from exc

    return model, RuntimeInfo(device_count, compute_types)


def _collect_transcript(
    model: Any,
    audio_path: Path,
    language: str | None,
) -> TranscriptResult:
    options: dict[str, Any] = {"vad_filter": True}
    if language:
        options["language"] = language

    try:
        segment_iterator, info = model.transcribe(str(audio_path), **options)
        segments = tuple(_normalize_segments(segment_iterator))
    except Exception as exc:
        raise MediaTranscriptionError(
            "GPU 语音转写失败；请检查 CUDA/cuDNN/显存和输入音频。"
            f"不会降级到 CPU。原始错误：{exc}"
        ) from exc

    detected_language = getattr(info, "language", None)
    probability = getattr(info, "language_probability", None)
    if probability is not None:
        probability = float(probability)
    return TranscriptResult(detected_language, probability, segments)


def _normalize_segments(segments: Iterable[Any]) -> Iterable[TranscriptSegment]:
    for segment in segments:
        text = str(segment.text).strip()
        if not text:
            continue
        start = max(0.0, float(segment.start))
        end = max(start, float(segment.end))
        yield TranscriptSegment(start=start, end=end, text=text)


def _project_chunk_segments(
    segments: Iterable[TranscriptSegment],
    chunk: AudioChunk,
    audio_duration: float,
) -> tuple[TranscriptSegment, ...]:
    projected: list[TranscriptSegment] = []
    for segment in segments:
        start = min(audio_duration, max(0.0, chunk.audio_start + segment.start))
        end = min(audio_duration, max(start, chunk.audio_start + segment.end))
        midpoint = (start + end) / 2.0
        belongs = chunk.core_start <= midpoint < chunk.core_end
        if chunk.index == chunk.total:
            belongs = chunk.core_start <= midpoint <= chunk.core_end
        if belongs:
            projected.append(TranscriptSegment(start=start, end=end, text=segment.text))
    return tuple(projected)


def _collect_chunked_transcript(
    model: Any,
    full_audio: Path,
    audio_duration: float,
    language: str | None,
    ffmpeg: str,
    runner: ProcessRunner,
) -> TranscriptResult:
    chunks = _plan_audio_chunks(audio_duration)
    if not chunks:
        raise MediaTranscriptionError("内部错误：短音频不应进入分块转写路径。")

    chunk_path = full_audio.parent / ".transcribe-chunk.flac"
    active_language = language
    result_language = language
    language_probability: float | None = None
    projected: list[TranscriptSegment] = []

    for chunk in chunks:
        print(f"正在转写分块 {chunk.index}/{chunk.total}", flush=True)
        try:
            _extract_audio_chunk(full_audio, chunk_path, chunk, ffmpeg, runner)
            chunk_result = _collect_transcript(model, chunk_path, active_language)
            if chunk.index == 1:
                result_language = chunk_result.language or language
                language_probability = chunk_result.language_probability
                if language is None:
                    if not chunk_result.language:
                        raise MediaTranscriptionError(
                            "首个音频块未返回检测语言，无法为后续块固定语言。"
                        )
                    active_language = chunk_result.language
            projected.extend(
                _project_chunk_segments(chunk_result.segments, chunk, audio_duration)
            )
        finally:
            chunk_path.unlink(missing_ok=True)

    projected.sort(key=lambda item: (item.start, item.end, item.text))
    return TranscriptResult(
        language=result_language,
        language_probability=language_probability,
        segments=tuple(projected),
    )


def _format_txt_timestamp(seconds: float) -> str:
    total_seconds = max(0, int(seconds))
    hours, remainder = divmod(total_seconds, 3600)
    minutes, secs = divmod(remainder, 60)
    return f"{hours:02d}:{minutes:02d}:{secs:02d}"


def _format_srt_timestamp(seconds: float) -> str:
    total_milliseconds = max(0, int(round(seconds * 1000)))
    hours, remainder = divmod(total_milliseconds, 3_600_000)
    minutes, remainder = divmod(remainder, 60_000)
    secs, milliseconds = divmod(remainder, 1000)
    return f"{hours:02d}:{minutes:02d}:{secs:02d},{milliseconds:03d}"


def _write_transcript_files(
    video: Path,
    staged: OutputPaths,
    result: TranscriptResult,
) -> None:
    text_content = "\n\n".join(
        f"[{_format_txt_timestamp(segment.start)}] {segment.text}"
        for segment in result.segments
    )
    if text_content:
        text_content += "\n"
    staged.text.write_text(text_content, encoding="utf-8")

    srt_blocks = []
    for index, segment in enumerate(result.segments, start=1):
        srt_blocks.append(
            "\n".join(
                [
                    str(index),
                    f"{_format_srt_timestamp(segment.start)} --> "
                    f"{_format_srt_timestamp(segment.end)}",
                    segment.text,
                ]
            )
        )
    srt_content = "\n\n".join(srt_blocks)
    if srt_content:
        srt_content += "\n"
    staged.srt.write_text(srt_content, encoding="utf-8")

    payload = {
        "source": video.name,
        "audio": staged.audio.name,
        "language": result.language,
        "language_probability": result.language_probability,
        "model": MODEL_NAME,
        "device": DEVICE,
        "compute_type": COMPUTE_TYPE,
        "vad_filter": True,
        "segments": [
            {
                "start": segment.start,
                "end": segment.end,
                "text": segment.text,
            }
            for segment in result.segments
        ],
    }
    staged.json.write_text(
        json.dumps(payload, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )


def _staged_paths(staging_dir: Path, final: OutputPaths) -> OutputPaths:
    return OutputPaths(
        directory=staging_dir,
        audio=staging_dir / final.audio.name,
        text=staging_dir / final.text.name,
        srt=staging_dir / final.srt.name,
        json=staging_dir / final.json.name,
    )


def _commit_outputs(staged: OutputPaths, final: OutputPaths) -> None:
    _ensure_output_available(final)
    created_directory = False
    moved: list[Path] = []
    try:
        if not final.directory.exists():
            final.directory.mkdir()
            created_directory = True
        for source, destination in zip(staged.files, final.files, strict=True):
            try:
                source.replace(destination)
            except BaseException:
                if destination.exists() and not source.exists():
                    moved.append(destination)
                raise
            else:
                moved.append(destination)
    except BaseException as exc:
        rollback_failures: list[str] = []
        for path in reversed(moved):
            try:
                path.unlink(missing_ok=True)
            except BaseException as cleanup_exc:
                rollback_failures.append(f"{path}: {cleanup_exc}")
        if created_directory:
            try:
                final.directory.rmdir()
            except BaseException as cleanup_exc:
                rollback_failures.append(f"{final.directory}: {cleanup_exc}")

        if not isinstance(exc, Exception):
            if rollback_failures:
                try:
                    setattr(
                        exc,
                        "_media_commit_rollback_failures",
                        tuple(rollback_failures),
                    )
                except Exception:
                    pass
            raise

        rollback_detail = ""
        if rollback_failures:
            rollback_detail = "；回滚不完整：" + "；".join(rollback_failures)
        raise MediaTranscriptionError(
            f"提交输出文件失败，已回滚已写结果：{exc}{rollback_detail}"
        ) from exc


def transcribe_video(
    video_path: str | Path,
    output_dir: str | Path | None = None,
    language: str | None = None,
    *,
    runner: ProcessRunner | None = None,
    model_loader: ModelLoader | None = None,
) -> tuple[OutputPaths, RuntimeInfo]:
    runner = runner or _run_process
    model_loader = model_loader or _load_cuda_model
    language = language.strip().lower() if language else None

    video = _resolve_video(video_path)
    outputs = _build_output_paths(video, output_dir)
    _ensure_output_available(outputs)
    ffprobe = _require_executable("ffprobe")
    ffmpeg = _require_executable("ffmpeg")
    _probe_first_audio_stream(video, ffprobe, runner)

    outputs.directory.parent.mkdir(parents=True, exist_ok=True)
    staging_dir = Path(
        tempfile.mkdtemp(
            prefix=f".{video.stem}.transcribe-",
            dir=str(outputs.directory.parent),
        )
    )
    staged = _staged_paths(staging_dir, outputs)

    try:
        print("正在提取第一条音轨（16 kHz、单声道 FLAC）……")
        _extract_first_audio_stream(video, staged.audio, ffmpeg, runner)
        audio_duration = _probe_audio_duration(staged.audio, ffprobe, runner)
        print("正在加载 large-v3（CUDA、float16）……")
        model, runtime = model_loader()
        print("正在使用 GPU 和 VAD 转写……")
        if audio_duration <= SHORT_AUDIO_LIMIT_SECONDS:
            transcript = _collect_transcript(model, staged.audio, language)
        else:
            transcript = _collect_chunked_transcript(
                model,
                staged.audio,
                audio_duration,
                language,
                ffmpeg,
                runner,
            )
        _write_transcript_files(video, staged, transcript)
        _commit_outputs(staged, outputs)
    finally:
        shutil.rmtree(staging_dir, ignore_errors=True)

    return outputs, runtime


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="用 FFmpeg 和 faster-whisper GPU 转写单个本地视频。"
    )
    parser.add_argument("video", help="本地视频文件路径")
    parser.add_argument("--output", help="可选输出目录；默认在视频旁创建 <文件名>.transcript")
    parser.add_argument("--language", help="可选语言代码，例如 zh；默认自动检测")
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = _build_parser().parse_args(argv)
    try:
        outputs, runtime = transcribe_video(args.video, args.output, args.language)
    except MediaTranscriptionError as exc:
        print(f"错误：{exc}", file=sys.stderr)
        return 1
    except KeyboardInterrupt as exc:
        rollback_failures = getattr(exc, "_media_commit_rollback_failures", ())
        if rollback_failures:
            detail = "\n".join(f"- {item}" for item in rollback_failures)
            print(
                "错误：用户中断；正式输出回滚不完整，请人工检查以下路径：\n"
                f"{detail}",
                file=sys.stderr,
            )
        else:
            print("错误：用户中断，未提交输出文件。", file=sys.stderr)
        return 130

    print(f"完成：{outputs.directory}")
    print(
        "GPU 自检："
        f"CUDA 设备数={runtime.cuda_device_count}，"
        f"计算类型={','.join(runtime.cuda_compute_types)}"
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())