client_file.py

import os
import io
import shutil
import subprocess
import asyncio
import json
import argparse
import sys
import websockets
import numpy as np
import soundfile as sf
import webrtcvad
from math import gcd
from scipy.signal import resample_poly
from datetime import datetime, timedelta
from pypdf import PdfReader
from tqdm import tqdm

# Audio settings
SAMPLING_RATE = 16000
CHANNELS = 1
MIN_SEGMENT_DURATION_S = 1.5  # skip VAD segments shorter than this (slide transitions, applause pops)
MAX_SEGMENT_DURATION_S = 30.0  # force-split segments longer than this at the lowest-energy point
CONTAMINATION_WINDOW = 6  # consecutive words that must appear verbatim in context to flag contamination


def is_url(s: str) -> bool:
    return s.startswith(("http://", "https://"))


def detect_browser_for_cookies() -> str | None:
    """Pick the first browser likely to have a Bilibili session, by checking common profile dirs."""
    home = os.path.expanduser("~")
    candidates = [
        ("firefox", os.path.join(home, ".mozilla", "firefox")),
        ("chrome", os.path.join(home, ".config", "google-chrome")),
        ("chromium", os.path.join(home, ".config", "chromium")),
        ("edge", os.path.join(home, ".config", "microsoft-edge")),
        ("brave", os.path.join(home, ".config", "BraveSoftware", "Brave-Browser")),
        # macOS paths
        ("firefox", os.path.join(home, "Library", "Application Support", "Firefox")),
        ("chrome", os.path.join(home, "Library", "Application Support", "Google", "Chrome")),
    ]
    for name, path in candidates:
        if os.path.isdir(path):
            return name
    return None


def download_url(url: str, output_dir: str = "downloads", cookies_from_browser: str | None = None) -> str:
    """
    Download a video/audio URL via yt-dlp (works for Bilibili, YouTube, etc.).
    Returns the path to the downloaded media file. Reuses an existing download if the
    output template already produced a file for this URL.
    """
    if shutil.which("yt-dlp") is None:
        raise RuntimeError("yt-dlp is not installed. Install with: pip install yt-dlp")

    os.makedirs(output_dir, exist_ok=True)

    # Use the video id as the stem so we can detect cached downloads
    id_proc = subprocess.run(
        ["yt-dlp", "--get-id", "--no-warnings", url],
        capture_output=True,
        text=True,
    )
    if id_proc.returncode != 0 or not id_proc.stdout.strip():
        raise RuntimeError(f"yt-dlp could not resolve URL: {id_proc.stderr.strip() or url}")
    video_id = id_proc.stdout.strip().splitlines()[-1]

    # Look for an existing file with that id
    for fname in os.listdir(output_dir):
        if fname.startswith(video_id + ".") and not fname.endswith(".part"):
            cached = os.path.join(output_dir, fname)
            print(f"Reusing cached download: {cached}")
            return cached

    output_template = os.path.join(output_dir, "%(id)s.%(ext)s")
    cmd = [
        "yt-dlp",
        "-f", "bestaudio/best",
        "--no-playlist",
        "-o", output_template,
        url,
    ]
    if cookies_from_browser is None:
        cookies_from_browser = detect_browser_for_cookies()
    if cookies_from_browser:
        print(f"Using cookies from {cookies_from_browser}")
        cmd[1:1] = ["--cookies-from-browser", cookies_from_browser]

    print(f"Downloading via yt-dlp: {url}")
    proc = subprocess.run(cmd, capture_output=True, text=True)
    if proc.returncode != 0:
        raise RuntimeError(f"yt-dlp failed:\n{proc.stderr}")

    # Locate the produced file (extension chosen by yt-dlp)
    for fname in os.listdir(output_dir):
        if fname.startswith(video_id + ".") and not fname.endswith(".part"):
            downloaded = os.path.join(output_dir, fname)
            print(f"Downloaded: {downloaded}")
            return downloaded

    raise FileNotFoundError(f"yt-dlp succeeded but no output file found for id={video_id} in {output_dir}")


