     1→from __future__ import annotations
     2→
     3→import argparse
     4→import io
     5→import json
     6→import math
     7→import os
     8→import random
     9→import tarfile
    10→import threading
    11→import time
    12→import zlib
    13→from collections import defaultdict
    14→from contextlib import contextmanager
    15→from dataclasses import dataclass
    16→from pathlib import Path
    17→from typing import Any, Optional
    18→
    19→import boto3
    20→import numpy as np
    21→import psycopg
    22→import soundfile as sf
    23→from boto3.s3.transfer import TransferConfig
    24→from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, as_completed, wait
    25→from scipy import signal
    26→from tqdm import tqdm
    27→
    28→try:
    29→    import pyloudnorm as pyln
    30→except Exception:  # pragma: no cover - optional dependency
    31→    pyln = None
    32→
    33→ENV_PATH = Path(".env")
    34→LANGUAGES = [
    35→    "telugu",
    36→    "malayalam",
    37→    "english",
    38→    "hindi",
    39→    "punjabi",
    40→    "tamil",
    41→    "kannada",
    42→    "gujarati",
    43→    "bengali",
    44→    "odia",
    45→    "marathi",
    46→    "assamese",
    47→]
    48→
    49→# Global timing stats (thread-safe)
    50→_timing_lock = threading.Lock()
    51→_timing_stats: dict[str, list[float]] = defaultdict(list)
    52→_timing_enabled = False
    53→
    54→
    55→@contextmanager
    56→def timer(name: str):
    57→    """Context manager for timing code blocks."""
    58→    if not _timing_enabled:
    59→        yield
    60→        return
    61→    start = time.perf_counter()
    62→    try:
    63→        yield
    64→    finally:
    65→        elapsed = time.perf_counter() - start
    66→        with _timing_lock:
    67→            _timing_stats[name].append(elapsed)
    68→
    69→
    70→def print_timing_stats():
    71→    """Print aggregated timing statistics."""
    72→    if not _timing_stats:
    73→        return
    74→    print("\n=== Timing Statistics ===")
    75→    total_time = sum(sum(times) for times in _timing_stats.values())
    76→    for name, times in sorted(_timing_stats.items()):
    77→        count = len(times)
    78→        total = sum(times)
    79→        avg = total / count if count > 0 else 0
    80→        percent = (total / total_time * 100) if total_time > 0 else 0
    81→        print(f"{name:35s}: {count:6d} calls, {total:8.2f}s total, {avg*1000:8.2f}ms avg, {percent:5.1f}%")
    82→    print(f"{'TOTAL':35s}:              {total_time:8.2f}s")
    83→    print("=" * 80)
    84→
    85→
    86→_window_cache: dict[int, np.ndarray] = {}
    87→_freq_cache: dict[tuple[int, int], np.ndarray] = {}
    88→_cache_lock = threading.Lock()
    89→_lufs_local = threading.local()
    90→
    91→
    92→def get_hann_window(length: int) -> np.ndarray:
    93→    window = _window_cache.get(length)
    94→    if window is not None:
    95→        return window
    96→    with _cache_lock:
    97→        window = _window_cache.get(length)
    98→        if window is None:
    99→            window = np.hanning(length)
   100→            _window_cache[length] = window
   101→    return window
   102→
   103→
   104→def get_rfftfreq(length: int, sr: int) -> np.ndarray:
   105→    key = (length, sr)
   106→    freqs = _freq_cache.get(key)
   107→    if freqs is not None:
   108→        return freqs
   109→    with _cache_lock:
   110→        freqs = _freq_cache.get(key)
   111→        if freqs is None:
   112→            freqs = np.fft.rfftfreq(length, d=1.0 / sr)
   113→            _freq_cache[key] = freqs
   114→    return freqs
   115→
   116→
   117→def get_lufs_meter(sr: int) -> Optional[Any]:
   118→    if pyln is None:
   119→        return None
   120→    meters = getattr(_lufs_local, "meters", None)
   121→    if meters is None:
   122→        meters = {}
   123→        _lufs_local.meters = meters
   124→    meter = meters.get(sr)
   125→    if meter is None:
   126→        meter = pyln.Meter(sr)
   127→        meters[sr] = meter
   128→    return meter
   129→
   130→
   131→@dataclass(frozen=True)
   132→class VideoRecord:
   133→    youtube_id: str
   134→    r2_tar_key: str
   135→
   136→
   137→@dataclass(frozen=True)
   138→class SegmentOutput:
   139→    sample_rate: int
   140→    audio_bytes: bytes
   141→    metadata: dict[str, Any]
   142→
   143→
   144→@dataclass(frozen=True)
   145→class ChunkConfig:
   146→    enabled: bool
   147→    min_sec: float
   148→    max_sec: float
   149→    threshold_sec: float
   150→    all_buckets: bool
   151→
   152→
   153→@dataclass(frozen=True)
   154→class SilenceConfig:
   155→    enabled: bool
   156→    frame_sec: float
   157→    hop_sec: float
   158→    search_sec: float
   159→    percentile: float
   160→    db_drop: float
   161→
   162→
   163→def load_env(env_path: Path) -> dict[str, str]:
   164→    env: dict[str, str] = {}
   165→    for line in env_path.read_text().splitlines():
   166→        stripped = line.strip()
   167→        if not stripped or stripped.startswith("#") or "=" not in stripped:
   168→            continue
   169→        if stripped.startswith("export "):
   170→            stripped = stripped[len("export ") :]
   171→        key, value = stripped.split("=", 1)
   172→        env[key.strip()] = value.strip()
   173→    return env
   174→
   175→
   176→def make_r2_client(env: dict[str, str]):
   177→    endpoint = env.get("R2_ENDPOINT_URL")
   178→    access_key = env.get("R2_ACCESS_KEY_ID")
   179→    secret_key = env.get("R2_SECRET_ACCESS_KEY")
   180→    if not endpoint or not access_key or not secret_key:
   181→        raise SystemExit("Missing R2 credentials in .env")
   182→    return boto3.client(
   183→        "s3",
   184→        endpoint_url=endpoint,
   185→        aws_access_key_id=access_key,
   186→        aws_secret_access_key=secret_key,
   187→    )
   188→
   189→
   190→def resolve_r2_key(
   191→    s3, bucket: str, r2_key: str, prefix: Optional[str]
   192→) -> Optional[str]:
   193→    if not r2_key:
   194→        return None
   195→    key = r2_key.lstrip("/")
   196→    candidates = [key]
   197→
   198→    if prefix:
   199→        cleaned_prefix = prefix.strip("/")
   200→        if not key.startswith(cleaned_prefix):
   201→            candidates.append(f"{cleaned_prefix}/{key}")
   202→
   203→    base = os.path.basename(key)
   204→    if base != key:
   205→        candidates.append(base)
   206→        if prefix:
   207→            candidates.append(f"{prefix.strip('/')}/{base}")
   208→
   209→    for candidate in candidates:
   210→        try:
   211→            with timer("r2_head_object"):
   212→                s3.head_object(Bucket=bucket, Key=candidate)
   213→            return candidate
   214→        except Exception:
   215→            continue
   216→    return None
   217→
   218→
   219→def parse_source_url(source_url: Optional[str], default_bucket: Optional[str]) -> Optional[tuple[str, str]]:
   220→    if not source_url or not default_bucket:
   221→        return None
   222→    cleaned = source_url.lstrip("/")
   223→    if cleaned.startswith(default_bucket + "/"):
   224→        return default_bucket, cleaned[len(default_bucket) + 1 :]
   225→    return default_bucket, cleaned
   226→
   227→
   228→def fetch_video_records(
   229→    db_url: str,
   230→    language: str,
   231→    limit: Optional[int],
   232→    partition_count: int = 1,
   233→    partition_index: int = 0,
   234→) -> list[VideoRecord]:
   235→    records: list[VideoRecord] = []
   236→    with psycopg.connect(db_url) as conn:
   237→        with conn.cursor() as cur:
   238→            query = """
   239→                select youtube_id, r2_tar_key
   240→                from videos
   241→                where lower(language) = %s
   242→                  and r2_tar_key is not null
   243→                  and status = 'COMPLETED'
   244→            """
   245→            params = [language.lower()]
   246→            if partition_count > 1:
   247→                query += " and mod(abs(hashtext(youtube_id)), %s) = %s"
   248→                params.extend([partition_count, partition_index])
   249→            query += " order by random()"
   250→            if limit:
   251→                query += " limit %s"
   252→                params.append(limit)
   253→            cur.execute(query, params)
   254→            for youtube_id, r2_tar_key in cur.fetchall():
   255→                if youtube_id and r2_tar_key:
   256→                    records.append(VideoRecord(youtube_id, r2_tar_key))
   257→    return records
   258→
   259→
   260→def normalize(value: float, min_val: float, max_val: float) -> float:
   261→    if max_val <= min_val:
   262→        return 0.0
   263→    return max(0.0, min(1.0, (value - min_val) / (max_val - min_val)))
   264→
   265→
   266→def compute_spectral_metrics(audio: np.ndarray, sr: int) -> tuple[float, float, float]:
   267→    if audio.size == 0:
   268→        return 0.0, 0.0, 0.0
   269→    with timer("cpu_spectral_compute"):
   270→        max_samples = min(audio.size, sr * 5)
   271→        window = get_hann_window(max_samples)
   272→        chunk = audio[:max_samples] * window
   273→        spec = np.fft.rfft(chunk)
   274→        power = np.abs(spec) ** 2
   275→        freqs = get_rfftfreq(max_samples, sr)
   276→        total_energy = float(np.sum(power)) + 1e-12
   277→        cumulative = np.cumsum(power)
   278→        rolloff_idx = int(np.searchsorted(cumulative, 0.85 * cumulative[-1]))
   279→        rolloff_hz = float(freqs[min(rolloff_idx, len(freqs) - 1)])
   280→
   281→        power_db = 10.0 * np.log10(power + 1e-12)
   282→        threshold = power_db.max() - 60.0
   283→        above = np.where(power_db >= threshold)[0]
   284→        effective_nyquist_hz = float(freqs[above[-1]]) if above.size > 0 else 0.0
   285→
   286→        high_freq_energy = float(np.sum(power[freqs >= 8000.0]))
   287→        high_freq_ratio = high_freq_energy / total_energy
   288→    return rolloff_hz, effective_nyquist_hz, high_freq_ratio
   289→
   290→
   291→def classify_native_bucket(effective_nyquist_hz: float) -> str:
   292→    """Classify audio bandwidth based on effective Nyquist frequency.
   293→
   294→    These buckets represent actual audio bandwidth, not container sample rates.
   295→    Thresholds are based on observed distribution in the dataset.
   296→    """
   297→    if effective_nyquist_hz < 8000:
   298→        return "<16k"
   299→    if effective_nyquist_hz < 12000:
   300→        return "16-24k"
   301→    if effective_nyquist_hz < 16000:
   302→        return "24-32k"
   303→    if effective_nyquist_hz < 22050:
   304→        return "32-44k"
   305→    return ">=44k"
   306→
   307→
   308→def compute_audio_metrics(audio: np.ndarray, sr: int) -> dict[str, float]:
   309→    with timer("audio_metrics_basic"):
   310→        audio = audio.astype(np.float32, copy=False)
   311→        rms = float(np.sqrt(np.mean(audio ** 2) + 1e-12))
   312→        loudness_rms_db = 20.0 * math.log10(rms + 1e-12)
   313→        clipping_fraction = float(np.mean(np.abs(audio) > 0.99))
   314→        dc_offset = float(np.mean(audio))
   315→        noise_floor = float(np.percentile(np.abs(audio), 10))
   316→        snr_proxy = 20.0 * math.log10((rms + 1e-9) / (noise_floor + 1e-9))
   317→
   318→    with timer("audio_metrics_lufs"):
   319→        loudness_lufs = -100.0
   320→        if pyln is not None:
   321→            try:
   322→                meter = get_lufs_meter(sr)
   323→                if meter is not None:
   324→                    loudness_lufs = float(meter.integrated_loudness(audio))
   325→            except Exception:
   326→                loudness_lufs = -100.0
   327→
   328→    rolloff_hz, effective_nyquist_hz, high_freq_ratio = compute_spectral_metrics(audio, sr)
   329→
   330→    return {
   331→        "loudness_lufs": loudness_lufs,
   332→        "loudness_rms_db": loudness_rms_db,
   333→        "clipping_fraction": clipping_fraction,
   334→        "dc_offset": dc_offset,
   335→        "snr_proxy": snr_proxy,
   336→        "spectral_rolloff_hz": rolloff_hz,
   337→        "effective_nyquist_hz": effective_nyquist_hz,
   338→        "high_freq_energy_ratio": high_freq_ratio,
   339→    }
   340→
   341→
   342→def compute_quality_score(metrics: dict[str, float], has_music: bool) -> float:
   343→    score = (
   344→        0.3 * (1.0 - metrics["clipping_fraction"])
   345→        + 0.2 * normalize(metrics["snr_proxy"], 10, 40)
   346→        + 0.2 * (1.0 if not has_music else 0.5)
   347→        + 0.2 * normalize(metrics["loudness_lufs"], -40, -14)
   348→        + 0.1 * (1.0 if abs(metrics["dc_offset"]) < 0.01 else 0.5)
   349→    )
   350→    return max(0.0, min(1.0, score))
   351→
   352→
   353→def resample_audio(
   354→    audio: np.ndarray,
   355→    src_sr: int,
   356→    dst_sr: int,
   357→) -> np.ndarray:
   358→    if src_sr == dst_sr:
   359→        return audio
   360→    with timer("cpu_resample"):
   361→        gcd = math.gcd(src_sr, dst_sr)
   362→        up = dst_sr // gcd
   363→        down = src_sr // gcd
   364→        return signal.resample_poly(audio, up, down).astype(np.float32, copy=False)
   365→
   366→
   367→def encode_audio(audio: np.ndarray, sr: int, output_format: str) -> bytes:
   368→    with timer("audio_encode"):
   369→        buf = io.BytesIO()
   370→        if output_format == "wav":
   371→            sf.write(buf, audio, sr, format="WAV", subtype="PCM_16")
   372→        else:
   373→            sf.write(buf, audio, sr, format="FLAC", subtype="PCM_16")
   374→        return buf.getvalue()
   375→
   376→
   377→def duration_bucket(duration_sec: float) -> Optional[tuple[str, float]]:
   378→    if duration_sec < 3.0:
   379→        return None
   380→    if duration_sec < 6.0:
   381→        return "3-6s", 0.3
   382→    if duration_sec < 12.0:
   383→        return "6-12s", 1.0
   384→    if duration_sec < 24.0:
   385→        return "12-24s", 0.8
   386→    if duration_sec <= 60.0:
   387→        return "24-60s", 0.5
   388→    return None
   389→
   390→
   391→def compute_frame_rms(
   392→    audio: np.ndarray, sr: int, frame_sec: float, hop_sec: float
   393→) -> tuple[np.ndarray, np.ndarray]:
   394→    frame_len = max(1, int(round(frame_sec * sr)))
   395→    hop_len = max(1, int(round(hop_sec * sr)))
   396→    if audio.size < frame_len:
   397→        return np.array([], dtype=np.float32), np.array([], dtype=np.float32)
   398→
   399→    if not audio.flags["C_CONTIGUOUS"]:
   400→        audio = np.ascontiguousarray(audio)
   401→
   402→    num_frames = 1 + (audio.size - frame_len) // hop_len
   403→    if num_frames <= 0:
   404→        return np.array([], dtype=np.float32), np.array([], dtype=np.float32)
   405→
   406→    stride = audio.strides[0]
   407→    frames = np.lib.stride_tricks.as_strided(
   408→        audio,
   409→        shape=(num_frames, frame_len),
   410→        strides=(hop_len * stride, stride),
   411→    )
   412→    rms = np.sqrt(np.mean(frames * frames, axis=1) + 1e-12)
   413→    times = (np.arange(num_frames, dtype=np.float32) * hop_len + (frame_len / 2)) / sr
   414→    return times.astype(np.float32), rms.astype(np.float32)
   415→
   416→
   417→def find_silence_boundary(
   418→    times: np.ndarray,
   419→    rms: np.ndarray,
   420→    search_start: float,
   421→    search_end: float,
   422→    threshold_rms: float,
   423→    median_db: float,
   424→    db_drop: float,
   425→) -> Optional[float]:
   426→    if times.size == 0:
   427→        return None
   428→    mask = (times >= search_start) & (times <= search_end)
   429→    if not np.any(mask):
   430→        return None
   431→    subset = rms[mask]
   432→    min_idx = int(np.argmin(subset))
   433→    min_rms = float(subset[min_idx])
   434→    min_time = float(times[mask][min_idx])
   435→    min_db = 20.0 * math.log10(min_rms + 1e-12)
   436→    if min_rms <= threshold_rms or (median_db - min_db) >= db_drop:
   437→        return min_time
   438→    return None
   439→
   440→
   441→def make_chunk_windows(
   442→    audio: np.ndarray,
   443→    sr: int,
   444→    cfg: ChunkConfig,
   445→    silence_cfg: SilenceConfig,
   446→    rng: random.Random,
   447→) -> list[tuple[float, float]]:
   448→    duration_sec = float(audio.shape[0] / sr)
   449→    if not cfg.enabled or duration_sec < cfg.threshold_sec:
   450→        return [(0.0, duration_sec)]
   451→
   452→    if not silence_cfg.enabled:
   453→        windows: list[tuple[float, float]] = []
   454→        offset = 0.0
   455→        while duration_sec - offset >= cfg.min_sec:
   456→            remaining = duration_sec - offset
   457→            if remaining <= cfg.max_sec:
   458→                length = remaining
   459→            else:
   460→                length = rng.uniform(cfg.min_sec, cfg.max_sec)
   461→                if remaining - length < cfg.min_sec:
   462→                    length = remaining
   463→            windows.append((offset, length))
   464→            offset += length
   465→        return windows
   466→
   467→    times, rms = compute_frame_rms(
   468→        audio, sr, silence_cfg.frame_sec, silence_cfg.hop_sec
   469→    )
   470→    if times.size == 0:
   471→        return [(0.0, duration_sec)]
   472→
   473→    threshold_rms = float(np.quantile(rms, silence_cfg.percentile / 100.0))
   474→    median_db = 20.0 * math.log10(float(np.median(rms)) + 1e-12)
   475→
   476→    windows: list[tuple[float, float]] = []
   477→    offset = 0.0
   478→    while duration_sec - offset > cfg.max_sec:
   479→        remaining = duration_sec - offset
   480→        latest_cut = duration_sec - cfg.min_sec
   481→        if offset + cfg.min_sec >= latest_cut:
   482→            break
   483→        target_len = rng.uniform(cfg.min_sec, cfg.max_sec)
   484→        if remaining - target_len < cfg.min_sec:
   485→            target_len = remaining - cfg.min_sec
   486→        target_end = offset + target_len
   487→        search_low = max(offset + cfg.min_sec, target_end - silence_cfg.search_sec)
   488→        search_high = min(offset + cfg.max_sec, target_end + silence_cfg.search_sec, latest_cut)
   489→        if search_high <= search_low:
   490→            search_low = offset + cfg.min_sec
   491→            search_high = min(offset + cfg.max_sec, latest_cut)
   492→        boundary = find_silence_boundary(
   493→            times,
   494→            rms,
   495→            search_low,
   496→            search_high,
   497→            threshold_rms,
   498→            median_db,
   499→            silence_cfg.db_drop,
   500→        )
   501→        if boundary is None:
   502→            if not windows:
   503→                return [(0.0, duration_sec)]
   504→            windows.append((offset, duration_sec - offset))
   505→            return windows
   506→        windows.append((offset, boundary - offset))
   507→        offset = boundary
   508→
   509→    windows.append((offset, duration_sec - offset))
   510→    return windows
   511→
   512→
   513→def make_segment_seed(base_seed: int, video_id: str, speaker: str, index: int) -> int:
   514→    key = f"{video_id}-{speaker}-{index}-{base_seed}"
   515→    return zlib.adler32(key.encode("utf-8")) & 0xFFFFFFFF
   516→
   517→
   518→def parse_segment_filename(filename: str) -> Optional[tuple[str, float]]:
   519→    """Parse segment filename to extract speaker and start time.
   520→
   521→    Filename format: segments/SPEAKER_XX_NNNN_start-end.flac
   522→    Example: segments/SPEAKER_00_0007_35.10-37.49.flac
   523→
   524→    Returns (speaker, start_time) for matching against metadata.
   525→    """
   526→    base = filename.split("/")[-1]
   527→    if not base.endswith(".flac"):
   528→        return None
   529→
   530→    # Remove .flac extension
   531→    base = base[:-5]
   532→
   533→    # Try to extract time range (last part after final underscore)
   534→    parts = base.split("_")
   535→    if len(parts) < 3:
   536→        return None
   537→
   538→    # Speaker is everything except last two parts (index and time range)
   539→    speaker = "_".join(parts[:-2])
   540→
   541→    # Time range is the last part: "35.10-37.49"
   542→    time_range = parts[-1]
   543→    if "-" not in time_range:
   544→        return None
   545→
   546→    try:
   547→        start_str = time_range.split("-")[0]
   548→        start_time = float(start_str)
   549→    except (ValueError, IndexError):
   550→        return None
   551→
   552→    return speaker, start_time
   553→
   554→
   555→def process_segment(
   556→    video_id: str,
   557→    language: str,
   558→    segment_bytes: bytes,
   559→    segment_meta: dict[str, Any],
   560→    speaker: str,
   561→    index: int,
   562→    chunk_cfg: ChunkConfig,
   563→    silence_cfg: SilenceConfig,
   564→    rng_seed: int,
   565→    output_format: str,
   566→) -> tuple[list[SegmentOutput], dict[str, int]]:
   567→    with timer("segment_total"):
   568→        stats = {
   569→            "dropped_too_short": 0,
   570→            "dropped_too_long": 0,
   571→            "dropped_clipped": 0,
   572→            "chunks_emitted": 0,
   573→            "segments_chunked": 0,
   574→        }
   575→        outputs: list[SegmentOutput] = []
   576→
   577→        with timer("flac_decode"):
   578→            try:
   579→                audio, sr = sf.read(io.BytesIO(segment_bytes), dtype="float32")
   580→            except Exception:
   581→                return outputs, stats
   582→
   583→        if audio.ndim > 1:
   584→            audio = np.mean(audio, axis=1)
   585→
   586→        rng = random.Random(rng_seed)
   587→        windows = make_chunk_windows(audio, sr, chunk_cfg, silence_cfg, rng)
   588→
   589→        music_stats = segment_meta.get("music_stats") or {}
   590→        has_music = bool(music_stats.get("has_music", False))
   591→
   592→        def emit_outputs(
   593→            audio_slice: np.ndarray,
   594→            metrics: dict[str, float],
   595→            sample_id: str,
   596→            segment_index: int,
   597→            chunk_index: Optional[int],
   598→            start: Optional[float],
   599→            end: Optional[float],
   600→            original_start: Optional[float],
   601→            original_end: Optional[float],
   602→            target_srs: list[int],
   603→        ) -> None:
   604→            slice_duration = float(audio_slice.shape[0] / sr)
   605→            bucket = duration_bucket(slice_duration)
   606→            if bucket is None:
   607→                if slice_duration < 3.0:
   608→                    stats["dropped_too_short"] += 1
   609→                else:
   610→                    stats["dropped_too_long"] += 1
   611→                return
   612→            if metrics["clipping_fraction"] > 0.1:
   613→                stats["dropped_clipped"] += 1
   614→                return
   615→
   616→            duration_bucket_name, base_weight = bucket
   617→            sampling_weight = base_weight * (0.5 if has_music else 1.0)
   618→            quality_score = compute_quality_score(metrics, has_music)
   619→            effective_nyquist = metrics["effective_nyquist_hz"]
   620→
   621→            base_metadata = {
   622→                "id": sample_id,
   623→                "video_id": video_id,
   624→                "language": language,
   625→                "segment_index": segment_index,
   626→                "chunk_index": chunk_index,
   627→                "start": start,
   628→                "end": end,
   629→                "original_start": original_start,
   630→                "original_end": original_end,
   631→                "duration_sec": slice_duration,
   632→                "native_sr": sr,
   633→                "effective_nyquist_hz": effective_nyquist,
   634→                "spectral_rolloff_hz": metrics["spectral_rolloff_hz"],
   635→                "high_freq_energy_ratio": metrics["high_freq_energy_ratio"],
   636→                "is_truly_wideband": bool(effective_nyquist >= 12000),
   637→                "native_sr_bucket": classify_native_bucket(effective_nyquist),
   638→                "loudness_lufs": metrics["loudness_lufs"],
   639→                "loudness_rms_db": metrics["loudness_rms_db"],
   640→                "clipping_fraction": metrics["clipping_fraction"],
   641→                "snr_proxy": metrics["snr_proxy"],
   642→                "quality_score": quality_score,
   643→                "music_prob": float(music_stats.get("music_mean", 0.0)),
   644→                "music_stats": music_stats,
   645→                "speaker": speaker,
   646→                # VAD regions: Full duration placeholder; random 6s crops happen during training.
   647→                "vad_speech_regions": [[0.0, slice_duration]],
   648→                "duration_bucket": duration_bucket_name,
   649→                "sampling_weight": sampling_weight,
   650→            }
   651→
   652→            for target_sr in target_srs:
   653→                resampled = resample_audio(audio_slice, sr, target_sr)
   654→                resampled = np.clip(resampled, -1.0, 1.0)
   655→                audio_bytes = encode_audio(resampled, target_sr, output_format)
   656→                meta = dict(base_metadata)
   657→                meta["sample_rate"] = target_sr
   658→                meta["resample_target_sr"] = target_sr
   659→                outputs.append(SegmentOutput(target_sr, audio_bytes, meta))
   660→
   661→            if target_srs == [16000]:
   662→                stats["chunks_emitted"] += 1
   663→
   664→        sample_id_base = f"{video_id}__{speaker}_{index:04d}"
   665→        use_chunk_suffix = len(windows) > 1
   666→        if use_chunk_suffix:
   667→            stats["segments_chunked"] += 1
   668→
   669→        segment_start = segment_meta.get("start")
   670→        segment_end = segment_meta.get("end")
   671→        original_start = segment_meta.get("original_start")
   672→        original_end = segment_meta.get("original_end")
   673→
   674→        # Always emit 16k outputs; for long segments, emit chunked windows to boost sample count.
   675→        for chunk_idx, (offset_sec, chunk_len_sec) in enumerate(windows):
   676→            start_idx = int(round(offset_sec * sr))
   677→            end_idx = int(round((offset_sec + chunk_len_sec) * sr))
   678→            audio_slice = audio[start_idx:end_idx]
   679→            if audio_slice.size == 0:
   680→                continue
   681→            slice_duration = float(audio_slice.shape[0] / sr)
   682→            slice_metrics = compute_audio_metrics(audio_slice, sr)
   683→
   684→            chunk_start = float(segment_start) + offset_sec if segment_start is not None else None
   685→            chunk_end = chunk_start + slice_duration if chunk_start is not None else None
   686→            chunk_original_start = (
   687→                float(original_start) + offset_sec if original_start is not None else None
   688→            )
   689→            chunk_original_end = (
   690→                chunk_original_start + slice_duration
   691→                if chunk_original_start is not None
   692→                else None
   693→            )
   694→
   695→            sample_id = sample_id_base
   696→            chunk_index = None
   697→            if use_chunk_suffix:
   698→                sample_id = f"{sample_id_base}_c{chunk_idx:02d}"
   699→                chunk_index = chunk_idx
   700→
   701→            emit_outputs(
   702→                audio_slice,
   703→                slice_metrics,
   704→                sample_id,
   705→                index,
   706→                chunk_index,
   707→                chunk_start,
   708→                chunk_end,
   709→                chunk_original_start,
   710→                chunk_original_end,
   711→                [16000],
   712→            )
   713→
   714→        # Optional high-rate buckets: use full segment unless configured otherwise.
   715→        full_metrics = compute_audio_metrics(audio, sr)
   716→        effective_nyquist = full_metrics["effective_nyquist_hz"]
   717→        targets_high: list[int] = []
   718→        if effective_nyquist >= 8000:
   719→            targets_high.append(24000)
   720→        if effective_nyquist >= 12000:
   721→            targets_high.append(32000)
   722→        if effective_nyquist >= 16000:
   723→            targets_high.append(44100)
   724→
   725→        if targets_high:
   726→            if chunk_cfg.all_buckets and use_chunk_suffix:
   727→                for chunk_idx, (offset_sec, chunk_len_sec) in enumerate(windows):
   728→                    start_idx = int(round(offset_sec * sr))
   729→                    end_idx = int(round((offset_sec + chunk_len_sec) * sr))
   730→                    audio_slice = audio[start_idx:end_idx]
   731→                    if audio_slice.size == 0:
   732→                        continue
   733→                    slice_duration = float(audio_slice.shape[0] / sr)
   734→                    slice_metrics = compute_audio_metrics(audio_slice, sr)
   735→                    slice_effective_nyquist = slice_metrics["effective_nyquist_hz"]
   736→                    slice_targets: list[int] = []
   737→                    if slice_effective_nyquist >= 8000:
   738→                        slice_targets.append(24000)
   739→                    if slice_effective_nyquist >= 12000:
   740→                        slice_targets.append(32000)
   741→                    if slice_effective_nyquist >= 16000:
   742→                        slice_targets.append(44100)
   743→                    if not slice_targets:
   744→                        continue
   745→
   746→                    chunk_start = (
   747→                        float(segment_start) + offset_sec if segment_start is not None else None
   748→                    )
   749→                    chunk_end = (
   750→                        chunk_start + slice_duration if chunk_start is not None else None
   751→                    )
   752→                    chunk_original_start = (
   753→                        float(original_start) + offset_sec if original_start is not None else None
   754→                    )
   755→                    chunk_original_end = (
   756→                        chunk_original_start + slice_duration
   757→                        if chunk_original_start is not None
   758→                        else None
   759→                    )
   760→
   761→                    sample_id = f"{sample_id_base}_c{chunk_idx:02d}"
   762→                    emit_outputs(
   763→                        audio_slice,
   764→                        slice_metrics,
   765→                        sample_id,
   766→                        index,
   767→                        chunk_idx,
   768→                        chunk_start,
   769→                        chunk_end,
   770→                        chunk_original_start,
   771→                        chunk_original_end,
   772→                        slice_targets,
   773→                    )
   774→            else:
   775→                emit_outputs(
   776→                    audio,
   777→                    full_metrics,
   778→                    sample_id_base,
   779→                    index,
   780→                    None,
   781→                    segment_start,
   782→                    segment_end,
   783→                    original_start,
   784→                    original_end,
   785→                    targets_high,
   786→                )
   787→
   788→        return outputs, stats
   789→
   790→
   791→class ShardWriter:
   792→    def __init__(
   793→        self, output_dir: Path, prefix: str, target_size_bytes: int, audio_ext: str
   794→    ) -> None:
   795→        self.output_dir = output_dir
   796→        self.prefix = prefix
   797→        self.target_size_bytes = target_size_bytes
   798→        self.audio_ext = audio_ext
   799→        self.shard_index = 0
   800→        self.current_tar: Optional[tarfile.TarFile] = None
   801→        self.current_jsonl = None
   802→        self.current_size = 0
   803→        self.current_count = 0
   804→        self.current_name = ""
   805→        self.output_dir.mkdir(parents=True, exist_ok=True)
   806→
   807→    def _open_new_shard(self) -> None:
   808→        if self.current_tar:
   809→            self.close()
   810→        shard_name = f"{self.prefix}-{self.shard_index:06d}"
   811→        tar_path = self.output_dir / f"{shard_name}.tar"
   812→        jsonl_path = self.output_dir / f"{shard_name}.jsonl"
   813→        self.current_tar = tarfile.open(tar_path, "w")
   814→        self.current_jsonl = jsonl_path.open("w", encoding="utf-8")
   815→        self.current_size = 0
   816→        self.current_count = 0
   817→        self.current_name = shard_name
   818→        self.shard_index += 1
   819→
   820→    def add_sample(self, sample_id: str, audio_bytes: bytes, metadata: dict[str, Any]) -> str:
   821→        if self.current_tar is None:
   822→            self._open_new_shard()
   823→        if (
   824→            self.current_count > 0
   825→            and self.current_size + len(audio_bytes) > self.target_size_bytes
   826→        ):
   827→            self._open_new_shard()
   828→
   829→        assert self.current_tar is not None
   830→        assert self.current_jsonl is not None
   831→
   832→        arcname = f"{sample_id}.{self.audio_ext}"
   833→        info = tarfile.TarInfo(name=arcname)
   834→        info.size = len(audio_bytes)
   835→        info.mtime = int(time.time())
   836→        self.current_tar.addfile(info, io.BytesIO(audio_bytes))
   837→        self.current_jsonl.write(json.dumps(metadata) + "\n")
   838→
   839→        self.current_size += len(audio_bytes)
   840→        self.current_count += 1
   841→        return self.current_name
   842→
   843→    def close(self) -> None:
   844→        if self.current_tar:
   845→            self.current_tar.close()
   846→        if self.current_jsonl:
   847→            self.current_jsonl.close()
   848→        self.current_tar = None
   849→        self.current_jsonl = None
   850→
   851→
   852→def create_prefix_placeholders(
   853→    s3,
   854→    bucket: str,
   855→    languages: list[str],
   856→    rates: list[int],
   857→    extra_prefixes: list[str],
   858→    subprefix: Optional[str],
   859→) -> None:
   860→    prefixes = []
   861→    for lang in languages:
   862→        for rate in rates:
   863→            base = f"{lang}_{rate}"
   864→            prefixes.append(base)
   865→            if subprefix:
   866→                prefixes.append(f"{base}/{subprefix.strip('/')}")
   867→    prefixes.extend(extra_prefixes)
   868→    if subprefix:
   869→        prefixes.extend([f"{p}/{subprefix.strip('/')}" for p in extra_prefixes])
   870→
   871→    for prefix in prefixes:
   872→        key = f"{prefix.strip('/')}/.keep"
   873→        try:
   874→            s3.put_object(Bucket=bucket, Key=key, Body=b"")
   875→        except Exception:
   876→            continue
   877→
   878→
   879→def upload_directory(
   880→    s3,
   881→    bucket: str,
   882→    local_dir: Path,
   883→    remote_prefix: str,
   884→    max_concurrency: int,
   885→) -> list[str]:
   886→    uploaded = []
   887→    config = TransferConfig(
   888→        multipart_threshold=8 * 1024 * 1024,
   889→        multipart_chunksize=64 * 1024 * 1024,
   890→        max_concurrency=max_concurrency,
   891→        use_threads=True,
   892→    )
   893→    for path in sorted(local_dir.glob("*")):
   894→        if not path.is_file():
   895→            continue
   896→        key = f"{remote_prefix}/{path.name}"
   897→        s3.upload_file(str(path), bucket, key, Config=config)
   898→        uploaded.append(key)
   899→    return uploaded
   900→
   901→
   902→def verify_uploaded_shard(
   903→    s3,
   904→    bucket: str,
   905→    remote_prefix: str,
   906→    shard_name: str,
   907→    sample_rate: int,
   908→    audio_ext: str,
   909→    tmp_dir: Path,
   910→) -> dict[str, Any]:
   911→    tar_key = f"{remote_prefix}/{shard_name}.tar"
   912→    jsonl_key = f"{remote_prefix}/{shard_name}.jsonl"
   913→    tar_path = tmp_dir / f"{shard_name}.tar"
   914→    jsonl_path = tmp_dir / f"{shard_name}.jsonl"
   915→
   916→    s3.download_file(bucket, tar_key, str(tar_path))
   917→    s3.download_file(bucket, jsonl_key, str(jsonl_path))
   918→
   919→    with tarfile.open(tar_path, "r") as tf:
   920→        tar_names = set(tf.getnames())
   921→        with jsonl_path.open("r", encoding="utf-8") as fh:
   922→            lines = [json.loads(line) for line in fh if line.strip()]
   923→
   924→        missing = 0
   925→        checked = 0
   926→        for entry in lines[:10]:
   927→            checked += 1
   928→            audio_name = f"{entry['id']}.{audio_ext}"
   929→            if audio_name not in tar_names:
   930→                missing += 1
   931→                continue
   932→            member = tf.getmember(audio_name)
   933→            audio_bytes = tf.extractfile(member).read()
   934→            audio, sr = sf.read(io.BytesIO(audio_bytes))
   935→            if sr != sample_rate:
   936→                missing += 1
   937→
   938→    return {
   939→        "checked_samples": checked,
   940→        "missing_or_mismatch": missing,
   941→        "total_entries": len(lines),
   942→    }
   943→
   944→
   945→def parse_args() -> argparse.Namespace:
   946→    parser = argparse.ArgumentParser(description="Prepare sharded datasets from cleaned tars.")
   947→    parser.add_argument("--language", required=True)
   948→    parser.add_argument("--limit", type=int, default=None)
   949→    parser.add_argument("--output-dir", type=Path, default=Path("processed_data"))
   950→    parser.add_argument("--tmp-dir", type=Path, default=Path("/tmp/xcodec_work"))
   951→    parser.add_argument("--target-shard-size-mb", type=int, default=5120)
   952→    parser.add_argument("--download-workers", type=int, default=16)
   953→    parser.add_argument("--segment-workers", type=int, default=8)
   954→    parser.add_argument(
   955→        "--segment-max-inflight",
   956→        type=int,
   957→        default=0,
   958→        help="Limit number of in-flight segment tasks (0=auto).",
   959→    )
   960→    parser.add_argument(
   961→        "--skip-downloads",
   962→        action="store_true",
   963→        help="Use only local tar files; do not download from R2.",
   964→    )
   965→    parser.add_argument(
   966→        "--local-tars-dir",
   967→        type=Path,
   968→        default=None,
   969→        help="Directory containing pre-downloaded tar files for this language.",
   970→    )
   971→    parser.add_argument(
   972→        "--partition-count",
   973→        type=int,
   974→        default=1,
   975→        help="Total number of partitions for parallel processing.",
   976→    )
   977→    parser.add_argument(
   978→        "--partition-index",
   979→        type=int,
   980→        default=0,
   981→        help="Partition index for this run (0-based).",
   982→    )
   983→    parser.add_argument("--val-ratio", type=float, default=0.01)
   984→    parser.add_argument("--seed", type=int, default=17)
   985→    parser.add_argument(
   986→        "--chunk-long",
   987→        dest="chunk_long",
   988→        action="store_true",
   989→        default=True,
   990→        help="Split long segments into 6-10s windows for the 16k bucket.",
   991→    )
   992→    parser.add_argument(
   993→        "--no-chunk-long",
   994→        dest="chunk_long",
   995→        action="store_false",
   996→        help="Disable chunking long segments.",
   997→    )
   998→    parser.add_argument("--chunk-min-sec", type=float, default=6.0)
   999→    parser.add_argument("--chunk-max-sec", type=float, default=10.0)
  1000→    parser.add_argument("--chunk-threshold-sec", type=float, default=12.0)
  1001→    parser.add_argument(
  1002→        "--chunk-all-buckets",
  1003→        action="store_true",
  1004→        help="Also chunk high-rate buckets (default: keep full segments).",
  1005→    )
  1006→    parser.add_argument(
  1007→        "--silence-snap",
  1008→        dest="silence_snap",
  1009→        action="store_true",
  1010→        default=True,
  1011→        help="Snap chunk boundaries to low-energy points when possible.",
  1012→    )
  1013→    parser.add_argument(
  1014→        "--no-silence-snap",
  1015→        dest="silence_snap",
  1016→        action="store_false",
  1017→        help="Disable silence-aware chunk boundaries.",
  1018→    )
  1019→    parser.add_argument("--silence-frame-ms", type=float, default=50.0)
  1020→    parser.add_argument("--silence-hop-ms", type=float, default=10.0)
  1021→    parser.add_argument("--silence-search-sec", type=float, default=0.5)
  1022→    parser.add_argument("--silence-percentile", type=float, default=20.0)
  1023→    parser.add_argument("--silence-db-drop", type=float, default=20.0)
  1024→    parser.add_argument("--upload", action="store_true")
  1025→    parser.add_argument("--verify-upload", action="store_true")
  1026→    parser.add_argument("--create-prefixes", action="store_true")
  1027→    parser.add_argument("--r2-subprefix", default="")
  1028→    parser.add_argument(
  1029→        "--profile-timing",
  1030→        action="store_true",
  1031→        help="Enable detailed timing instrumentation for profiling.",
  1032→    )
  1033→    parser.add_argument(
  1034→        "--output-format",
  1035→        choices=["flac", "wav"],
  1036→        default="flac",
  1037→        help="Output audio format for shards (default: flac).",
  1038→    )
  1039→    return parser.parse_args()
  1040→
  1041→
  1042→def main() -> None:
  1043→    global _timing_enabled
  1044→    args = parse_args()
  1045→    _timing_enabled = args.profile_timing
  1046→    env = load_env(ENV_PATH)
  1047→    db_url = env.get("DATABASE_URL")
  1048→    if not db_url:
  1049→        raise SystemExit("DATABASE_URL not found in .env")
  1050→
  1051→    cleaned_bucket = env.get("R2_ClEANED_BUCKET") or env.get("R2_CLEANED_BUCKET")
  1052→    shard_bucket = env.get("R2_BUCKET_SHARDS")
  1053→    prefix = env.get("R2_PREFIX")
  1054→    if not cleaned_bucket:
  1055→        raise SystemExit("R2 cleaned bucket not found in .env")
  1056→
  1057→    if args.upload and not shard_bucket:
  1058→        raise SystemExit("R2_BUCKET_SHARDS not found in .env")
  1059→
  1060→    if args.chunk_min_sec <= 0 or args.chunk_max_sec <= 0:
  1061→        raise SystemExit("Chunk durations must be positive.")
  1062→    if args.chunk_min_sec > args.chunk_max_sec:
  1063→        raise SystemExit("chunk-min-sec must be <= chunk-max-sec.")
  1064→    if args.chunk_threshold_sec < args.chunk_min_sec:
  1065→        raise SystemExit("chunk-threshold-sec must be >= chunk-min-sec.")
  1066→    if args.silence_frame_ms <= 0 or args.silence_hop_ms <= 0:
  1067→        raise SystemExit("Silence frame/hop must be positive.")
  1068→    if args.silence_search_sec <= 0:
  1069→        raise SystemExit("silence-search-sec must be positive.")
  1070→    if not 0.0 < args.silence_percentile < 100.0:
  1071→        raise SystemExit("silence-percentile must be between 0 and 100.")
  1072→    if args.silence_db_drop <= 0:
  1073→        raise SystemExit("silence-db-drop must be positive.")
  1074→
  1075→    if args.partition_count < 1:
  1076→        raise SystemExit("partition-count must be >= 1.")
  1077→    if not 0 <= args.partition_index < args.partition_count:
  1078→        raise SystemExit("partition-index must be between 0 and partition-count - 1.")
  1079→
  1080→    output_format = args.output_format
  1081→    audio_ext = "wav" if output_format == "wav" else "flac"
  1082→
  1083→    partition_tag = ""
  1084→    if args.partition_count > 1:
  1085→        partition_tag = f"p{args.partition_index:02d}"
  1086→
  1087→    chunk_cfg = ChunkConfig(
  1088→        enabled=args.chunk_long,
  1089→        min_sec=args.chunk_min_sec,
  1090→        max_sec=args.chunk_max_sec,
  1091→        threshold_sec=args.chunk_threshold_sec,
  1092→        all_buckets=args.chunk_all_buckets,
  1093→    )
  1094→    silence_cfg = SilenceConfig(
  1095→        enabled=args.silence_snap,
  1096→        frame_sec=args.silence_frame_ms / 1000.0,
  1097→        hop_sec=args.silence_hop_ms / 1000.0,
  1098→        search_sec=args.silence_search_sec,
  1099→        percentile=args.silence_percentile,
  1100→        db_drop=args.silence_db_drop,
  1101→    )
  1102→    language = args.language.strip().lower()
  1103→    with timer("db_fetch_records"):
  1104→        records = fetch_video_records(
  1105→            db_url,
  1106→            language,
  1107→            args.limit,
  1108→            args.partition_count,
  1109→            args.partition_index,
  1110→        )
  1111→    if not records:
  1112→        raise SystemExit(f"No records found for language={language}")
  1113→
  1114→    s3 = make_r2_client(env)
  1115→    if args.create_prefixes and args.upload:
  1116→        create_prefix_placeholders(
  1117→            s3,
  1118→            shard_bucket,
  1119→            LANGUAGES,
  1120→            [16, 24, 32, 44],
  1121→            ["manifests", "stats", "reports", "checks"],
  1122→            args.r2_subprefix or None,
  1123→        )
  1124→
  1125→    tmp_dir = args.tmp_dir
  1126→    output_dir = args.output_dir
  1127→    if partition_tag:
  1128→        if tmp_dir.name != partition_tag:
  1129→            tmp_dir = tmp_dir / partition_tag
  1130→        if output_dir.name != partition_tag:
  1131→            output_dir = output_dir / partition_tag
  1132→    tmp_dir.mkdir(parents=True, exist_ok=True)
  1133→    download_dir = tmp_dir / "downloads"
  1134→    download_dir.mkdir(parents=True, exist_ok=True)
  1135→
  1136→    config = TransferConfig(
  1137→        multipart_threshold=8 * 1024 * 1024,
  1138→        multipart_chunksize=64 * 1024 * 1024,
  1139→        max_concurrency=max(4, args.download_workers),
  1140→        use_threads=True,
  1141→    )
  1142→
  1143→    def download_record(record: VideoRecord) -> Path:
  1144→        with timer("r2_download_total"):
  1145→            resolved = resolve_r2_key(s3, cleaned_bucket, record.r2_tar_key, prefix)
  1146→            if not resolved:
  1147→                raise FileNotFoundError(f"Missing tar in R2: {record.r2_tar_key}")
  1148→            local_path = download_dir / f"{record.youtube_id}.tar"
  1149→            if not local_path.exists():
  1150→                s3.download_file(cleaned_bucket, resolved, str(local_path), Config=config)
  1151→            return local_path
  1152→
  1153→    download_paths: dict[str, Path] = {}
  1154→    if args.skip_downloads:
  1155→        local_dir = args.local_tars_dir or download_dir
  1156→        if not local_dir.exists():
  1157→            raise SystemExit(f"Local tar directory not found: {local_dir}")
  1158→        missing = 0
  1159→        for record in tqdm(records, desc="local"):
  1160→            local_path = local_dir / f"{record.youtube_id}.tar"
  1161→            if local_path.exists():
  1162→                download_paths[record.youtube_id] = local_path
  1163→            else:
  1164→                missing += 1
  1165→        if missing:
  1166→            print(f"Missing local tars: {missing} (skipping)")
  1167→    else:
  1168→        with ThreadPoolExecutor(max_workers=args.download_workers) as executor:
  1169→            futures = {
  1170→                executor.submit(download_record, record): record.youtube_id
  1171→                for record in records
  1172→            }
  1173→            for future in tqdm(as_completed(futures), total=len(futures), desc="download"):
  1174→                youtube_id = futures[future]
  1175→                try:
  1176→                    download_paths[youtube_id] = future.result()
  1177→                except Exception as exc:
  1178→                    print(f"download failed {youtube_id}: {exc}")
  1179→
  1180→    if not download_paths:
  1181→        raise SystemExit("No downloads succeeded.")
  1182→
  1183→    video_ids = list(download_paths.keys())
  1184→    random.seed(args.seed)
  1185→    val_count = max(1, int(len(video_ids) * args.val_ratio)) if len(video_ids) > 1 else 0
  1186→    val_ids = set(random.sample(video_ids, val_count)) if val_count > 0 else set()
  1187→
  1188→    manifest_dir = output_dir / "manifests"
  1189→    stats_dir = output_dir / "stats"
  1190→    manifest_dir.mkdir(parents=True, exist_ok=True)
  1191→    stats_dir.mkdir(parents=True, exist_ok=True)
  1192→
  1193→    shard_prefix = "shard" if not partition_tag else f"shard-{partition_tag}"
  1194→    writers = {
  1195→        16000: ShardWriter(
  1196→            output_dir / f"{language}_16",
  1197→            shard_prefix,
  1198→            args.target_shard_size_mb * 1024 * 1024,
  1199→            audio_ext,
  1200→        ),
  1201→        24000: ShardWriter(
  1202→            output_dir / f"{language}_24",
  1203→            shard_prefix,
  1204→            args.target_shard_size_mb * 1024 * 1024,
  1205→            audio_ext,
  1206→        ),
  1207→        32000: ShardWriter(
  1208→            output_dir / f"{language}_32",
  1209→            shard_prefix,
  1210→            args.target_shard_size_mb * 1024 * 1024,
  1211→            audio_ext,
  1212→        ),
  1213→        44100: ShardWriter(
  1214→            output_dir / f"{language}_44",
  1215→            shard_prefix,
  1216→            args.target_shard_size_mb * 1024 * 1024,
  1217→            audio_ext,
  1218→        ),
  1219→    }
  1220→
  1221→    manifest_suffix = f"_{partition_tag}" if partition_tag else ""
  1222→    train_manifest = (manifest_dir / f"{language}_train{manifest_suffix}.jsonl").open(
  1223→        "w", encoding="utf-8"
  1224→    )
  1225→    val_manifest = (manifest_dir / f"{language}_val{manifest_suffix}.jsonl").open(
  1226→        "w", encoding="utf-8"
  1227→    )
  1228→
  1229→    stats = {
  1230→        "videos": 0,
  1231→        "segments_total": 0,
  1232→        "segments_kept": {16000: 0, 24000: 0, 32000: 0, 44100: 0},
  1233→        "dropped_too_short": 0,
  1234→        "dropped_too_long": 0,
  1235→        "dropped_clipped": 0,
  1236→        "segments_chunked": 0,
  1237→        "chunks_emitted": 0,
  1238→        "native_sr_buckets": {},
  1239→    }
  1240→
  1241→    segment_workers = max(1, args.segment_workers)
  1242→    max_inflight = (
  1243→        args.segment_max_inflight
  1244→        if args.segment_max_inflight > 0
  1245→        else max(128, segment_workers * 2)
  1246→    )
  1247→    pending: set[Any] = set()
  1248→    future_info: dict[Any, bool] = {}
  1249→
  1250→    def handle_future(future) -> None:
  1251→        is_val = future_info.pop(future, False)
  1252→        outputs, seg_stats = future.result()
  1253→        stats["dropped_too_short"] += seg_stats["dropped_too_short"]
  1254→        stats["dropped_too_long"] += seg_stats["dropped_too_long"]
  1255→        stats["dropped_clipped"] += seg_stats["dropped_clipped"]
  1256→        stats["segments_chunked"] += seg_stats.get("segments_chunked", 0)
  1257→        stats["chunks_emitted"] += seg_stats.get("chunks_emitted", 0)
  1258→        for out in outputs:
  1259→            stats["segments_kept"][out.sample_rate] += 1
  1260→            if out.sample_rate == 16000:
  1261→                bucket = out.metadata.get("native_sr_bucket")
  1262→                if bucket:
  1263→                    stats["native_sr_buckets"][bucket] = (
  1264→                        stats["native_sr_buckets"].get(bucket, 0) + 1
  1265→                    )
  1266→            writer = writers[out.sample_rate]
  1267→            with timer("shard_write"):
  1268→                shard_name = writer.add_sample(
  1269→                    out.metadata["id"], out.audio_bytes, out.metadata
  1270→                )
  1271→            manifest_entry = dict(out.metadata)
  1272→            manifest_entry["shard"] = f"{shard_name}.tar"
  1273→            manifest_entry["audio"] = f"{out.metadata['id']}.{audio_ext}"
  1274→            line = json.dumps(manifest_entry) + "\n"
  1275→            if is_val:
  1276→                val_manifest.write(line)
  1277→            else:
  1278→                train_manifest.write(line)
  1279→
  1280→    def drain_completed() -> None:
  1281→        nonlocal pending
  1282→        if not pending:
  1283→            return
  1284→        done, pending = wait(pending, return_when=FIRST_COMPLETED)
  1285→        for future in done:
  1286→            handle_future(future)
  1287→
  1288→    with ThreadPoolExecutor(max_workers=segment_workers) as executor:
  1289→        for video_id, tar_path in tqdm(download_paths.items(), desc="process"):
  1290→            stats["videos"] += 1
  1291→            is_val = video_id in val_ids
  1292→            with timer("tar_open"):
  1293→                tf = tarfile.open(tar_path, "r")
  1294→            with tf:
  1295→                with timer("tar_meta_load"):
  1296→                    meta = json.loads(
  1297→                        tf.extractfile(f"{video_id}/metadata.json").read()
  1298→                    )
  1299→                usable_segments = [
  1300→                    s for s in meta.get("segments", []) if s.get("status") == "usable"
  1301→                ]
  1302→                # Create lookup by (speaker, start) for robust matching
  1303→                # The filename contains 'start' timestamp (relative to intro-skipped audio),
  1304→                # not 'original_start' (absolute time in original video)
  1305→                with timer("tar_build_segment_map"):
  1306→                    segment_map = {
  1307→                        (seg.get("speaker"), round(seg.get("start", 0.0), 2)): seg
  1308→                        for seg in usable_segments
  1309→                    }
  1310→
  1311→                with timer("tar_list_members"):
  1312→                    members = [m for m in tf.getmembers() if m.name.endswith(".flac")]
  1313→                stats["segments_total"] += len(members)
  1314→
  1315→                for member in members:
  1316→                    parsed = parse_segment_filename(member.name)
  1317→                    if not parsed:
  1318→                        continue
  1319→                    speaker, start_time = parsed
  1320→                    # Match by (speaker, start) with rounding for float precision
  1321→                    seg_meta = segment_map.get((speaker, round(start_time, 2)))
  1322→                    if not seg_meta:
  1323→                        # Try fuzzy match for edge cases (tolerance 0.1s)
  1324→                        matches = [
  1325→                            s for s in usable_segments
  1326→                            if s.get("speaker") == speaker
  1327→                            and abs(s.get("start", 0.0) - start_time) < 0.1
  1328→                        ]
  1329→                        if not matches:
  1330→                            continue
  1331→                        seg_meta = matches[0]
  1332→
  1333→                    while len(pending) >= max_inflight:
  1334→                        drain_completed()
  1335→
  1336→                    with timer("tar_extract_segment"):
  1337→                        segment_bytes = tf.extractfile(member).read()
  1338→                    # Extract index from filename for ID generation (cosmetic only)
  1339→                    filename_parts = member.name.split("_")
  1340→                    try:
  1341→                        index = int(filename_parts[-2])
  1342→                    except (ValueError, IndexError):
  1343→                        index = 0
  1344→                    rng_seed = make_segment_seed(args.seed, video_id, speaker, index)
  1345→
  1346→                    future = executor.submit(
  1347→                        process_segment,
  1348→                        video_id,
  1349→                        language,
  1350→                        segment_bytes,
  1351→                        seg_meta,
  1352→                        speaker,
  1353→                        index,
  1354→                        chunk_cfg,
  1355→                        silence_cfg,
  1356→                        rng_seed,
  1357→                        output_format,
  1358→                    )
  1359→                    future_info[future] = is_val
  1360→                    pending.add(future)
  1361→
  1362→        while pending:
  1363→            drain_completed()
  1364→
  1365→    for writer in writers.values():
  1366→        writer.close()
  1367→    train_manifest.close()
  1368→    val_manifest.close()
  1369→
  1370→    stats_suffix = f"_{partition_tag}" if partition_tag else ""
  1371→    stats_path = stats_dir / f"{language}_stats{stats_suffix}.json"
  1372→    stats_path.write_text(json.dumps(stats, indent=2))
  1373→
  1374→    if args.upload:
  1375→        r2_subprefix = args.r2_subprefix.strip("/")
  1376→        upload_map = {
  1377→            16000: f"{language}_16",
  1378→            24000: f"{language}_24",
  1379→            32000: f"{language}_32",
  1380→            44100: f"{language}_44",
  1381→        }
  1382→        for sr, local_prefix in upload_map.items():
  1383→            remote_prefix = local_prefix if not r2_subprefix else f"{local_prefix}/{r2_subprefix}"
  1384→            upload_directory(
  1385→                s3,
  1386→                shard_bucket,
  1387→                output_dir / local_prefix,
  1388→                remote_prefix,
  1389→                args.download_workers,
  1390→            )
  1391→        upload_directory(
  1392→            s3,
  1393→            shard_bucket,
  1394→            manifest_dir,
  1395→            "manifests" if not r2_subprefix else f"manifests/{r2_subprefix}",
  1396→            args.download_workers,
  1397→        )
  1398→        upload_directory(
  1399→            s3,
  1400→            shard_bucket,
  1401→            stats_dir,
  1402→            "stats" if not r2_subprefix else f"stats/{r2_subprefix}",
  1403→            args.download_workers,
  1404→        )
  1405→
  1406→        if args.verify_upload:
  1407→            verify_dir = tmp_dir / "verify"
  1408→            verify_dir.mkdir(parents=True, exist_ok=True)
  1409→            for sr, local_prefix in upload_map.items():
  1410→                shard_dir = output_dir / local_prefix
  1411→                shard_files = sorted(p.stem for p in shard_dir.glob("*.tar"))
  1412→                if not shard_files:
  1413→                    continue
  1414→                shard_name = shard_files[0]
  1415→                remote_prefix = local_prefix if not r2_subprefix else f"{local_prefix}/{r2_subprefix}"
  1416→                result = verify_uploaded_shard(
  1417→                    s3,
  1418→                    shard_bucket,
  1419→                    remote_prefix,
  1420→                    shard_name,
  1421→                    sr,
  1422→                    audio_ext,
  1423→                    verify_dir,
  1424→                )
  1425→                print(
  1426→                    f"verify {language} {sr}k: "
  1427→                    f"checked={result['checked_samples']} "
  1428→                    f"missing_or_mismatch={result['missing_or_mismatch']} "
  1429→                    f"entries={result['total_entries']}"
  1430→                )
  1431→
  1432→    if _timing_enabled:
  1433→        print_timing_stats()
  1434→
  1435→    print(f"Done. Stats written to {stats_path}")
  1436→
  1437→
  1438→if __name__ == "__main__":
  1439→    main()
  1440→

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

