#!/usr/bin/env python3
"""
Qwen3-ASR-1.7B 智能分段转录 v3
- Silero VAD ONNX 在自然停顿处切分
- 合并小段为 30~60s 大块，边界完整
- 多线程并行转录
- 输出 TXT / SRT / JSON
"""
import argparse
import json
import os
from pathlib import Path
import re
import shutil
import subprocess
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path

sys.stdout = __import__("io").TextIOWrapper(
    sys.stdout.buffer, encoding="utf-8", errors="replace"
)

import numpy as np
import soundfile as sf
import torch
import silero_vad

import requests

MODEL = "ggml-org/Qwen3-ASR-1.7B-GGUF:Q8_0"
PORT = os.environ.get("QWEN3_ASR_PORT", "18003")
API_URL = f"http://localhost:{PORT}/v1/audio/transcriptions"
REPORT_DIR = Path(os.path.dirname(__file__)) / "reports"

_http = requests.Session()
_http.headers.update({"Accept": "application/json"})


@dataclass
class Segment:
    start: float
    end: float

@dataclass
class Chunk:
    segments: list[Segment]
    text: str = ""

    @property
    def start(self) -> float:
        return self.segments[0].start

    @property
    def end(self) -> float:
        return self.segments[-1].end

    @property
    def audio_duration(self) -> float:
        return sum(s.end - s.start for s in self.segments)

    @property
    def wall_duration(self) -> float:
        return self.end - self.start


def ffmpeg_run(args: list[str], check=True):
    subprocess.run(
        ["ffmpeg", "-y", "-hide_banner", "-loglevel", "error"] + args,
        capture_output=True, check=check
    )


def extract_wav(src: str, dst: str, start: float, duration: float):
    ffmpeg_run([
        "-i", src, "-ss", str(start), "-t", str(duration),
        "-vn", "-acodec", "pcm_s16le", "-ar", "16000", "-ac", "1", dst
    ])


def find_speech_segments(audio_path: str) -> list[Segment]:
    """用 Silero VAD 检出所有语音片段"""
    data, sr = sf.read(audio_path)
    audio_t = torch.from_numpy(data).float()
    model = silero_vad.load_silero_vad(onnx=True)
    raw = silero_vad.get_speech_timestamps(
        audio_t, model,
        threshold=0.5,
        min_speech_duration_ms=300,
        min_silence_duration_ms=400,
        return_seconds=True
    )
    return [Segment(s["start"], s["end"]) for s in raw]


def group_chunks(segments: list[Segment],
                 target: float = 45,
                 min_size: float = 20,
                 max_size: float = 75) -> list[Chunk]:
    """将语音片段合并为合适的转录块"""
    chunks = []
    cur: list[Segment] = []
    cur_audio = 0.0

    for seg in segments:
        dur = seg.end - seg.start
        if cur_audio + dur > max_size and cur_audio >= min_size:
            chunks.append(Chunk(cur))
            cur = []
            cur_audio = 0.0
        cur.append(seg)
        cur_audio += dur

    if cur:
        if cur_audio < min_size and chunks:
            chunks[-1].segments.extend(cur)
        else:
            chunks.append(Chunk(cur))

    return chunks


def transcribe_chunk(audio_path: str, chunk: Chunk, tmp_dir: str) -> str:
    wav = os.path.join(tmp_dir, f"c{chunk.start:07.1f}.wav")
    extract_wav(audio_path, wav, chunk.start, chunk.wall_duration)
    with open(wav, "rb") as f:
        r = _http.post(API_URL, data={"model": MODEL}, files={"file": f}, timeout=300)
    text = r.text.strip()
    try:
        parsed = r.json()
        text = parsed.get("text", text)
    except Exception:
        pass
    return text


def strip_prefix(text: str) -> str:
    idx = text.find("<asr_text>")
    return text[idx + 10:].strip() if idx >= 0 else text.strip()


def to_srt(chunks: list[Chunk]) -> str:
    lines = []
    for i, c in enumerate(chunks, 1):
        text = strip_prefix(c.text)
        start_s = format_ts(c.start)
        end_s = format_ts(c.end)
        lines.append(f"{i}")
        lines.append(f"{start_s} --> {end_s}")
        lines.append(text)
        lines.append("")
    return "\n".join(lines)