def extract_context(file_path: str, max_chars: int = 1000) -> str:
    """Extract text from a PDF or Markdown file up to max_chars to use as vocabulary context."""
    ext = os.path.splitext(file_path)[1].lower()
    if ext == ".pdf":
        reader = PdfReader(file_path)
        text = ""
        for page in reader.pages:
            page_text = page.extract_text()
            if page_text:
                text += page_text + "\n"
            if len(text) > max_chars:
                text = text[:max_chars]
                break
    elif ext in (".md", ".markdown"):
        with open(file_path, "r", encoding="utf-8") as f:
            text = f.read(max_chars)
    else:
        raise ValueError(f"Unsupported context file type: {ext!r}. Expected .pdf or .md/.markdown.")

    instruction = "The following is background reference text to help with vocabulary and spelling. Do NOT read or transcribe this text. Only transcribe the audio. Reference: "
    final_context = instruction + text.replace("\n", " ")
    print(f"Extracted {len(text)} chars from {file_path} as reference context.")
    return final_context


def extract_vocals(audio_path: str, device: str = "cuda", separated_dir: str | None = None) -> str:
    """
    Run demucs (htdemucs) to isolate the vocal stem.
    Tracks are saved persistently under separated_dir (default: ./separated/ next to the audio
    file) so subsequent runs reuse the existing output and skip separation.
    Defaults to CPU to avoid competing with the ASR server's GPU memory.
    Returns vocals_wav_path.
    """
    if separated_dir is None:
        separated_dir = os.path.join(os.path.dirname(os.path.abspath(audio_path)), "separated")

    stem = os.path.splitext(os.path.basename(audio_path))[0]
    vocals_path = os.path.join(separated_dir, "htdemucs", stem, "vocals.wav")

    if os.path.exists(vocals_path):
        print(f"Reusing cached vocals: {vocals_path}")
        return vocals_path

    os.makedirs(separated_dir, exist_ok=True)
    print(f"Extracting vocals with demucs (htdemucs) on {device} — this may take a few minutes...")
    env = os.environ.copy()
    try:
        runner = os.path.join(os.path.dirname(os.path.abspath(__file__)), "demucs_runner.py")
        subprocess.run(
            [
                sys.executable,
                runner,
                "--two-stems",
                "vocals",
                "-d",
                device,
                "-o",
                separated_dir,
                audio_path,
            ],
            check=True,
            env=env,
        )
    except subprocess.CalledProcessError as e:
        raise RuntimeError(f"Demucs failed: {e}") from e

    if not os.path.exists(vocals_path):
        raise FileNotFoundError(f"Demucs did not produce expected output at: {vocals_path}\n" f"Contents of {separated_dir}: {os.listdir(separated_dir)}")

    print(f"Vocals saved to: {vocals_path}")
    return vocals_path


def resample_to_16k(audio_int16: np.ndarray, src_sr: int) -> np.ndarray:
    """Resample int16 audio from src_sr to 16000 Hz using polyphase filtering."""
    if src_sr == SAMPLING_RATE:
        return audio_int16
    g = gcd(src_sr, SAMPLING_RATE)
    up = SAMPLING_RATE // g
    down = src_sr // g
    resampled = resample_poly(audio_int16.astype(np.float64), up, down)
    return np.clip(resampled, -32768, 32767).astype(np.int16)


