from __future__ import annotations

import argparse
import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any


REGULAR_MODEL_CONFIGS: dict[str, dict[str, Any]] = {
    "sensevoice": {
        "model": "iic/SenseVoiceSmall",
        "vad_model": "fsmn-vad",
        "vad_kwargs": {"max_single_segment_time": 30000},
    },
    "paraformer": {
        "model": "paraformer-zh",
        "vad_model": "fsmn-vad",
        "punc_model": "ct-punc",
    },
    "paraformer-en": {
        "model": "paraformer-en",
        "vad_model": "fsmn-vad",
    },
}

NANO_MODEL_CONFIG: dict[str, Any] = {
    "model": "FunAudioLLM/Fun-ASR-Nano-2512",
    "vad_model": "fsmn-vad",
}


@dataclass
class AudioTrack:
    index: int
    codec: str | None
    channels: int | None
    channel_layout: str | None
    language: str | None
    title: str | None


def require_command(name: str) -> None:
    if shutil.which(name) is None:
        raise SystemExit(f"缺少命令: {name}")


def run_command(cmd: list[str], *, capture: bool = True) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        cmd,
        check=True,
        text=True,
        stdout=subprocess.PIPE if capture else None,
        stderr=subprocess.PIPE if capture else None,
    )


def list_audio_tracks(media_path: Path) -> list[AudioTrack]:
    require_command("ffprobe")
    proc = run_command(
        [
            "ffprobe",
            "-v",
            "error",
            "-select_streams",
            "a",
            "-show_entries",
            "stream=index,codec_name,channels,channel_layout:stream_tags=language,title",
            "-of",
            "json",
            str(media_path),
        ]
    )
    data = json.loads(proc.stdout)
    tracks: list[AudioTrack] = []
    for stream in data.get("streams", []):
        tags = stream.get("tags") or {}
        tracks.append(
            AudioTrack(
                index=int(stream["index"]),
                codec=stream.get("codec_name"),
                channels=stream.get("channels"),
                channel_layout=stream.get("channel_layout"),
                language=tags.get("language"),
                title=tags.get("title"),
            )
        )
    return tracks


def print_audio_tracks(media_path: Path) -> None:
    tracks = list_audio_tracks(media_path)
    if not tracks:
        print("未找到音频轨")
        return
    for i, track in enumerate(tracks, 1):
        parts = [
            f"#{i}",
            f"stream_index={track.index}",
            f"codec={track.codec or '?'}",
            f"channels={track.channels or '?'}",
        ]
        if track.channel_layout:
            parts.append(f"layout={track.channel_layout}")
        if track.language:
            parts.append(f"lang={track.language}")
        if track.title:
            parts.append(f"title={track.title}")
        print("  ".join(parts))


def resolve_track(media_path: Path, requested: int | None) -> int:
    tracks = list_audio_tracks(media_path)
    if not tracks:
        raise SystemExit(f"未找到音频轨: {media_path}")
    if requested is None:
        return tracks[0].index
    valid_indexes = {track.index for track in tracks}
    if requested in valid_indexes:
        return requested
    raise SystemExit(
        f"音轨 stream index 不存在: {requested}\n"
        f"可用音轨: {', '.join(str(t.index) for t in tracks)}\n"
        "提示: --track 使用 ffprobe/ffmpeg 的 stream index，不是第几条音轨的序号。"
    )


def extract_audio(media_path: Path, wav_path: Path, track_index: int) -> None:
    require_command("ffmpeg")
    wav_path.parent.mkdir(parents=True, exist_ok=True)
    cmd = [
        "ffmpeg",
        "-hide_banner",
        "-y",
        "-i",
        str(media_path),
        "-map",
        f"0:{track_index}",
        "-vn",
        "-ar",
        "16000",
        "-ac",
        "1",
        "-c:a",
        "pcm_s16le",
        str(wav_path),
    ]
    subprocess.run(cmd, check=True)


def clean_text(text: str, *, sensevoice: bool) -> str:
    if sensevoice:
        try:
            from funasr.utils.postprocess_utils import rich_transcription_postprocess

            text = rich_transcription_postprocess(text)
        except Exception:
            # Keep transcription usable even if FunASR changes this helper.
            pass
    return re.sub(r"<\|[^|]*\|>", "", text).strip()


def get_audio_duration(audio_path: Path) -> float | None:
    try:
        import soundfile as sf

        return round(float(sf.info(str(audio_path)).duration), 3)
    except Exception:
        return None


def normalize_result(
    raw_result: list[dict[str, Any]],
    *,
    audio_path: Path,
    source_media: Path,
    track_index: int,
    model_name: str,
    model_id: str,
    language: str | None,
    diarize: bool,
    elapsed: float,
    sensevoice: bool,
) -> dict[str, Any]:
    first = raw_result[0] if raw_result else {}
    text = clean_text(str(first.get("text", "")), sensevoice=sensevoice)
    segments: list[dict[str, Any]] = []
    for seg in first.get("sentence_info") or []:
        item: dict[str, Any] = {
            "start": seg.get("start", 0),
            "end": seg.get("end", 0),
            "text": clean_text(str(seg.get("sentence") or seg.get("text") or ""), sensevoice=sensevoice),
        }
        if diarize and "spk" in seg:
            item["speaker"] = seg["spk"]
        segments.append(item)

    output: dict[str, Any] = {
        "text": text,
        "segments": segments,
        "file": source_media.name,
        "audio_file": audio_path.name,
        "track": track_index,
        "model": model_name,
        "model_id": model_id,
        "language": language or "auto",
        "audio_duration_s": get_audio_duration(audio_path),
        "processing_s": round(elapsed, 3),
    }
    for key in ("timestamps", "timestamp"):
        if key in first:
            output[key] = first[key]
    return output