from __future__ import annotations

import argparse
import io
import json
import math
import os
import random
import tarfile
import threading
import time
import zlib
from collections import defaultdict
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Optional

import boto3
import numpy as np
import psycopg
import soundfile as sf
from boto3.s3.transfer import TransferConfig
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, as_completed, wait
from scipy import signal
from tqdm import tqdm

try:
    import pyloudnorm as pyln
except Exception:  # pragma: no cover - optional dependency
    pyln = None

ENV_PATH = Path(".env")
LANGUAGES = [
    "telugu",
    "malayalam",
    "english",
    "hindi",
    "punjabi",
    "tamil",
    "kannada",
    "gujarati",
    "bengali",
    "odia",
    "marathi",
    "assamese",
]

# Global timing stats (thread-safe)
_timing_lock = threading.Lock()
_timing_stats: dict[str, list[float]] = defaultdict(list)
_timing_enabled = False

@contextmanager
def timer(name: str):
    """Context manager for timing code blocks."""
    if not _timing_enabled:
        yield
        return
    start = time.perf_counter()
    try:
        yield
    finally:
        elapsed = time.perf_counter() - start
        with _timing_lock:
            _timing_stats[name].append(elapsed)

def print_timing_stats():
    """Print aggregated timing statistics."""
    if not _timing_stats:
        return
    print("\n=== Timing Statistics ===")
    total_time = sum(sum(times) for times in _timing_stats.values())
    for name, times in sorted(_timing_stats.items()):
        count = len(times)
        total = sum(times)
        avg = total / count if count > 0 else 0
        percent = (total / total_time * 100) if total_time > 0 else 0
        print(f"{name:35s}: {count:6d} calls, {total:8.2f}s total, {avg*1000:8.2f}ms avg, {percent:5.1f}%")
    print(f"{'TOTAL':35s}:              {total_time:8.2f}s")
    print("=" * 80)