def collapse_repetitions(text: str, min_ngram: int = 2, max_ngram: int = 8) -> str:
    """
    Detect and collapse ASR hallucination loops — e.g. "ski skills, ski skills, ski skills…"
    For each n-gram size (from largest to smallest), find runs of the same n-gram repeated
    3+ times consecutively and reduce them to a single occurrence.
    Returns the cleaned text, or an empty string if the result is mostly repetition (>60% removed).
    """
    words = text.split()
    if len(words) < min_ngram * 3:
        return text

    changed = True
    while changed:
        changed = False
        for n in range(max_ngram, min_ngram - 1, -1):
            i = 0
            new_words = []
            while i < len(words):
                # Check whether the n-gram starting at i repeats at i+n, i+2n, ...
                ngram = words[i : i + n]
                if len(ngram) < n:
                    new_words.extend(words[i:])
                    i = len(words)
                    break
                repeat_count = 1
                j = i + n
                while words[j : j + n] == ngram:
                    repeat_count += 1
                    j += n
                if repeat_count >= 3:
                    new_words.extend(ngram)  # keep one copy
                    i = j
                    changed = True
                else:
                    new_words.append(words[i])
                    i += 1
            words = new_words

    result = " ".join(words)
    # If we stripped more than 60% of the original words, the segment was mostly hallucination
    if len(words) < len(text.split()) * 0.4:
        return result  # still return what's left — caller can decide
    return result


def is_near_duplicate(text: str, prev: str, threshold: float = 0.7) -> bool:
    """
    Returns True if text is a near-duplicate of prev — same segment re-sent by the server.
    Uses word-level Jaccard similarity on bigrams so minor wording differences still match.
    """
    if not prev or not text:
        return False

    def bigrams(s):
        w = s.lower().split()
        return set(zip(w, w[1:])) if len(w) > 1 else set()

    a, b = bigrams(text), bigrams(prev)
    if not a and not b:
        return text.lower().strip() == prev.lower().strip()
    if not a or not b:
        return False
    return len(a & b) / len(a | b) >= threshold


def is_context_contamination(text: str, context: str) -> bool:
    """
    Returns True if CONTAMINATION_WINDOW or more consecutive words from text appear
    verbatim in context, meaning the model generated from the context document rather
    than from actual audio.
    """
    if not context or not text:
        return False
    words = text.lower().split()
    if len(words) < CONTAMINATION_WINDOW:
        return False
    context_lower = context.lower()
    for i in range(len(words) - CONTAMINATION_WINDOW + 1):
        phrase = " ".join(words[i : i + CONTAMINATION_WINDOW])
        if phrase in context_lower:
            return True
    return False


def parse_time_str(time_str: str) -> timedelta:
    parts = time_str.split(":")
    if len(parts) == 3:
        h, m, s = map(int, parts)
    elif len(parts) == 2:
        h, m = map(int, parts)
        s = 0
    else:
        h, m, s = int(parts[0]), 0, 0
    return timedelta(hours=h, minutes=m, seconds=s)


def int16_to_bytes(audio_int16: np.ndarray) -> bytes:
    return audio_int16.tobytes()


async def stream_audio_segment(ws, audio_int16: np.ndarray, sample_rate: int, context: str = ""):
    # Handshake — include context if provided
    start_msg: dict = {
        "type": "start",
        "format": "pcm_s16le",
        "sample_rate_hz": sample_rate,
        "channels": CHANNELS,
    }
    if context:
        start_msg["context"] = context
    await ws.send(json.dumps(start_msg))

    # 100ms chunks to send over websocket
    chunk_size = int(sample_rate * 0.1)

    for i in range(0, len(audio_int16), chunk_size):
        chunk = audio_int16[i : i + chunk_size]
        await ws.send(int16_to_bytes(chunk))
        # Small sleep to yield control and simulate streaming / avoid overwhelming server buffers
        await asyncio.sleep(0.01)

    await ws.send(json.dumps({"type": "stop"}))