def write_json(path: Path, data: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")


def make_output_path(output_dir: Path, media_path: Path, track_index: int, suffix: str) -> Path:
    return output_dir / f"{media_path.stem}_track{track_index}_{suffix}.json"


def resolve_output_dir(output_dir: Path | None, media_path: Path) -> Path:
    if output_dir is None:
        return media_path.parent
    return output_dir.expanduser().resolve()


def add_common_args(parser: argparse.ArgumentParser) -> None:
    parser.add_argument("media", type=Path, help="视频/音频文件路径")
    parser.add_argument("--track", type=int, default=None, help="ffmpeg stream index；默认第一条音频轨")
    parser.add_argument("--list-tracks", action="store_true", help="列出音轨后退出")
    parser.add_argument("--output-dir", "-o", type=Path, default=None, help="输出目录；默认源视频/音频文件所在目录")
    parser.add_argument("--language", "-l", default="auto", help="语言，如 auto/zh/en/ja/ko/yue；默认 auto")
    parser.add_argument("--device", default=None, help="设备，如 cuda:0/cpu；默认自动")
    parser.add_argument("--no-diarize", action="store_true", help="禁用 cam++ 说话人分离")
    parser.add_argument("--keep-wav", action="store_true", help="保留提取出的 16k wav 文件")
    parser.add_argument("--wav-dir", type=Path, default=None, help="临时 wav 输出目录；默认系统临时目录")
    parser.add_argument("--verbose", "-v", action="store_true", help="输出更多进度信息")


def prepare_audio(args: argparse.Namespace) -> tuple[Path, int, tempfile.TemporaryDirectory[str] | None]:
    media_path: Path = args.media.expanduser().resolve()
    if not media_path.exists():
        raise SystemExit(f"文件不存在: {media_path}")

    if args.list_tracks:
        print_audio_tracks(media_path)
        raise SystemExit(0)

    track_index = resolve_track(media_path, args.track)
    temp_dir: tempfile.TemporaryDirectory[str] | None = None
    if args.keep_wav:
        wav_dir = args.wav_dir.expanduser().resolve() if args.wav_dir else resolve_output_dir(args.output_dir, media_path)
        wav_dir.mkdir(parents=True, exist_ok=True)
    elif args.wav_dir is not None:
        wav_dir = args.wav_dir.expanduser().resolve()
        wav_dir.mkdir(parents=True, exist_ok=True)
    else:
        temp_dir = tempfile.TemporaryDirectory(prefix="funasr-script-")
        wav_dir = Path(temp_dir.name)

    wav_path = wav_dir / f"{media_path.stem}_track{track_index}.wav"
    print(f"[1/2] 提取音轨 stream_index={track_index} -> {wav_path}", file=sys.stderr)
    extract_audio(media_path, wav_path, track_index)
    return wav_path, track_index, temp_dir


def auto_device(device: str | None) -> str:
    if device:
        return device
    import torch

    return "cuda:0" if torch.cuda.is_available() else "cpu"


def transcribe_with_config(
    *,
    wav_path: Path,
    media_path: Path,
    track_index: int,
    output_path: Path,
    model_name: str,
    config: dict[str, Any],
    language: str | None,
    device: str | None,
    diarize: bool,
    use_itn: bool,
    sensevoice: bool,
    verbose: bool,
) -> dict[str, Any]:
    from funasr import AutoModel

    config = config.copy()
    if diarize and "spk_model" not in config:
        config["spk_model"] = "cam++"
    resolved_device = auto_device(device)

    print(f"[2/2] 加载模型 {model_name} ({config['model']}) on {resolved_device}", file=sys.stderr)
    load_start = time.time()
    model = AutoModel(device=resolved_device, disable_update=True, **config)
    if verbose:
        print(f"模型加载耗时: {time.time() - load_start:.1f}s", file=sys.stderr)

    gen_kw: dict[str, Any] = {"input": str(wav_path), "batch_size": 1}
    if language and language != "auto":
        gen_kw["language"] = language
    if use_itn:
        gen_kw["use_itn"] = True

    start = time.time()
    raw_result = model.generate(**gen_kw)
    elapsed = time.time() - start
    output = normalize_result(
        raw_result,
        audio_path=wav_path,
        source_media=media_path,
        track_index=track_index,
        model_name=model_name,
        model_id=str(config["model"]),
        language=language,
        diarize=diarize,
        elapsed=elapsed,
        sensevoice=sensevoice,
    )
    write_json(output_path, output)
    print(f"完成: {output_path}", file=sys.stderr)
    print(f"转录耗时: {elapsed:.2f}s", file=sys.stderr)
    return output


def run_pipeline(args: argparse.Namespace, *, model_name: str, config: dict[str, Any], suffix: str, use_itn: bool, sensevoice: bool) -> Path:
    media_path = args.media.expanduser().resolve()
    temp_dir: tempfile.TemporaryDirectory[str] | None = None
    try:
        wav_path, track_index, temp_dir = prepare_audio(args)
        output_path = make_output_path(resolve_output_dir(args.output_dir, media_path), media_path, track_index, suffix)
        transcribe_with_config(
            wav_path=wav_path,
            media_path=media_path,
            track_index=track_index,
            output_path=output_path,
            model_name=model_name,
            config=config,
            language=args.language,
            device=args.device,
            diarize=not args.no_diarize,
            use_itn=use_itn,
            sensevoice=sensevoice,
            verbose=args.verbose,
        )
        print(output_path)
        return output_path
    finally:
        if temp_dir is not None:
            temp_dir.cleanup()