_window_cache: dict[int, np.ndarray] = {}
_freq_cache: dict[tuple[int, int], np.ndarray] = {}
_cache_lock = threading.Lock()
_lufs_local = threading.local()

def get_hann_window(length: int) -> np.ndarray:
    window = _window_cache.get(length)
    if window is not None:
        return window
    with _cache_lock:
        window = _window_cache.get(length)
        if window is None:
            window = np.hanning(length)
            _window_cache[length] = window
    return window

def get_rfftfreq(length: int, sr: int) -> np.ndarray:
    key = (length, sr)
    freqs = _freq_cache.get(key)
    if freqs is not None:
        return freqs
    with _cache_lock:
        freqs = _freq_cache.get(key)
        if freqs is None:
            freqs = np.fft.rfftfreq(length, d=1.0 / sr)
            _freq_cache[key] = freqs
    return freqs

def get_lufs_meter(sr: int) -> Optional[Any]:
    if pyln is None:
        return None
    meters = getattr(_lufs_local, "meters", None)
    if meters is None:
        meters = {}
        _lufs_local.meters = meters
    meter = meters.get(sr)
    if meter is None:
        meter = pyln.Meter(sr)
        meters[sr] = meter
    return meter

@dataclass(frozen=True)
class VideoRecord:
    youtube_id: str
    r2_tar_key: str