async def receiver(ws) -> str:
    """Receives websocket messages and returns the final transcribed text."""
    final_text = ""
    last_partial = ""
    try:
        async for message in ws:
            try:
                timestamp = datetime.now().strftime("%H:%M:%S.%f")[:-3]
                evt = json.loads(message)
                msg_type = evt.get("type")
                text = evt.get("text", "")
                lang = evt.get("language", "")

                if msg_type == "ready":
                    pass
                elif msg_type == "partial":
                    sys.stdout.write(f"\r[Partial] ({lang}): {text}")
                    sys.stdout.flush()
                    last_partial = text
                elif msg_type == "final":
                    sys.stdout.write(f"\r[Final] ({lang}): {text}\n")
                    sys.stdout.flush()
                    final_text = text
                    last_partial = ""
                elif msg_type == "error":
                    print(f"\n[{timestamp}] [Error]: {evt.get('message')}")
                elif msg_type == "info":
                    pass

            except json.JSONDecodeError:
                print(f"\n[Raw]: {message}")
    except Exception:
        pass

    # Fall back to last partial if connection closed before "final" was received
    return final_text or last_partial


def apply_vad(audio_int16: np.ndarray, sample_rate: int, min_silence_duration_s: float = 1.0) -> list[tuple[int, int]]:
    """
    Applies WebRTC VAD to find non-silent segments separated by at least `min_silence_duration_s` of silence.
    Aggressiveness=2 rejects more non-speech (music, noise) than the previous value of 1.
    Returns a list of (start_idx, end_idx) tuples.
    """
    vad = webrtcvad.Vad(2)  # 0–3; 2 = balanced, filters music/noise better than 1

    frame_duration_ms = 30  # WebRTC VAD requires 10, 20, or 30 ms frames
    frame_size = int(sample_rate * (frame_duration_ms / 1000.0))
    min_silence_frames = int((min_silence_duration_s * 1000.0) / frame_duration_ms)

    segments = []
    current_start = -1
    silence_counter = 0

    for i in range(0, len(audio_int16) - frame_size, frame_size):
        frame = audio_int16[i : i + frame_size]
        if len(frame) < frame_size:
            break

        is_speech = vad.is_speech(frame.tobytes(), sample_rate)

        if is_speech:
            silence_counter = 0
            if current_start == -1:
                current_start = max(0, i - frame_size * 5)  # Include a bit of lead-in
        else:
            if current_start != -1:
                silence_counter += 1
                if silence_counter > min_silence_frames:
                    current_end = min(len(audio_int16), i + frame_size * 5)  # Include a bit of lead-out
                    segments.append((current_start, current_end))
                    current_start = -1
                    silence_counter = 0

    if current_start != -1:
        segments.append((current_start, len(audio_int16)))

    return segments