def format_ts(seconds: float) -> str:
    h = int(seconds // 3600)
    m = int((seconds % 3600) // 60)
    s = seconds % 60
    return f"{h:02d}:{m:02d}:{s:06.3f}".replace(".", ",")


def transcribe_video(video_path: str,
                     target: int = 45,
                     workers: int = 4,
                     output: str = "txt"):
    if not os.path.isfile(video_path):
        print(f"[ERROR] File not found: {video_path}")
        return

    t_start = time.time()
    video_name = Path(video_path).stem
    tmp = os.path.join(os.environ.get("TEMP", "C:\\Temp"), "qwen3_vad")
    os.makedirs(tmp, exist_ok=True)

    # ── 1. 提取完整 WAV ──
    full_wav = os.path.join(tmp, "full.wav")
    print("Extracting audio...", end=" ", flush=True)
    t0 = time.time()
    extract_wav(video_path, full_wav, 0, 999999)
    print(f"{time.time()-t0:.2f}s")

    # ── 2. VAD 检测 ──
    print("VAD detecting speech...", end=" ", flush=True)
    t0 = time.time()
    segs = find_speech_segments(full_wav)
    print(f"{time.time()-t0:.2f}s ({len(segs)} segments)")

    # ── 3. 分组 ──
    chunks = group_chunks(segs, target=target)
    print(f"Grouped into {len(chunks)} chunks")

    # ── 4. 并行转写 ──
    done = [0]
    lock = threading.Lock()
    errs = []

    def process(c: Chunk):
        try:
            text = transcribe_chunk(full_wav, c, tmp)
            with lock:
                c.text = text
        except Exception as e:
            with lock:
                errs.append((c.start, str(e)))
        with lock:
            done[0] += 1
            sys.stdout.write(f"\r  [{done[0]}/{len(chunks)}]")
            sys.stdout.flush()

    t1 = time.time()
    with ThreadPoolExecutor(max_workers=workers) as pool:
        for f in as_completed([pool.submit(process, c) for c in chunks]):
            f.result()

    trans_time = time.time() - t1
    trans_audio = sum(c.audio_duration for c in chunks)

    # ── 5. 合并结果 ──
    full_text = "".join(strip_prefix(c.text) for c in chunks if c.text)

    # ── 6. 报告 ──
    total = time.time() - t_start
    rtf = round(total / trans_audio, 4)
    speed = round(trans_audio / trans_time, 2)

    print(f"\n\nAudio: {trans_audio:.0f}s ({trans_audio/60:.1f} min)")
    print(f"Transcribe: {trans_time:.1f}s  Total: {total:.1f}s")
    print(f"RTF: {rtf}  Speed: {speed}x")
    print(f"\n--- Result ({len(full_text)} chars) ---")
    print(full_text[:600] + "..." if len(full_text) > 600 else full_text)

    # ── 7. 输出文件 ──
    os.makedirs(REPORT_DIR, exist_ok=True)
    now = datetime.now().strftime("%Y%m%d-%H%M%S")
    base = os.path.join(REPORT_DIR, f"{video_name}-{now}")

    if output in ("txt", "all"):
        p = f"{base}.txt"
        with open(p, "w", encoding="utf-8") as f:
            f.write(full_text)
        print(f"TXT:  {p}")

    if output in ("srt", "all"):
        p = f"{base}.srt"
        with open(p, "w", encoding="utf-8") as f:
            f.write(to_srt(chunks))
        print(f"SRT:  {p}")

    if output in ("json", "all"):
        data = {
            "model": MODEL, "source": video_path,
            "duration": trans_audio,
            "rtf": rtf, "speed": speed,
            "chunks": [
                {"start": c.start, "end": c.end,
                 "text": strip_prefix(c.text)}
                for c in chunks if c.text
            ],
            "text": full_text
        }
        p = f"{base}.json"
        with open(p, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        print(f"JSON: {p}")

    # 清理
    try:
        shutil.rmtree(tmp)
    except Exception:
        pass

    return full_text


def main():
    p = argparse.ArgumentParser(description="Qwen3-ASR 智能分段转录 v3")
    p.add_argument("--video", default=None)
    p.add_argument("--target", type=int, default=45,
                   help="目标每段时长(秒), 默认45")
    p.add_argument("--workers", type=int, default=4)
    p.add_argument("--output", choices=["txt", "srt", "json", "all"],
                   default="txt", help="输出格式")
    args = p.parse_args()

    video = args.video or (r"G:\Download\BaiduNetdiskDownload\个人成长--合集\打造AI时代的终身学习力：重构被异化的学习更新中\6-模型预测：模型预测未见-1080P 高清-AVC.mp4"
    )
    transcribe_video(video, args.target, args.workers, args.output)


if __name__ == "__main__":
    main()