@dataclass(frozen=True)
class SegmentOutput:
    sample_rate: int
    audio_bytes: bytes
    metadata: dict[str, Any]

@dataclass(frozen=True)
class ChunkConfig:
    enabled: bool
    min_sec: float
    max_sec: float
    threshold_sec: float
    all_buckets: bool

@dataclass(frozen=True)
class SilenceConfig:
    enabled: bool
    frame_sec: float
    hop_sec: float
    search_sec: float
    percentile: float
    db_drop: float

def load_env(env_path: Path) -> dict[str, str]:
    env: dict[str, str] = {}
    for line in env_path.read_text().splitlines():
        stripped = line.strip()
        if not stripped or stripped.startswith("#") or "=" not in stripped:
            continue
        if stripped.startswith("export "):
            stripped = stripped[len("export ") :]
        key, value = stripped.split("=", 1)
        env[key.strip()] = value.strip()
    return env

def make_r2_client(env: dict[str, str]):
    endpoint = env.get("R2_ENDPOINT_URL")
    access_key = env.get("R2_ACCESS_KEY_ID")
    secret_key = env.get("R2_SECRET_ACCESS_KEY")
    if not endpoint or not access_key or not secret_key:
        raise SystemExit("Missing R2 credentials in .env")
    return boto3.client(
        "s3",
        endpoint_url=endpoint,
        aws_access_key_id=access_key,
        aws_secret_access_key=secret_key,
    )