def split_long_segment(audio_int16: np.ndarray, sample_rate: int, max_duration_s: float) -> list[tuple[int, int]]:
    """
    Recursively split a segment that exceeds max_duration_s by cutting at the lowest-energy
    point within a search window around the midpoint. Returns a list of (start, end) sample indices
    relative to the start of the input array.
    """
    total_samples = len(audio_int16)
    max_samples = int(max_duration_s * sample_rate)

    if total_samples <= max_samples:
        return [(0, total_samples)]

    # Search for the lowest-energy point in a ±10% window around the midpoint
    mid = total_samples // 2
    search_half = max(1, total_samples // 10)
    search_start = max(0, mid - search_half)
    search_end = min(total_samples, mid + search_half)

    # Use 30ms frames to find lowest RMS energy
    frame_size = int(sample_rate * 0.03)
    best_idx = mid
    best_energy = float("inf")
    for i in range(search_start, search_end - frame_size, frame_size):
        frame = audio_int16[i : i + frame_size].astype(np.float32)
        energy = float(np.mean(frame ** 2))
        if energy < best_energy:
            best_energy = energy
            best_idx = i + frame_size // 2

    left = split_long_segment(audio_int16[:best_idx], sample_rate, max_duration_s)
    right = split_long_segment(audio_int16[best_idx:], sample_rate, max_duration_s)
    return left + [(best_idx + s, best_idx + e) for s, e in right]


def read_audio(file_path: str) -> tuple[np.ndarray, int]:
    """Read audio file, using ffmpeg for formats soundfile can't handle (e.g. mp3, m4a)."""
    try:
        audio, sr = sf.read(file_path, dtype="int16")
        return audio, sr
    except Exception:
        proc = subprocess.run(
            ["ffmpeg", "-y", "-i", file_path, "-f", "wav", "-"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        if proc.returncode != 0:
            raise RuntimeError(f"FFmpeg failed: {proc.stderr.decode(errors='ignore')}")
        audio, sr = sf.read(io.BytesIO(proc.stdout), dtype="int16")
        return audio, sr


async def process_file(
    file_path: str,
    endpoint: str,
    context: str = "",
    output: str | None = None,
    vocal_extraction: bool = False,
    demucs_device: str = "cuda",
    separated_dir: str | None = None,
    offset: timedelta = timedelta(),
):
    # ── Step 1: Vocal extraction (optional) ───────────────────────────────────
    audio_source = file_path
    if vocal_extraction:
        audio_source = extract_vocals(file_path, device=demucs_device, separated_dir=separated_dir)

    # ── Step 2: Load audio ────────────────────────────────────────────────
    if True:
        print(f"Loading audio: {audio_source}...")
        audio_data, sr = read_audio(audio_source)

        # Downmix to mono (demucs outputs stereo)
        if len(audio_data.shape) > 1:
            audio_data = audio_data[:, 0]

        # ── Step 3: Resample to 16kHz ─────────────────────────────────────────
        if sr != SAMPLING_RATE:
            print(f"Resampling from {sr} Hz → {SAMPLING_RATE} Hz...")
            audio_data = resample_to_16k(audio_data, sr)
            sr = SAMPLING_RATE

        duration_s = len(audio_data) / sr
        print(f"Audio ready: {duration_s:.1f}s @ {sr}Hz mono")

        # ── Step 4: VAD segmentation ──────────────────────────────────────────
        print("Applying VAD to split into speech segments...")
        segments = apply_vad(audio_data, SAMPLING_RATE, min_silence_duration_s=0.8)

        if not segments:
            print("No speech detected after vocal extraction + VAD.")
            return

        print(f"Found {len(segments)} speech segments.")

        # ── Step 5: Transcribe each segment ──────────────────────────────────
        stem = os.path.splitext(file_path)[0]
        txt_output_file = output if output else f"{stem}.txt"
        print(f"Results will be written line by line to {txt_output_file}")

        with open(txt_output_file, "w", encoding="utf-8") as f:
            pass  # truncate

        pbar = tqdm(total=len(segments), desc="Transcribing")
        last_text = ""

        for start_idx, end_idx in segments:
            segment_duration = (end_idx - start_idx) / SAMPLING_RATE
            start_time = timedelta(seconds=start_idx / SAMPLING_RATE)
            end_time = timedelta(seconds=end_idx / SAMPLING_RATE)
            timestamp_str = f"[{start_time} - {end_time}]"

            if segment_duration < MIN_SEGMENT_DURATION_S:
                tqdm.write(f"Skipping {timestamp_str} (too short: {segment_duration:.1f}s)")
                pbar.update(1)
                continue

            # Force-split oversized segments at lowest-energy points
            segment_audio_full = audio_data[start_idx:end_idx]
            sub_splits = split_long_segment(segment_audio_full, SAMPLING_RATE, MAX_SEGMENT_DURATION_S)
            if len(sub_splits) > 1:
                tqdm.write(f"\nSplitting {timestamp_str} ({segment_duration:.0f}s) into {len(sub_splits)} sub-segments...")

            for sub_start, sub_end in sub_splits:
                abs_start = start_idx + sub_start
                abs_end = start_idx + sub_end
                sub_start_time = timedelta(seconds=abs_start / SAMPLING_RATE)
                sub_timestamp_str = f"[{timedelta(seconds=abs_start/SAMPLING_RATE)} - {timedelta(seconds=abs_end/SAMPLING_RATE)}]"
                tqdm.write(f"\nProcessing {sub_timestamp_str}...")
                segment_audio = audio_data[abs_start:abs_end]

                try:
                    async with websockets.connect(endpoint, max_size=None, ping_interval=None) as ws:
                        sender_task = asyncio.create_task(stream_audio_segment(ws, segment_audio, SAMPLING_RATE, context=context))
                        receiver_task = asyncio.create_task(receiver(ws))

                        _, text = await asyncio.gather(sender_task, receiver_task)

                        text = text.strip()
                        if not text:
                            pass
                        elif is_context_contamination(text, context):
                            tqdm.write(f"Skipping {sub_timestamp_str} (context contamination detected)")
                        elif is_near_duplicate(text, last_text):
                            tqdm.write(f"Skipping {sub_timestamp_str} (near-duplicate of previous segment)")
                        else:
                            cleaned = collapse_repetitions(text)
                            if cleaned != text:
                                tqdm.write(f"Collapsed repetitions in {sub_timestamp_str}")
                            ts_str = str((sub_start_time + offset)).split(".")[0]
                            with open(txt_output_file, "a", encoding="utf-8") as f:
                                f.write(f"[{ts_str}] {cleaned}\n")
                            last_text = cleaned

                except Exception as e:
                    tqdm.write(f"Failed to process segment {sub_timestamp_str}: {e}")

            pbar.update(1)

        pbar.close()
        print(f"\n--- Final transcription written to {txt_output_file} ---")


async def main():
    parser = argparse.ArgumentParser(description="Qwen3-ASR File Streaming Transcriber with Vocal Extraction, VAD & Context (outputs TXT)", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("audio", help="Path to input audio file (.mp3, .wav, .m4a, ...) or URL (Bilibili, YouTube, etc.)")
    parser.add_argument("-e", "--endpoint", default="ws://localhost:9002/transcribe-streaming", help="WebSocket Endpoint URL")
    parser.add_argument("-l", "--language", default="English", help="Forced language full name (e.g. English, Chinese, Japanese)")
    parser.add_argument("-o", "--output", default=None, help="Output TXT file path (default: <audio_stem>.txt)")
    parser.add_argument("--context", default=None, help="Path to PDF or Markdown reference document for vocabulary context")
    parser.add_argument(
        "--vocal-extraction",
        action="store_true",
        help="Run demucs vocal extraction before ASR. Useful for event recordings with background music. Separated tracks are cached in --separated-dir for reuse.",
    )
    parser.add_argument("--demucs-device", default="cuda", help="Device for demucs: cuda (default) or cpu / cuda:N.")
    parser.add_argument("--separated-dir", default=None, help="Where to store/reuse demucs output (default: separated/ next to the audio file).")
    parser.add_argument("--offset", default=None, help="Start time offset added to all timestamps. Format: hh:mm:ss, hh:mm (hours:minutes), or hh (hours). E.g. 18:00 = 18 hours, 1:30:00 = 1.5 hours.")
    parser.add_argument("--download-dir", default="downloads", help="Where to store downloaded files when the input is a URL.")
    parser.add_argument("--cookies-from-browser", default=None, help="Override browser used for cookies (firefox, chrome, chromium, edge, brave). Auto-detected if omitted.")
    args = parser.parse_args()

    # Resolve URL input to a local file before the rest of the pipeline
    audio_path = args.audio
    if is_url(audio_path):
        audio_path = download_url(audio_path, output_dir=args.download_dir, cookies_from_browser=args.cookies_from_browser)

    endpoint = args.endpoint
    if args.language:
        sep = "&" if "?" in endpoint else "?"
        endpoint += f"{sep}language={args.language}"

    context = ""
    if args.context:
        context = extract_context(args.context)

    offset = parse_time_str(args.offset) if args.offset else timedelta()

    await process_file(
        audio_path,
        endpoint,
        context=context,
        output=args.output,
        vocal_extraction=args.vocal_extraction,
        demucs_device=args.demucs_device,
        separated_dir=args.separated_dir,
        offset=offset,
    )


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\nStopped by user.")