def resolve_r2_key(
    s3, bucket: str, r2_key: str, prefix: Optional[str]
) -> Optional[str]:
    if not r2_key:
        return None
    key = r2_key.lstrip("/")
    candidates = [key]

    if prefix:
        cleaned_prefix = prefix.strip("/")
        if not key.startswith(cleaned_prefix):
            candidates.append(f"{cleaned_prefix}/{key}")

    base = os.path.basename(key)
    if base != key:
        candidates.append(base)
        if prefix:
            candidates.append(f"{prefix.strip('/')}/{base}")

    for candidate in candidates:
        try:
            with timer("r2_head_object"):
                s3.head_object(Bucket=bucket, Key=candidate)
            return candidate
        except Exception:
            continue
    return None

def parse_source_url(source_url: Optional[str], default_bucket: Optional[str]) -> Optional[tuple[str, str]]:
    if not source_url or not default_bucket:
        return None
    cleaned = source_url.lstrip("/")
    if cleaned.startswith(default_bucket + "/"):
        return default_bucket, cleaned[len(default_bucket) + 1 :]
    return default_bucket, cleaned

def fetch_video_records(
    db_url: str,
    language: str,
    limit: Optional[int],
    partition_count: int = 1,
    partition_index: int = 0,
) -> list[VideoRecord]:
    records: list[VideoRecord] = []
    with psycopg.connect(db_url) as conn:
        with conn.cursor() as cur:
            query = """
                select youtube_id, r2_tar_key
                from videos
                where lower(language) = %s
                  and r2_tar_key is not null
                  and status = 'COMPLETED'
            """
            params = [language.lower()]
            if partition_count > 1:
                query += " and mod(abs(hashtext(youtube_id)), %s) = %s"
                params.extend([partition_count, partition_index])
            query += " order by random()"
            if limit:
                query += " limit %s"
                params.append(limit)
            cur.execute(query, params)
            for youtube_id, r2_tar_key in cur.fetchall():
                if youtube_id and r2_tar_key:
                    records.append(VideoRecord(youtube_id, r2_tar_key))
    return records

def normalize(value: float, min_val: float, max_val: float) -> float:
    if max_val <= min_val:
        return 0.0
    return max(0.0, min(1.0, (value - min_val) / (max_val - min_val)))

def compute_spectral_metrics(audio: np.ndarray, sr: int) -> tuple[float, float, float]:
    if audio.size == 0:
        return 0.0, 0.0, 0.0
    with timer("cpu_spectral_compute"):
        max_samples = min(audio.size, sr * 5)
        window = get_hann_window(max_samples)
        chunk = audio[:max_samples] * window
        spec = np.fft.rfft(chunk)
        power = np.abs(spec) ** 2
        freqs = get_rfftfreq(max_samples, sr)
        total_energy = float(np.sum(power)) + 1e-12
        cumulative = np.cumsum(power)
        rolloff_idx = int(np.searchsorted(cumulative, 0.85 * cumulative[-1]))
        rolloff_hz = float(freqs[min(rolloff_idx, len(freqs) - 1)])

        power_db = 10.0 * np.log10(power + 1e-12)
        threshold = power_db.max() - 60.0
        above = np.where(power_db >= threshold)[0]
        effective_nyquist_hz = float(freqs[above[-1]]) if above.size > 0 else 0.0

        high_freq_energy = float(np.sum(power[freqs >= 8000.0]))
        high_freq_ratio = high_freq_energy / total_energy
    return rolloff_hz, effective_nyquist_hz, high_freq_ratio

def classify_native_bucket(effective_nyquist_hz: float) -> str:
    """Classify audio bandwidth based on effective Nyquist frequency.

    These buckets represent actual audio bandwidth, not container sample rates.
    Thresholds are based on observed distribution in the dataset.
    """
    if effective_nyquist_hz < 8000:
        return "<16k"
    if effective_nyquist_hz < 12000:
        return "16-24k"
    if effective_nyquist_hz < 16000:
        return "24-32k"
    if effective_nyquist_hz < 22050:
        return "32-44k"
    return ">=44k"

def compute_audio_metrics(audio: np.ndarray, sr: int) -> dict[str, float]:
    with timer("audio_metrics_basic"):
        audio = audio.astype(np.float32, copy=False)
        rms = float(np.sqrt(np.mean(audio ** 2) + 1e-12))
        loudness_rms_db = 20.0 * math.log10(rms + 1e-12)
        clipping_fraction = float(np.mean(np.abs(audio) > 0.99))
        dc_offset = float(np.mean(audio))
        noise_floor = float(np.percentile(np.abs(audio), 10))
        snr_proxy = 20.0 * math.log10((rms + 1e-9) / (noise_floor + 1e-9))

    with timer("audio_metrics_lufs"):
        loudness_lufs = -100.0
        if pyln is not None:
            try:
                meter = get_lufs_meter(sr)
                if meter is not None:
                    loudness_lufs = float(meter.integrated_loudness(audio))
            except Exception:
                loudness_lufs = -100.0

    rolloff_hz, effective_nyquist_hz, high_freq_ratio = compute_spectral_metrics(audio, sr)

    return {
        "loudness_lufs": loudness_lufs,
        "loudness_rms_db": loudness_rms_db,
        "clipping_fraction": clipping_fraction,
        "dc_offset": dc_offset,
        "snr_proxy": snr_proxy,
        "spectral_rolloff_hz": rolloff_hz,
        "effective_nyquist_hz": effective_nyquist_hz,
        "high_freq_energy_ratio": high_freq_ratio,
    }

def compute_quality_score(metrics: dict[str, float], has_music: bool) -> float:
    score = (
        0.3 * (1.0 - metrics["clipping_fraction"])
        + 0.2 * normalize(metrics["snr_proxy"], 10, 40)
        + 0.2 * (1.0 if not has_music else 0.5)
        + 0.2 * normalize(metrics["loudness_lufs"], -40, -14)
        + 0.1 * (1.0 if abs(metrics["dc_offset"]) < 0.01 else 0.5)
    )
    return max(0.0, min(1.0, score))

def resample_audio(
    audio: np.ndarray,
    src_sr: int,
    dst_sr: int,
) -> np.ndarray:
    if src_sr == dst_sr:
        return audio
    with timer("cpu_resample"):
        gcd = math.gcd(src_sr, dst_sr)
        up = dst_sr // gcd
        down = src_sr // gcd
        return signal.resample_poly(audio, up, down).astype(np.float32, copy=False)

def encode_audio(audio: np.ndarray, sr: int, output_format: str) -> bytes:
    with timer("audio_encode"):
        buf = io.BytesIO()
        if output_format == "wav":
            sf.write(buf, audio, sr, format="WAV", subtype="PCM_16")
        else:
            sf.write(buf, audio, sr, format="FLAC", subtype="PCM_16")
        return buf.getvalue()

def duration_bucket(duration_sec: float) -> Optional[tuple[str, float]]:
    if duration_sec < 3.0:
        return None
    if duration_sec < 6.0:
        return "3-6s", 0.3
    if duration_sec < 12.0:
        return "6-12s", 1.0
    if duration_sec < 24.0:
        return "12-24s", 0.8
    if duration_sec <= 60.0:
        return "24-60s", 0.5
    return None

def compute_frame_rms(
    audio: np.ndarray, sr: int, frame_sec: float, hop_sec: float
) -> tuple[np.ndarray, np.ndarray]:
    frame_len = max(1, int(round(frame_sec * sr)))
    hop_len = max(1, int(round(hop_sec * sr)))
    if audio.size < frame_len:
        return np.array([], dtype=np.float32), np.array([], dtype=np.float32)

    if not audio.flags["C_CONTIGUOUS"]:
        audio = np.ascontiguousarray(audio)

    num_frames = 1 + (audio.size - frame_len) // hop_len
    if num_frames <= 0:
        return np.array([], dtype=np.float32), np.array([], dtype=np.float32)

    stride = audio.strides[0]
    frames = np.lib.stride_tricks.as_strided(
        audio,
        shape=(num_frames, frame_len),
        strides=(hop_len * stride, stride),
    )
    rms = np.sqrt(np.mean(frames * frames, axis=1) + 1e-12)
    times = (np.arange(num_frames, dtype=np.float32) * hop_len + (frame_len / 2)) / sr
    return times.astype(np.float32), rms.astype(np.float32)

def find_silence_boundary(
    times: np.ndarray,
    rms: np.ndarray,
    search_start: float,
    search_end: float,
    threshold_rms: float,
    median_db: float,
    db_drop: float,
) -> Optional[float]:
    if times.size == 0:
        return None
    mask = (times >= search_start) & (times <= search_end)
    if not np.any(mask):
        return None
    subset = rms[mask]
    min_idx = int(np.argmin(subset))
    min_rms = float(subset[min_idx])
    min_time = float(times[mask][min_idx])
    min_db = 20.0 * math.log10(min_rms + 1e-12)
    if min_rms <= threshold_rms or (median_db - min_db) >= db_drop:
        return min_time
    return None

def make_chunk_windows(
    audio: np.ndarray,
    sr: int,
    cfg: ChunkConfig,
    silence_cfg: SilenceConfig,
    rng: random.Random,
) -> list[tuple[float, float]]:
    duration_sec = float(audio.shape[0] / sr)
    if not cfg.enabled or duration_sec < cfg.threshold_sec:
        return [(0.0, duration_sec)]

    if not silence_cfg.enabled:
        windows: list[tuple[float, float]] = []
        offset = 0.0
        while duration_sec - offset >= cfg.min_sec:
            remaining = duration_sec - offset
            if remaining <= cfg.max_sec:
                length = remaining
            else:
                length = rng.uniform(cfg.min_sec, cfg.max_sec)
                if remaining - length < cfg.min_sec:
                    length = remaining
            windows.append((offset, length))
            offset += length
        return windows

    times, rms = compute_frame_rms(
        audio, sr, silence_cfg.frame_sec, silence_cfg.hop_sec
    )
    if times.size == 0:
        return [(0.0, duration_sec)]

    threshold_rms = float(np.quantile(rms, silence_cfg.percentile / 100.0))
    median_db = 20.0 * math.log10(float(np.median(rms)) + 1e-12)

    windows: list[tuple[float, float]] = []
    offset = 0.0
    while duration_sec - offset > cfg.max_sec:
        remaining = duration_sec - offset
        latest_cut = duration_sec - cfg.min_sec
        if offset + cfg.min_sec >= latest_cut:
            break
        target_len = rng.uniform(cfg.min_sec, cfg.max_sec)
        if remaining - target_len < cfg.min_sec:
            target_len = remaining - cfg.min_sec
        target_end = offset + target_len
        search_low = max(offset + cfg.min_sec, target_end - silence_cfg.search_sec)
        search_high = min(offset + cfg.max_sec, target_end + silence_cfg.search_sec, latest_cut)
        if search_high <= search_low:
            search_low = offset + cfg.min_sec
            search_high = min(offset + cfg.max_sec, latest_cut)
        boundary = find_silence_boundary(
            times,
            rms,
            search_low,
            search_high,
            threshold_rms,
            median_db,
            silence_cfg.db_drop,
        )
        if boundary is None:
            if not windows:
                return [(0.0, duration_sec)]
            windows.append((offset, duration_sec - offset))
            return windows
        windows.append((offset, boundary - offset))
        offset = boundary

    windows.append((offset, duration_sec - offset))
    return windows

def make_segment_seed(base_seed: int, video_id: str, speaker: str, index: int) -> int:
    key = f"{video_id}-{speaker}-{index}-{base_seed}"
    return zlib.adler32(key.encode("utf-8")) & 0xFFFFFFFF

def parse_segment_filename(filename: str) -> Optional[tuple[str, float]]:
    """Parse segment filename to extract speaker and start time.

    Filename format: segments/SPEAKER_XX_NNNN_start-end.flac
    Example: segments/SPEAKER_00_0007_35.10-37.49.flac

    Returns (speaker, start_time) for matching against metadata.
    """
    base = filename.split("/")[-1]
    if not base.endswith(".flac"):
        return None

    # Remove .flac extension
    base = base[:-5]

    # Try to extract time range (last part after final underscore)
    parts = base.split("_")
    if len(parts) < 3:
        return None

    # Speaker is everything except last two parts (index and time range)
    speaker = "_".join(parts[:-2])

    # Time range is the last part: "35.10-37.49"
    time_range = parts[-1]
    if "-" not in time_range:
        return None

    try:
        start_str = time_range.split("-")[0]
        start_time = float(start_str)
    except (ValueError, IndexError):
        return None

    return speaker, start_time

def process_segment(
    video_id: str,
    language: str,
    segment_bytes: bytes,
    segment_meta: dict[str, Any],
    speaker: str,
    index: int,
    chunk_cfg: ChunkConfig,
    silence_cfg: SilenceConfig,
    rng_seed: int,
    output_format: str,
) -> tuple[list[SegmentOutput], dict[str, int]]:
    with timer("segment_total"):
        stats = {
            "dropped_too_short": 0,
            "dropped_too_long": 0,
            "dropped_clipped": 0,
            "chunks_emitted": 0,
            "segments_chunked": 0,
        }
        outputs: list[SegmentOutput] = []

        with timer("flac_decode"):
            try:
                audio, sr = sf.read(io.BytesIO(segment_bytes), dtype="float32")
            except Exception:
                return outputs, stats

        if audio.ndim > 1:
            audio = np.mean(audio, axis=1)

        rng = random.Random(rng_seed)
        windows = make_chunk_windows(audio, sr, chunk_cfg, silence_cfg, rng)

        music_stats = segment_meta.get("music_stats") or {}
        has_music = bool(music_stats.get("has_music", False))

        def emit_outputs(
            audio_slice: np.ndarray,
            metrics: dict[str, float],
            sample_id: str,
            segment_index: int,
            chunk_index: Optional[int],
            start: Optional[float],
            end: Optional[float],
            original_start: Optional[float],
            original_end: Optional[float],
            target_srs: list[int],
        ) -> None:
            slice_duration = float(audio_slice.shape[0] / sr)
            bucket = duration_bucket(slice_duration)
            if bucket is None:
                if slice_duration < 3.0:
                    stats["dropped_too_short"] += 1
                else:
                    stats["dropped_too_long"] += 1
                return
            if metrics["clipping_fraction"] > 0.1:
                stats["dropped_clipped"] += 1
                return

            duration_bucket_name, base_weight = bucket
            sampling_weight = base_weight * (0.5 if has_music else 1.0)
            quality_score = compute_quality_score(metrics, has_music)
            effective_nyquist = metrics["effective_nyquist_hz"]

            base_metadata = {
                "id": sample_id,
                "video_id": video_id,
                "language": language,
                "segment_index": segment_index,
                "chunk_index": chunk_index,
                "start": start,
                "end": end,
                "original_start": original_start,
                "original_end": original_end,
                "duration_sec": slice_duration,
                "native_sr": sr,
                "effective_nyquist_hz": effective_nyquist,
                "spectral_rolloff_hz": metrics["spectral_rolloff_hz"],
                "high_freq_energy_ratio": metrics["high_freq_energy_ratio"],
                "is_truly_wideband": bool(effective_nyquist >= 12000),
                "native_sr_bucket": classify_native_bucket(effective_nyquist),
                "loudness_lufs": metrics["loudness_lufs"],
                "loudness_rms_db": metrics["loudness_rms_db"],
                "clipping_fraction": metrics["clipping_fraction"],
                "snr_proxy": metrics["snr_proxy"],
                "quality_score": quality_score,
                "music_prob": float(music_stats.get("music_mean", 0.0)),
                "music_stats": music_stats,
                "speaker": speaker,
                # VAD regions: Full duration placeholder; random 6s crops happen during training.
                "vad_speech_regions": [[0.0, slice_duration]],
                "duration_bucket": duration_bucket_name,
                "sampling_weight": sampling_weight,
            }

            for target_sr in target_srs:
                resampled = resample_audio(audio_slice, sr, target_sr)
                resampled = np.clip(resampled, -1.0, 1.0)
                audio_bytes = encode_audio(resampled, target_sr, output_format)
                meta = dict(base_metadata)
                meta["sample_rate"] = target_sr
                meta["resample_target_sr"] = target_sr
                outputs.append(SegmentOutput(target_sr, audio_bytes, meta))

            if target_srs == [16000]:
                stats["chunks_emitted"] += 1

        sample_id_base = f"{video_id}__{speaker}_{index:04d}"
        use_chunk_suffix = len(windows) > 1
        if use_chunk_suffix:
            stats["segments_chunked"] += 1

        segment_start = segment_meta.get("start")
        segment_end = segment_meta.get("end")
        original_start = segment_meta.get("original_start")
        original_end = segment_meta.get("original_end")

        # Always emit 16k outputs; for long segments, emit chunked windows to boost sample count.
        for chunk_idx, (offset_sec, chunk_len_sec) in enumerate(windows):
            start_idx = int(round(offset_sec * sr))
            end_idx = int(round((offset_sec + chunk_len_sec) * sr))
            audio_slice = audio[start_idx:end_idx]
            if audio_slice.size == 0:
                continue
            slice_duration = float(audio_slice.shape[0] / sr)
            slice_metrics = compute_audio_metrics(audio_slice, sr)

            chunk_start = float(segment_start) + offset_sec if segment_start is not None else None
            chunk_end = chunk_start + slice_duration if chunk_start is not None else None
            chunk_original_start = (
                float(original_start) + offset_sec if original_start is not None else None
            )
            chunk_original_end = (
                chunk_original_start + slice_duration
                if chunk_original_start is not None
                else None
            )

            sample_id = sample_id_base
            chunk_index = None
            if use_chunk_suffix:
                sample_id = f"{sample_id_base}_c{chunk_idx:02d}"
                chunk_index = chunk_idx

            emit_outputs(
                audio_slice,
                slice_metrics,
                sample_id,
                index,
                chunk_index,
                chunk_start,
                chunk_end,
                chunk_original_start,
                chunk_original_end,
                [16000],
            )

        # Optional high-rate buckets: use full segment unless configured otherwise.
        full_metrics = compute_audio_metrics(audio, sr)
        effective_nyquist = full_metrics["effective_nyquist_hz"]
        targets_high: list[int] = []
        if effective_nyquist >= 8000:
            targets_high.append(24000)
        if effective_nyquist >= 12000:
            targets_high.append(32000)
        if effective_nyquist >= 16000:
            targets_high.append(44100)

        if targets_high:
            if chunk_cfg.all_buckets and use_chunk_suffix:
                for chunk_idx, (offset_sec, chunk_len_sec) in enumerate(windows):
                    start_idx = int(round(offset_sec * sr))
                    end_idx = int(round((offset_sec + chunk_len_sec) * sr))
                    audio_slice = audio[start_idx:end_idx]
                    if audio_slice.size == 0:
                        continue
                    slice_duration = float(audio_slice.shape[0] / sr)
                    slice_metrics = compute_audio_metrics(audio_slice, sr)
                    slice_effective_nyquist = slice_metrics["effective_nyquist_hz"]
                    slice_targets: list[int] = []
                    if slice_effective_nyquist >= 8000:
                        slice_targets.append(24000)
                    if slice_effective_nyquist >= 12000:
                        slice_targets.append(32000)
                    if slice_effective_nyquist >= 16000:
                        slice_targets.append(44100)
                    if not slice_targets:
                        continue

                    chunk_start = (
                        float(segment_start) + offset_sec if segment_start is not None else None
                    )
                    chunk_end = (
                        chunk_start + slice_duration if chunk_start is not None else None
                    )
                    chunk_original_start = (
                        float(original_start) + offset_sec if original_start is not None else None
                    )
                    chunk_original_end = (
                        chunk_original_start + slice_duration
                        if chunk_original_start is not None
                        else None
                    )

                    sample_id = f"{sample_id_base}_c{chunk_idx:02d}"
                    emit_outputs(
                        audio_slice,
                        slice_metrics,
                        sample_id,
                        index,
                        chunk_idx,
                        chunk_start,
                        chunk_end,
                        chunk_original_start,
                        chunk_original_end,
                        slice_targets,
                    )
            else:
                emit_outputs(
                    audio,
                    full_metrics,
                    sample_id_base,
                    index,
                    None,
                    segment_start,
                    segment_end,
                    original_start,
                    original_end,
                    targets_high,
                )

        return outputs, stats

class ShardWriter:
    def __init__(
        self, output_dir: Path, prefix: str, target_size_bytes: int, audio_ext: str
    ) -> None:
        self.output_dir = output_dir
        self.prefix = prefix
        self.target_size_bytes = target_size_bytes
        self.audio_ext = audio_ext
        self.shard_index = 0
        self.current_tar: Optional[tarfile.TarFile] = None
        self.current_jsonl = None
        self.current_size = 0
        self.current_count = 0
        self.current_name = ""
        self.output_dir.mkdir(parents=True, exist_ok=True)

    def _open_new_shard(self) -> None:
        if self.current_tar:
            self.close()
        shard_name = f"{self.prefix}-{self.shard_index:06d}"
        tar_path = self.output_dir / f"{shard_name}.tar"
        jsonl_path = self.output_dir / f"{shard_name}.jsonl"
        self.current_tar = tarfile.open(tar_path, "w")
        self.current_jsonl = jsonl_path.open("w", encoding="utf-8")
        self.current_size = 0
        self.current_count = 0
        self.current_name = shard_name
        self.shard_index += 1

    def add_sample(self, sample_id: str, audio_bytes: bytes, metadata: dict[str, Any]) -> str:
        if self.current_tar is None:
            self._open_new_shard()
        if (
            self.current_count > 0
            and self.current_size + len(audio_bytes) > self.target_size_bytes
        ):
            self._open_new_shard()

        assert self.current_tar is not None
        assert self.current_jsonl is not None

        arcname = f"{sample_id}.{self.audio_ext}"
        info = tarfile.TarInfo(name=arcname)
        info.size = len(audio_bytes)
        info.mtime = int(time.time())
        self.current_tar.addfile(info, io.BytesIO(audio_bytes))
        self.current_jsonl.write(json.dumps(metadata) + "\n")

        self.current_size += len(audio_bytes)
        self.current_count += 1
        return self.current_name

    def close(self) -> None:
        if self.current_tar:
            self.current_tar.close()
        if self.current_jsonl:
            self.current_jsonl.close()
        self.current_tar = None
        self.current_jsonl = None

def create_prefix_placeholders(
    s3,
    bucket: str,
    languages: list[str],
    rates: list[int],
    extra_prefixes: list[str],
    subprefix: Optional[str],
) -> None:
    prefixes = []
    for lang in languages:
        for rate in rates:
            base = f"{lang}_{rate}"
            prefixes.append(base)
            if subprefix:
                prefixes.append(f"{base}/{subprefix.strip('/')}")
    prefixes.extend(extra_prefixes)
    if subprefix:
        prefixes.extend([f"{p}/{subprefix.strip('/')}" for p in extra_prefixes])

    for prefix in prefixes:
        key = f"{prefix.strip('/')}/.keep"
        try:
            s3.put_object(Bucket=bucket, Key=key, Body=b"")
        except Exception:
            continue

def upload_directory(
    s3,
    bucket: str,
    local_dir: Path,
    remote_prefix: str,
    max_concurrency: int,
) -> list[str]:
    uploaded = []
    config = TransferConfig(
        multipart_threshold=8 * 1024 * 1024,
        multipart_chunksize=64 * 1024 * 1024,
        max_concurrency=max_concurrency,
        use_threads=True,
    )
    for path in sorted(local_dir.glob("*")):
        if not path.is_file():
            continue
        key = f"{remote_prefix}/{path.name}"
        s3.upload_file(str(path), bucket, key, Config=config)
        uploaded.append(key)
    return uploaded

def verify_uploaded_shard(
    s3,
    bucket: str,
    remote_prefix: str,
    shard_name: str,
    sample_rate: int,
    audio_ext: str,
    tmp_dir: Path,
) -> dict[str, Any]:
    tar_key = f"{remote_prefix}/{shard_name}.tar"
    jsonl_key = f"{remote_prefix}/{shard_name}.jsonl"
    tar_path = tmp_dir / f"{shard_name}.tar"
    jsonl_path = tmp_dir / f"{shard_name}.jsonl"

    s3.download_file(bucket, tar_key, str(tar_path))
    s3.download_file(bucket, jsonl_key, str(jsonl_path))

    with tarfile.open(tar_path, "r") as tf:
        tar_names = set(tf.getnames())
        with jsonl_path.open("r", encoding="utf-8") as fh:
            lines = [json.loads(line) for line in fh if line.strip()]

        missing = 0
        checked = 0
        for entry in lines[:10]:
            checked += 1
            audio_name = f"{entry['id']}.{audio_ext}"
            if audio_name not in tar_names:
                missing += 1
                continue
            member = tf.getmember(audio_name)
            audio_bytes = tf.extractfile(member).read()
            audio, sr = sf.read(io.BytesIO(audio_bytes))
            if sr != sample_rate:
                missing += 1

    return {
        "checked_samples": checked,
        "missing_or_mismatch": missing,
        "total_entries": len(lines),
    }

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Prepare sharded datasets from cleaned tars.")
    parser.add_argument("--language", required=True)
    parser.add_argument("--limit", type=int, default=None)
    parser.add_argument("--output-dir", type=Path, default=Path("processed_data"))
    parser.add_argument("--tmp-dir", type=Path, default=Path("/tmp/xcodec_work"))
    parser.add_argument("--target-shard-size-mb", type=int, default=5120)
    parser.add_argument("--download-workers", type=int, default=16)
    parser.add_argument("--segment-workers", type=int, default=8)
    parser.add_argument(
        "--segment-max-inflight",
        type=int,
        default=0,
        help="Limit number of in-flight segment tasks (0=auto).",
    )
    parser.add_argument(
        "--skip-downloads",
        action="store_true",
        help="Use only local tar files; do not download from R2.",
    )
    parser.add_argument(
        "--local-tars-dir",
        type=Path,
        default=None,
        help="Directory containing pre-downloaded tar files for this language.",
    )
    parser.add_argument(
        "--partition-count",
        type=int,
        default=1,
        help="Total number of partitions for parallel processing.",
    )
    parser.add_argument(
        "--partition-index",
        type=int,
        default=0,
        help="Partition index for this run (0-based).",
    )
    parser.add_argument("--val-ratio", type=float, default=0.01)
    parser.add_argument("--seed", type=int, default=17)
    parser.add_argument(
        "--chunk-long",
        dest="chunk_long",
        action="store_true",
        default=True,
        help="Split long segments into 6-10s windows for the 16k bucket.",
    )
    parser.add_argument(
        "--no-chunk-long",
        dest="chunk_long",
        action="store_false",
        help="Disable chunking long segments.",
    )
    parser.add_argument("--chunk-min-sec", type=float, default=6.0)
    parser.add_argument("--chunk-max-sec", type=float, default=10.0)
    parser.add_argument("--chunk-threshold-sec", type=float, default=12.0)
    parser.add_argument(
        "--chunk-all-buckets",
        action="store_true",
        help="Also chunk high-rate buckets (default: keep full segments).",
    )
    parser.add_argument(
        "--silence-snap",
        dest="silence_snap",
        action="store_true",
        default=True,
        help="Snap chunk boundaries to low-energy points when possible.",
    )
    parser.add_argument(
        "--no-silence-snap",
        dest="silence_snap",
        action="store_false",
        help="Disable silence-aware chunk boundaries.",
    )
    parser.add_argument("--silence-frame-ms", type=float, default=50.0)
    parser.add_argument("--silence-hop-ms", type=float, default=10.0)
    parser.add_argument("--silence-search-sec", type=float, default=0.5)
    parser.add_argument("--silence-percentile", type=float, default=20.0)
    parser.add_argument("--silence-db-drop", type=float, default=20.0)
    parser.add_argument("--upload", action="store_true")
    parser.add_argument("--verify-upload", action="store_true")
    parser.add_argument("--create-prefixes", action="store_true")
    parser.add_argument("--r2-subprefix", default="")
    parser.add_argument(
        "--profile-timing",
        action="store_true",
        help="Enable detailed timing instrumentation for profiling.",
    )
    parser.add_argument(
        "--output-format",
        choices=["flac", "wav"],
        default="flac",
        help="Output audio format for shards (default: flac).",
    )
    return parser.parse_args()

def main() -> None:
    global _timing_enabled
    args = parse_args()
    _timing_enabled = args.profile_timing
    env = load_env(ENV_PATH)
    db_url = env.get("DATABASE_URL")
    if not db_url:
        raise SystemExit("DATABASE_URL not found in .env")

    cleaned_bucket = env.get("R2_ClEANED_BUCKET") or env.get("R2_CLEANED_BUCKET")
    shard_bucket = env.get("R2_BUCKET_SHARDS")
    prefix = env.get("R2_PREFIX")
    if not cleaned_bucket:
        raise SystemExit("R2 cleaned bucket not found in .env")

    if args.upload and not shard_bucket:
        raise SystemExit("R2_BUCKET_SHARDS not found in .env")

    if args.chunk_min_sec <= 0 or args.chunk_max_sec <= 0:
        raise SystemExit("Chunk durations must be positive.")
    if args.chunk_min_sec > args.chunk_max_sec:
        raise SystemExit("chunk-min-sec must be <= chunk-max-sec.")
    if args.chunk_threshold_sec < args.chunk_min_sec:
        raise SystemExit("chunk-threshold-sec must be >= chunk-min-sec.")
    if args.silence_frame_ms <= 0 or args.silence_hop_ms <= 0:
        raise SystemExit("Silence frame/hop must be positive.")
    if args.silence_search_sec <= 0:
        raise SystemExit("silence-search-sec must be positive.")
    if not 0.0 < args.silence_percentile < 100.0:
        raise SystemExit("silence-percentile must be between 0 and 100.")
    if args.silence_db_drop <= 0:
        raise SystemExit("silence-db-drop must be positive.")

    if args.partition_count < 1:
        raise SystemExit("partition-count must be >= 1.")
    if not 0 <= args.partition_index < args.partition_count:
        raise SystemExit("partition-index must be between 0 and partition-count - 1.")

    output_format = args.output_format
    audio_ext = "wav" if output_format == "wav" else "flac"

    partition_tag = ""
    if args.partition_count > 1:
        partition_tag = f"p{args.partition_index:02d}"

    chunk_cfg = ChunkConfig(
        enabled=args.chunk_long,
        min_sec=args.chunk_min_sec,
        max_sec=args.chunk_max_sec,
        threshold_sec=args.chunk_threshold_sec,
        all_buckets=args.chunk_all_buckets,
    )
    silence_cfg = SilenceConfig(
        enabled=args.silence_snap,
        frame_sec=args.silence_frame_ms / 1000.0,
        hop_sec=args.silence_hop_ms / 1000.0,
        search_sec=args.silence_search_sec,
        percentile=args.silence_percentile,
        db_drop=args.silence_db_drop,
    )
    language = args.language.strip().lower()
    with timer("db_fetch_records"):
        records = fetch_video_records(
            db_url,
            language,
            args.limit,
            args.partition_count,
            args.partition_index,
        )
    if not records:
        raise SystemExit(f"No records found for language={language}")

    s3 = make_r2_client(env)
    if args.create_prefixes and args.upload:
        create_prefix_placeholders(
            s3,
            shard_bucket,
            LANGUAGES,
            [16, 24, 32, 44],
            ["manifests", "stats", "reports", "checks"],
            args.r2_subprefix or None,
        )

    tmp_dir = args.tmp_dir
    output_dir = args.output_dir
    if partition_tag:
        if tmp_dir.name != partition_tag:
            tmp_dir = tmp_dir / partition_tag
        if output_dir.name != partition_tag:
            output_dir = output_dir / partition_tag
    tmp_dir.mkdir(parents=True, exist_ok=True)
    download_dir = tmp_dir / "downloads"
    download_dir.mkdir(parents=True, exist_ok=True)

    config = TransferConfig(
        multipart_threshold=8 * 1024 * 1024,
        multipart_chunksize=64 * 1024 * 1024,
        max_concurrency=max(4, args.download_workers),
        use_threads=True,
    )

    def download_record(record: VideoRecord) -> Path:
        with timer("r2_download_total"):
            resolved = resolve_r2_key(s3, cleaned_bucket, record.r2_tar_key, prefix)
            if not resolved:
                raise FileNotFoundError(f"Missing tar in R2: {record.r2_tar_key}")
            local_path = download_dir / f"{record.youtube_id}.tar"
            if not local_path.exists():
                s3.download_file(cleaned_bucket, resolved, str(local_path), Config=config)
            return local_path

    download_paths: dict[str, Path] = {}
    if args.skip_downloads:
        local_dir = args.local_tars_dir or download_dir
        if not local_dir.exists():
            raise SystemExit(f"Local tar directory not found: {local_dir}")
        missing = 0
        for record in tqdm(records, desc="local"):
            local_path = local_dir / f"{record.youtube_id}.tar"
            if local_path.exists():
                download_paths[record.youtube_id] = local_path
            else:
                missing += 1
        if missing:
            print(f"Missing local tars: {missing} (skipping)")
    else:
        with ThreadPoolExecutor(max_workers=args.download_workers) as executor:
            futures = {
                executor.submit(download_record, record): record.youtube_id
                for record in records
            }
            for future in tqdm(as_completed(futures), total=len(futures), desc="download"):
                youtube_id = futures[future]
                try:
                    download_paths[youtube_id] = future.result()
                except Exception as exc:
                    print(f"download failed {youtube_id}: {exc}")

    if not download_paths:
        raise SystemExit("No downloads succeeded.")

    video_ids = list(download_paths.keys())
    random.seed(args.seed)
    val_count = max(1, int(len(video_ids) * args.val_ratio)) if len(video_ids) > 1 else 0
    val_ids = set(random.sample(video_ids, val_count)) if val_count > 0 else set()

    manifest_dir = output_dir / "manifests"
    stats_dir = output_dir / "stats"
    manifest_dir.mkdir(parents=True, exist_ok=True)
    stats_dir.mkdir(parents=True, exist_ok=True)

    shard_prefix = "shard" if not partition_tag else f"shard-{partition_tag}"
    writers = {
        16000: ShardWriter(
            output_dir / f"{language}_16",
            shard_prefix,
            args.target_shard_size_mb * 1024 * 1024,
            audio_ext,
        ),
        24000: ShardWriter(
            output_dir / f"{language}_24",
            shard_prefix,
            args.target_shard_size_mb * 1024 * 1024,
            audio_ext,
        ),
        32000: ShardWriter(
            output_dir / f"{language}_32",
            shard_prefix,
            args.target_shard_size_mb * 1024 * 1024,
            audio_ext,
        ),
        44100: ShardWriter(
            output_dir / f"{language}_44",
            shard_prefix,
            args.target_shard_size_mb * 1024 * 1024,
            audio_ext,
        ),
    }

    manifest_suffix = f"_{partition_tag}" if partition_tag else ""
    train_manifest = (manifest_dir / f"{language}_train{manifest_suffix}.jsonl").open(
        "w", encoding="utf-8"
    )
    val_manifest = (manifest_dir / f"{language}_val{manifest_suffix}.jsonl").open(
        "w", encoding="utf-8"
    )

    stats = {
        "videos": 0,
        "segments_total": 0,
        "segments_kept": {16000: 0, 24000: 0, 32000: 0, 44100: 0},
        "dropped_too_short": 0,
        "dropped_too_long": 0,
        "dropped_clipped": 0,
        "segments_chunked": 0,
        "chunks_emitted": 0,
        "native_sr_buckets": {},
    }

    segment_workers = max(1, args.segment_workers)
    max_inflight = (
        args.segment_max_inflight
        if args.segment_max_inflight > 0
        else max(128, segment_workers * 2)
    )
    pending: set[Any] = set()
    future_info: dict[Any, bool] = {}

    def handle_future(future) -> None:
        is_val = future_info.pop(future, False)
        outputs, seg_stats = future.result()
        stats["dropped_too_short"] += seg_stats["dropped_too_short"]
        stats["dropped_too_long"] += seg_stats["dropped_too_long"]
        stats["dropped_clipped"] += seg_stats["dropped_clipped"]
        stats["segments_chunked"] += seg_stats.get("segments_chunked", 0)
        stats["chunks_emitted"] += seg_stats.get("chunks_emitted", 0)
        for out in outputs:
            stats["segments_kept"][out.sample_rate] += 1
            if out.sample_rate == 16000:
                bucket = out.metadata.get("native_sr_bucket")
                if bucket:
                    stats["native_sr_buckets"][bucket] = (
                        stats["native_sr_buckets"].get(bucket, 0) + 1
                    )
            writer = writers[out.sample_rate]
            with timer("shard_write"):
                shard_name = writer.add_sample(
                    out.metadata["id"], out.audio_bytes, out.metadata
                )
            manifest_entry = dict(out.metadata)
            manifest_entry["shard"] = f"{shard_name}.tar"
            manifest_entry["audio"] = f"{out.metadata['id']}.{audio_ext}"
            line = json.dumps(manifest_entry) + "\n"
            if is_val:
                val_manifest.write(line)
            else:
                train_manifest.write(line)

    def drain_completed() -> None:
        nonlocal pending
        if not pending:
            return
        done, pending = wait(pending, return_when=FIRST_COMPLETED)
        for future in done:
            handle_future(future)

    with ThreadPoolExecutor(max_workers=segment_workers) as executor:
        for video_id, tar_path in tqdm(download_paths.items(), desc="process"):
            stats["videos"] += 1
            is_val = video_id in val_ids
            with timer("tar_open"):
                tf = tarfile.open(tar_path, "r")
            with tf:
                with timer("tar_meta_load"):
                    meta = json.loads(
                        tf.extractfile(f"{video_id}/metadata.json").read()
                    )
                usable_segments = [
                    s for s in meta.get("segments", []) if s.get("status") == "usable"
                ]
                # Create lookup by (speaker, start) for robust matching
                # The filename contains 'start' timestamp (relative to intro-skipped audio),
                # not 'original_start' (absolute time in original video)
                with timer("tar_build_segment_map"):
                    segment_map = {
                        (seg.get("speaker"), round(seg.get("start", 0.0), 2)): seg
                        for seg in usable_segments
                    }

                with timer("tar_list_members"):
                    members = [m for m in tf.getmembers() if m.name.endswith(".flac")]
                stats["segments_total"] += len(members)

                for member in members:
                    parsed = parse_segment_filename(member.name)
                    if not parsed:
                        continue
                    speaker, start_time = parsed
                    # Match by (speaker, start) with rounding for float precision
                    seg_meta = segment_map.get((speaker, round(start_time, 2)))
                    if not seg_meta:
                        # Try fuzzy match for edge cases (tolerance 0.1s)
                        matches = [
                            s for s in usable_segments
                            if s.get("speaker") == speaker
                            and abs(s.get("start", 0.0) - start_time) < 0.1
                        ]
                        if not matches:
                            continue
                        seg_meta = matches[0]

                    while len(pending) >= max_inflight:
                        drain_completed()

                    with timer("tar_extract_segment"):
                        segment_bytes = tf.extractfile(member).read()
                    # Extract index from filename for ID generation (cosmetic only)
                    filename_parts = member.name.split("_")
                    try:
                        index = int(filename_parts[-2])
                    except (ValueError, IndexError):
                        index = 0
                    rng_seed = make_segment_seed(args.seed, video_id, speaker, index)

                    future = executor.submit(
                        process_segment,
                        video_id,
                        language,
                        segment_bytes,
                        seg_meta,
                        speaker,
                        index,
                        chunk_cfg,
                        silence_cfg,
                        rng_seed,
                        output_format,
                    )
                    future_info[future] = is_val
                    pending.add(future)

        while pending:
            drain_completed()

    for writer in writers.values():
        writer.close()
    train_manifest.close()
    val_manifest.close()

    stats_suffix = f"_{partition_tag}" if partition_tag else ""
    stats_path = stats_dir / f"{language}_stats{stats_suffix}.json"
    stats_path.write_text(json.dumps(stats, indent=2))

    if args.upload:
        r2_subprefix = args.r2_subprefix.strip("/")
        upload_map = {
            16000: f"{language}_16",
            24000: f"{language}_24",
            32000: f"{language}_32",
            44100: f"{language}_44",
        }
        for sr, local_prefix in upload_map.items():
            remote_prefix = local_prefix if not r2_subprefix else f"{local_prefix}/{r2_subprefix}"
            upload_directory(
                s3,
                shard_bucket,
                output_dir / local_prefix,
                remote_prefix,
                args.download_workers,
            )
        upload_directory(
            s3,
            shard_bucket,
            manifest_dir,
            "manifests" if not r2_subprefix else f"manifests/{r2_subprefix}",
            args.download_workers,
        )
        upload_directory(
            s3,
            shard_bucket,
            stats_dir,
            "stats" if not r2_subprefix else f"stats/{r2_subprefix}",
            args.download_workers,
        )

        if args.verify_upload:
            verify_dir = tmp_dir / "verify"
            verify_dir.mkdir(parents=True, exist_ok=True)
            for sr, local_prefix in upload_map.items():
                shard_dir = output_dir / local_prefix
                shard_files = sorted(p.stem for p in shard_dir.glob("*.tar"))
                if not shard_files:
                    continue
                shard_name = shard_files[0]
                remote_prefix = local_prefix if not r2_subprefix else f"{local_prefix}/{r2_subprefix}"
                result = verify_uploaded_shard(
                    s3,
                    shard_bucket,
                    remote_prefix,
                    shard_name,
                    sr,
                    audio_ext,
                    verify_dir,
                )
                print(
                    f"verify {language} {sr}k: "
                    f"checked={result['checked_samples']} "
                    f"missing_or_mismatch={result['missing_or_mismatch']} "
                    f"entries={result['total_entries']}"
                )

    if _timing_enabled:
        print_timing_stats()

    print(f"Done. Stats written to {stats_path}")

if __name__ == "__main__":
    main()