"""
src/step2_extract_features.py

Step 2: Pre-extract and Cache Audio Features from Genuine Pretrained 300M Audio Encoder
Loads the official google/embeddinggemma-2 audio_tower (304.8M) + embed_audio (0.8M)
+ language_model (271.0M) in torch.bfloat16 and freezes all parameters.
Includes Lecture Room Augmentation (Ambient Noise, p=0.4).
Supports three streaming data sources:
- Mode A: Hugging Face Streaming (Bingsu/zeroth-korean)
- Mode B: AI Hub ZIP In-Memory Reader (KsponSpeech / Lecture speech)
- Mode C: AI Hub Relay Reader (100GB split zip volumes TS.z01, TS.z02, ... or folders)
Slices 768d output to 256d (MRL) and caches bfloat16 .pt files in data/cache_features/.

Optimizations (100% Bit-Exact & Training-Compatible):
1. 3-Stage Asynchronous Pipeline:
   - Stage 1 (CPU Producer Thread): Archive streaming + resampling + noise augmentation + Mel-spectrogram extraction & bucket padding + early CTC length validation (T >= L).
   - Stage 2 (Main MPS GPU Thread): Dedicated 100% to GPU forward pass without CPU/disk stalls or per-sample empty_cache() thrashing.
   - Stage 3 (Async Disk Writer Thread): Non-blocking atomic .pt.tmp -> .pt disk writes.
2. Bit-Exact Model & Kernel Acceleration:
   - Bypasses redundant AutoTokenizer string construction/parsing per audio sample and extracts only valid Mel frames via zero-copy strided STFT.
   - Caches 5D boolean Audio Attention Masks, 4D LM Full/Sliding Masks, RoPE (cos, sin), and prompt token embeddings per (n_audio, n_audio_pad) bucket.
   - Fuses Q/K/V projections across all 12 Audio Attention layers and all 24 LM Attention layers into single contiguous GEMMs, and fuses Gate/Up projections across all 24 LM MLP layers.
   - Replaces Tensor-bound clamps in Gemma4ClippableLinear with direct scalar-bound clamps + F.linear, and dispatches native F.scaled_dot_product_attention directly.
   - Supports dynamic same-bucket GPU batching (--batch_size) to amortize 577M model weight reads across multiple samples sharing the same 64-token bucket.
   - Precomputes frozen Gemma4AudioAttention relative key projections and softplus scales across all 12 audio layers.
   - Pre-caches float32 RMSNorm weights with native torch.rsqrt across all 302 RMSNorm modules and eliminates redundant 1e10 activation clamps in eval mode.
   - Pre-slices LM final embedding_projection weights to 256 dims (MRL) before D2H copy.
3. O(1) In-Memory Set Resume & Multi-Volume Split ZIP Scanner:
   - Eliminates hundreds of thousands of disk stat() calls and SQLite lookups when resuming.
   - Continuously scans across split ZIP boundaries (e.g. TS.z02) where leading bytes contain continuation data from the previous volume.
"""

from __future__ import annotations

import argparse
import datetime
import gc
import io
import json
import math
import os
import queue
import random
import struct
import sys
import threading
import time
import zipfile
import zlib
from collections import defaultdict
from pathlib import Path
from typing import Generator, Tuple, Dict, Any, Optional, Iterator, Set, List

# Auto-detect and switch to project virtualenv if running outside of it
_project_root = Path(__file__).resolve().parent.parent
_venv_python = _project_root / "venv" / "bin" / "python"
if _venv_python.exists():
    try:
        if Path(sys.executable).resolve() != _venv_python.resolve():
            print(f"🔄 [Auto-Env] Switching Python interpreter:\n    System:  {sys.executable}\n    Virtual: {_venv_python}\n", flush=True)
            os.execv(str(_venv_python), [str(_venv_python)] + sys.argv)
    except Exception as _e:
        print(f"⚠️ [Auto-Env] Warning: Could not auto-switch to venv: {_e}", flush=True)

_src_dir = str(Path(__file__).resolve().parent)
_proj_dir = str(Path(__file__).resolve().parent.parent)
if _src_dir not in sys.path:
    sys.path.insert(0, _src_dir)
if _proj_dir not in sys.path:
    sys.path.insert(0, _proj_dir)

# Strictly cap macOS MPS Unified Memory cache to prevent swap explosion
os.environ.setdefault("PYTORCH_MPS_HIGH_WATERMARK_RATIO", "0.7")
os.environ.setdefault("PYTORCH_MPS_LOW_WATERMARK_RATIO", "0.5")

import numpy as np
import soundfile as sf
import torch
import torch.nn.functional as F
import torchaudio
import torchaudio.functional as AF
import torchaudio.transforms as T
import yaml
from tqdm import tqdm

from dataset import normalize_text


# ==============================================================================
# 1. Lecture Room Augmenter (Reverb & Ambient Noise)
# ==============================================================================

class LectureRoomAugmenter:
    """
    Simulates lecture hall ambient acoustics:
    - Probabilistic ambient HVAC / audience background noise with controllable SNR
    - (Reverb / RIR removed per configuration)
    """
    def __init__(
        self,
        sample_rate: int = 16000,
        rir_prob: float = 0.0,
        noise_prob: float = 0.4,
        snr_db_low: float = 15.0,
        snr_db_high: float = 30.0
    ):
        self.sr = sample_rate
        self.rir_prob = 0.0
        self.noise_prob = noise_prob
        self.snr_db_low = snr_db_low
        self.snr_db_high = snr_db_high

    def add_background_noise(self, waveform: torch.Tensor, snr_db: float) -> torch.Tensor:
        """Adds ambient white/pink noise matching lecture hall background."""
        signal_power = torch.mean(waveform ** 2) + 1e-12
        snr_linear = 10.0 ** (snr_db / 10.0)
        noise_power = signal_power / snr_linear

        noise = torch.randn_like(waveform)
        noise = 0.7 * noise + 0.3 * torch.roll(noise, 1, dims=-1)
        current_noise_power = torch.mean(noise ** 2) + 1e-12
        scaled_noise = noise * torch.sqrt(noise_power / current_noise_power)
        return waveform + scaled_noise

    def __call__(self, waveform: torch.Tensor) -> torch.Tensor:
        if waveform.dim() == 1:
            waveform = waveform.unsqueeze(0)

        if random.random() < self.noise_prob:
            snr = random.uniform(self.snr_db_low, self.snr_db_high)
            waveform = self.add_background_noise(waveform, snr_db=snr)

        max_val = torch.max(torch.abs(waveform))
        if max_val > 0.95:
            waveform = waveform / (max_val / 0.95)

        return waveform


# ==============================================================================
# 2. Mode A: Hugging Face Streaming Loader (Zeroth-Korean)
# ==============================================================================

def stream_zeroth_korean(
    dataset_name: str = "Bingsu/zeroth-korean",
    split: str = "train",
    max_samples: Optional[int] = None
) -> Generator[Tuple[torch.Tensor, int, str], None, None]:
    """
    Streams audio directly from Hugging Face Datasets without downloading whole dataset to disk.
    Yields (waveform_tensor, sample_rate, transcription_text).
    Uses Audio(decode=False) + soundfile to decode raw bytes cleanly without torchcodec dependency.
    """
    print(f"[Step 2] Initializing Hugging Face Streaming on '{dataset_name}' (split='{split}')...")
    from datasets import load_dataset, Audio

    ds = load_dataset(dataset_name, split=split, streaming=True)
    # Disable torchcodec auto-decoder by receiving raw audio bytes
    ds = ds.cast_column("audio", Audio(decode=False))
    count = 0

    for sample in ds:
        audio_info = sample.get("audio", {})
        raw_bytes = audio_info.get("bytes")
        text = sample.get("text", "").strip()

        if not text:
            continue

        if raw_bytes is not None:
            try:
                data, sr = sf.read(io.BytesIO(raw_bytes), dtype="float32")
                if data.ndim > 1:
                    data = data.mean(axis=1)
                wav = torch.from_numpy(data).unsqueeze(0)
            except Exception:
                continue
        elif "array" in audio_info:
            wav = torch.from_numpy(np.array(audio_info["array"], dtype=np.float32)).unsqueeze(0)
            sr = audio_info.get("sampling_rate", 16000)
        else:
            continue

        yield wav, sr, text

        count += 1
        if max_samples is not None and count >= max_samples:
            break


class PrefetchedStream:
    """
    Asynchronous background prefetcher that pipelines I/O and CPU preprocessing
    with MPS GPU inference, eliminating CPU and disk/network stalls.
    """
    def __init__(self, generator: Iterator, max_queue_size: int = 16):
        self.generator = generator
        self.q: queue.Queue = queue.Queue(maxsize=max_queue_size)
        self.stop_event = threading.Event()
        self.worker = threading.Thread(target=self._produce, daemon=True)
        self.worker.start()

    def _produce(self):
        try:
            for item in self.generator:
                if self.stop_event.is_set():
                    break
                self.q.put(item)
        except Exception as e:
            print(f"[Step 2 Prefetch Warning] Streaming producer exception: {e}")
        finally:
            self.q.put(None)

    def __iter__(self):
        while True:
            item = self.q.get()
            if item is None:
                break
            yield item

    def close(self):
        self.stop_event.set()


class AsyncDiskSaver:
    """
    Asynchronous background disk writer that performs atomic .pt.tmp -> .pt saves
    and optional source file cleanup off the main MPS GPU thread.
    """
    def __init__(self, max_queue_size: int = 32, delete_source_after_extract: bool = False):
        self.q: queue.Queue = queue.Queue(maxsize=max_queue_size)
        self.delete_source = delete_source_after_extract
        self.error: Optional[Exception] = None
        self.worker = threading.Thread(target=self._consume, daemon=True)
        self.worker.start()

    def _consume(self):
        while True:
            item = self.q.get()
            if item is None:
                self.q.task_done()
                break
            payload, cache_file, src_path = item
            try:
                tmp_file = cache_file.with_suffix(".pt.tmp")
                torch.save(payload, tmp_file)
                os.replace(tmp_file, cache_file)
                if self.delete_source and src_path is not None:
                    try:
                        src_path.unlink()
                    except Exception as de:
                        print(f"[Step 2 Warning] Could not delete {src_path}: {de}")
            except Exception as e:
                self.error = e
                print(f"[Step 2 Disk Write Error] Failed to save {cache_file}: {e}")
            finally:
                self.q.task_done()

    def save(self, payload: Dict[str, Any], cache_file: Path, src_path: Optional[Path] = None):
        if self.error is not None:
            raise RuntimeError(f"AsyncDiskSaver encountered fatal write error: {self.error}")
        self.q.put((payload, cache_file, src_path))

    def close(self):
        self.q.put(None)
        self.q.join()
        self.worker.join(timeout=30.0)


# ==============================================================================
# 3. Mode B & C: AI Hub ZIP & Split-Volume Relay Readers
# ==============================================================================

class AIHubZipReader:
    """
    In-memory reader for large AI Hub speech archives (KsponSpeech, Lecture speech).
    Directly extracts and decodes .wav / .pcm audio and corresponding .txt transcripts
    from memory buffer without extracting archives to disk.
    """
    def __init__(self, zip_path: str):
        self.zip_path = Path(zip_path)
        if not self.zip_path.exists():
            raise FileNotFoundError(f"AI Hub ZIP archive not found at {self.zip_path}")

        print(f"[Step 2] Indexing AI Hub ZIP archive: {self.zip_path.name}...")
        self.zf = zipfile.ZipFile(self.zip_path, "r")
        self.file_list = self.zf.namelist()

        self.pairs: Dict[str, Dict[str, str]] = {}
        for fname in self.file_list:
            if fname.endswith("/") or fname.startswith("__MACOSX"):
                continue
            base, ext = os.path.splitext(fname)
            ext = ext.lower()
            if ext in [".wav", ".pcm", ".flac"]:
                self.pairs.setdefault(base, {})["audio"] = fname
            elif ext in [".txt", ".trn"]:
                self.pairs.setdefault(base, {})["text"] = fname

        self.valid_keys = [k for k, v in self.pairs.items() if "audio" in v and "text" in v]
        print(f"[Step 2] Discovered {len(self.valid_keys):,} valid audio-text pairs in ZIP.")

    def __len__(self) -> int:
        return len(self.valid_keys)

    def stream_samples(self, max_samples: Optional[int] = None) -> Generator[Tuple[torch.Tensor, int, str], None, None]:
        keys = self.valid_keys
        if max_samples is not None:
            keys = keys[:max_samples]

        for k in keys:
            audio_fname = self.pairs[k]["audio"]
            text_fname = self.pairs[k]["text"]

            try:
                raw_txt = self.zf.read(text_fname)
                try:
                    text = raw_txt.decode("utf-8").strip()
                except UnicodeDecodeError:
                    text = raw_txt.decode("cp949", errors="ignore").strip()
            except Exception:
                continue

            try:
                raw_audio = self.zf.read(audio_fname)
                if audio_fname.lower().endswith(".pcm"):
                    pcm_arr = np.frombuffer(raw_audio, dtype=np.int16).astype(np.float32) / 32768.0
                    wav = torch.from_numpy(pcm_arr).unsqueeze(0)
                    sr = 16000
                else:
                    audio_buf = io.BytesIO(raw_audio)
                    data, sr = sf.read(audio_buf, dtype="float32")
                    if data.ndim > 1:
                        data = data.mean(axis=1)
                    wav = torch.from_numpy(data).unsqueeze(0)
            except Exception:
                continue

            yield wav, sr, text


class AIHubLectureRelayReader:
    """
    Reader for AI Hub Korean College Lecture dataset in relay extraction mode.
    Matches audio files (.wav in folder or in split zip / zip) with labels from aihub_label_index.db / TL.zip.
    - Uses O(1) in-memory set lookup (existing_files) to skip already cached features with zero disk I/O.
    - Continuously scans across split ZIP volume boundaries (e.g. TS.z02) where leading bytes
      contain compressed data from the previous volume.
    - Strictly skips junk files (*.irx*, *.part, .DS_Store, __MACOSX).
    - Yields (waveform, sr, text, short_key, optional_source_path).
    """
    def __init__(
        self,
        source_dir: str,
        labels_source: Optional[str] = None,
        cache_dir: Optional[str] = None,
        max_samples: Optional[int] = None,
        existing_files: Optional[Set[str]] = None
    ):
        self.source_dir = Path(source_dir)
        self.max_samples = max_samples
        self.cache_dir = Path(cache_dir) if cache_dir else None
        self.existing_files = existing_files

        # Auto-resolve parent Training folder (contains 01.원천데이터 and 02.라벨링데이터)
        if self.source_dir.is_dir():
            candidate_raw = self.source_dir / "01.원천데이터"
            candidate_lbl = self.source_dir / "02.라벨링데이터" / "TL.zip"
            if not labels_source and candidate_lbl.exists():
                labels_source = str(candidate_lbl)
            if candidate_raw.is_dir():
                print(f"[Step 2] Auto-detected AI Hub raw audio folder: {candidate_raw}")
                self.source_dir = candidate_raw

        # Load label index
        if not labels_source:
            default_labels_zip = Path("/Users/yeongchan/Downloads/137.한국어 대학 강의 데이터/01-1.정식개방데이터/Training/02.라벨링데이터/TL.zip")
            labels_source = str(default_labels_zip)

        print(f"[Step 2] Loading AI Hub label index from {labels_source}...")
        from dataset import load_or_build_aihub_label_index
        self.labels_map = load_or_build_aihub_label_index(tl_zip_path=labels_source)
        print(f"[Step 2] Label index ready with {len(self.labels_map):,} entries.")

    @staticmethod
    def _seek_next_local_header(f: io.BufferedReader, initial_bytes: bytes = b"") -> bool:
        """
        Scans forward in 1MB chunks until a valid ZIP local file header (PK\\x03\\x04) is found,
        positioned right at the start of the 30-byte header.
        Returns True if found, False on EOF or Central Directory (PK\\x01\\x02).
        """
        buf = initial_bytes
        chunk_size = 1024 * 1024  # 1MB scan chunks
        while True:
            pos = 0
            while True:
                idx = buf.find(b"PK\x03\x04", pos)
                if idx == -1:
                    break
                # Ensure we have at least 30 bytes to validate header fields
                if len(buf) - idx < 30:
                    extra = f.read(30)
                    buf += extra
                    if len(buf) - idx < 30:
                        return False
                hdr_cand = buf[idx : idx + 30]
                _, ver, flag, method, _, _, _, _, _, fname_len, extra_len = struct.unpack("<4sHHHHHIIIHH", hdr_cand)
                if method in (0, 8) and 0 < fname_len <= 512 and extra_len <= 4096:
                    # Valid local file header found; seek exactly to its start
                    f.seek(f.tell() - len(buf) + idx)
                    return True
                pos = idx + 4

            # Check if we reached Central Directory at the end of the final volume
            if b"PK\x01\x02" in buf:
                cd_idx = buf.find(b"PK\x01\x02")
                if cd_idx != -1:
                    hdr_check = buf[cd_idx : cd_idx + 4]
                    if hdr_check == b"PK\x01\x02":
                        return False

            chunk = f.read(chunk_size)
            if not chunk:
                return False
            buf = buf[-30:] + chunk if len(buf) >= 30 else buf + chunk

    def stream_samples(self) -> Generator[Tuple[torch.Tensor, int, str, str, Optional[Path]], None, None]:
        if not self.source_dir.exists():
            print(f"[Step 2 Warning] Source dir {self.source_dir} does not exist.")
            return

        import re
        def _is_zip_archive(name: str) -> bool:
            nl = name.lower()
            return nl.endswith(".zip") or bool(re.match(r"^.+\.z\d{2,3}$", nl))

        # 1. Discover split zip volumes / zip files (*.z01, *.z02, *.zip, etc.), strictly skipping .irx, .part, .DS_Store
        archive_candidates = []
        if self.source_dir.is_file():
            archive_candidates.append(self.source_dir)
        elif self.source_dir.is_dir():
            for f in sorted(self.source_dir.iterdir()):
                fname = f.name
                if fname.startswith(".") or ".irx" in fname or ".part" in fname:
                    continue
                if _is_zip_archive(fname):
                    archive_candidates.append(f)

        count = 0

        # Stream from archives (split zip / zip) without disk extraction
        for arch in archive_candidates:
            print(f"[Step 2] Streaming audio from archive: {arch.name}...")
            try:
                with open(arch, "rb", buffering=4 * 1024 * 1024) as f:
                    sig = f.read(4)
                    if sig != b"PK\x07\x08":
                        f.seek(0)

                    while True:
                        hdr = f.read(30)
                        if len(hdr) < 30:
                            break
                        if not hdr.startswith(b"PK\x03\x04"):
                            if not self._seek_next_local_header(f, hdr):
                                break
                            continue

                        sig, ver, flag, method, t_val, d_val, crc, comp_size, uncomp_size, fname_len, extra_len = struct.unpack("<4sHHHHHIIIHH", hdr)
                        if method not in (0, 8) or fname_len == 0 or fname_len > 512 or extra_len > 4096:
                            if not self._seek_next_local_header(f, hdr[4:]):
                                break
                            continue

                        fname_bytes = f.read(fname_len)
                        if len(fname_bytes) < fname_len:
                            break
                        if extra_len > 0:
                            f.seek(extra_len, os.SEEK_CUR)
                        fname = fname_bytes.decode("utf-8", errors="replace")

                        if not fname.lower().endswith((".wav", ".pcm", ".flac")) or fname.startswith("__MACOSX"):
                            if comp_size > 0:
                                f.seek(comp_size, os.SEEK_CUR)
                            continue

                        # Extract relative key and short key
                        rel_key = fname.rsplit(".", 1)[0]
                        parts = rel_key.split("/")
                        short_key = f"{parts[-2]}_{parts[-1]}" if len(parts) >= 2 else parts[-1]
                        target_pt_name = f"feat_aihub_{short_key}.pt"

                        # Fast O(1) in-memory Resume check: skip reading audio & SQLite lookup if already cached!
                        if self.existing_files is not None:
                            if target_pt_name in self.existing_files:
                                f.seek(comp_size, os.SEEK_CUR)
                                continue
                        elif self.cache_dir and (self.cache_dir / target_pt_name).exists():
                            f.seek(comp_size, os.SEEK_CUR)
                            continue

                        # Silent skip if label not in index (short_key is already the canonical key)
                        text = self.labels_map.get(short_key)
                        if not text:
                            f.seek(comp_size, os.SEEK_CUR)
                            continue

                        try:
                            comp_data = f.read(comp_size)
                            if len(comp_data) < comp_size:
                                # Truncated entry at the end of a split volume
                                break
                            decomp = zlib.decompress(comp_data, -15) if method == 8 else comp_data
                            data, sr = sf.read(io.BytesIO(decomp), dtype="float32")
                            if data.ndim > 1:
                                data = data.mean(axis=1)
                            wav = torch.from_numpy(data).unsqueeze(0)
                        except Exception:
                            # Silent skip on decode errors (e.g., split boundary)
                            continue

                        yield wav, sr, text, short_key, None
                        count += 1
                        if self.max_samples and count >= self.max_samples:
                            return
            except Exception as e:
                print(f"[Step 2 Warning] Error reading archive {arch}: {e}. Proceeding...")

        # 2. Discover uncompressed standalone .wav files in directory
        if self.source_dir.is_dir():
            import glob
            wav_patterns = [str(self.source_dir / "**/*.wav"), str(self.source_dir / "**/*.pcm")]
            wav_files = []
            for pat in wav_patterns:
                for fpath in glob.glob(pat, recursive=True):
                    if ".irx" not in fpath and ".part" not in fpath and ".DS_Store" not in fpath and "__MACOSX" not in fpath:
                        wav_files.append(Path(fpath))

            wav_files.sort()
            for wf in wav_files:
                try:
                    rel_key = str(wf.relative_to(self.source_dir)).rsplit(".", 1)[0]
                except Exception:
                    rel_key = wf.stem
                parts = rel_key.split("/")
                short_key = f"{parts[-2]}_{parts[-1]}" if len(parts) >= 2 else parts[-1]
                target_pt_name = f"feat_aihub_{short_key}.pt"

                # Fast O(1) Resume check
                if self.existing_files is not None:
                    if target_pt_name in self.existing_files:
                        continue
                elif self.cache_dir and (self.cache_dir / target_pt_name).exists():
                    continue

                # Silent skip if label not found
                text = self.labels_map.get(short_key)
                if not text:
                    continue

                try:
                    data, sr = sf.read(str(wf), dtype="float32")
                    if data.ndim > 1:
                        data = data.mean(axis=1)
                    wav = torch.from_numpy(data).unsqueeze(0)
                except Exception:
                    continue

                yield wav, sr, text, short_key, wf
                count += 1
                if self.max_samples and count >= self.max_samples:
                    return


# ==============================================================================
# 4. Genuine 300M Audio Encoder Feature Extractor (Bit-Exact Accelerated)
# ==============================================================================

class GenuineAudioFeatureExtractor:
    """
    Extracts features using the official google/embeddinggemma-2 multimodal pipeline:
    - Audio processed by audio_tower (304.8M) + embed_audio (0.8M)
    - Passed through language_model (271.0M) into the shared 768-dim multimodal embedding space
    - Audio token representations sliced to 256-dim (MRL) and L2 normalized
    This guarantees audio and text representations share the exact same geometric vector space,
    producing bit-exact identical outputs while accelerating execution via precomputation and caching.
    """
    def __init__(self, device: torch.device, model_name: str = "google/embeddinggemma-2"):
        self.device = device
        self.target_sr = 16000
        self.audio_bucket = 64

        print(f"[Step 2] Loading genuine '{model_name}' Multimodal Model on {device} (bfloat16)...")
        try:
            from transformers import AutoModel, AutoProcessor
            self.processor = AutoProcessor.from_pretrained(model_name)
            self.feature_extractor = self.processor.feature_extractor
            self.model = AutoModel.from_pretrained(model_name, dtype=torch.bfloat16)

            if not hasattr(self.model, "audio_tower") or not hasattr(self.model, "embed_audio") or not hasattr(self.model, "language_model"):
                raise AttributeError(f"Model missing required modules. Available: {[n for n, _ in self.model.named_children()]}")

            # Delete vision_tower to save unified memory
            if hasattr(self.model, "vision_tower"):
                del self.model.vision_tower
            gc.collect()

            if hasattr(self.model, "config"):
                self.model.config.use_cache = False
            if hasattr(self.model, "language_model") and hasattr(self.model.language_model, "config"):
                self.model.language_model.config.use_cache = False

            self.model = self.model.to(device)
            self.model.eval()

            #Strictly freeze all parameters
            for p in self.model.parameters():
                p.requires_grad = False

            total_p = sum(p.numel() for p in self.model.parameters())
            self._apply_bit_exact_optimizations()
            print(f"[Step 2] Successfully loaded, FROZE, and optimized genuine Multimodal Model ({total_p:,} params, {total_p/1e6:.1f}M)!")

            if torch.backends.mps.is_available():
                torch.mps.empty_cache()

        except Exception as e:
            raise RuntimeError(
                f"\n[FATAL ERROR] Failed to load genuine '{model_name}' Multimodal Model: {e}\n"
                f"Underlying error: {e}\n"
                f"Ensure you are running with updated transformers: pip install git+https://github.com/huggingface/transformers.git"
            )

    def _apply_bit_exact_optimizations(self) -> None:
        """
        Applies mathematically and numerically bit-exact optimizations to the frozen
        Audio Tower and Language Model layers:
        1. Pre-caches float32 RMSNorm weights and uses native torch.rsqrt across all 302 RMSNorm modules.
        2. Optimizes Gemma4ClippableLinear with pre-extracted scalar bounds and direct F.linear.
        3. Fuses Q/K/V weights and precomputes relative positional key projections + softplus scales across all 12 Audio Attention layers.
        4. Removes redundant 1e10 activation clamps in eval mode (values are already bounded by ClippableLinear).
        5. Fuses Q/K/V weights + direct SDPA dispatch in all 24 LM Attention layers, and fuses Gate/Up weights in all 24 LM MLP layers.
        6. Precomputes prompt token embeddings ([BOS, BOA], [EOA, EOS], [PAD]) and MRL 256d projection slice.
        """
        from transformers.models.gemma4.modeling_gemma4 import (
            Gemma4RMSNorm,
            Gemma4ClippableLinear,
            Gemma4AudioAttention,
            Gemma4AudioLayer,
            Gemma4AudioFeedForward,
            Gemma4AudioLightConv1d,
            create_bidirectional_mask as g4_create_mask,
            sliding_window_mask_function,
        )
        from transformers.models.embedding_gemma2.modeling_embedding_gemma2 import (
            EmbeddingGemma2RMSNorm,
            EmbeddingGemma2Attention,
            EmbeddingGemma2MLP,
            EmbeddingGemma2TextPLEBlock,
            EmbeddingGemma2EncoderLayer,
            apply_rotary_pos_emb,
            create_bidirectional_mask as eg2_create_mask,
            create_bidirectional_sliding_window_mask,
        )

        self._g4_create_mask = g4_create_mask
        self._sliding_window_mask_fn = sliding_window_mask_function
        self._eg2_create_mask = eg2_create_mask
        self._eg2_create_sliding_mask = create_bidirectional_sliding_window_mask

        # 1. Pre-cache float32 weights in all RMSNorm layers & replace torch.pow(-0.5) with torch.rsqrt
        for m in self.model.modules():
            if isinstance(m, (Gemma4RMSNorm, EmbeddingGemma2RMSNorm)):
                if getattr(m, "with_scale", True) and hasattr(m, "weight"):
                    m.register_buffer("_weight_f32", m.weight.detach().float(), persistent=False)

        def _exact_fast_rmsnorm(norm_self, hidden_states: torch.Tensor) -> torch.Tensor:
            x = hidden_states.float()
            normed = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + norm_self.eps)
            if norm_self.with_scale:
                normed = normed * norm_self._weight_f32
            return normed.type_as(hidden_states)

        Gemma4RMSNorm.forward = _exact_fast_rmsnorm
        EmbeddingGemma2RMSNorm.forward = _exact_fast_rmsnorm

        # 2. Optimize Gemma4ClippableLinear with Python float scalar bounds & direct F.linear
        at = self.model.audio_tower
        lm = self.model.language_model
        for m in at.modules():
            if isinstance(m, Gemma4ClippableLinear):
                m._w = m.linear.weight
                if m.use_clipped_linears:
                    m._imin = float(m.input_min.item())
                    m._imax = float(m.input_max.item())
                    m._omin = float(m.output_min.item())
                    m._omax = float(m.output_max.item())

        def _fast_clippable_linear(c_self, hidden_states: torch.Tensor) -> torch.Tensor:
            if c_self.use_clipped_linears:
                hidden_states = torch.clamp(hidden_states, c_self._imin, c_self._imax)
                return torch.clamp(F.linear(hidden_states, c_self._w), c_self._omin, c_self._omax)
            return F.linear(hidden_states, c_self._w)

        Gemma4ClippableLinear.forward = _fast_clippable_linear

        # 3. Precompute constant relative position embeddings & fuse Q/K/V projections in Audio Tower
        with torch.no_grad():
            self._pos_emb_const = at.rel_pos_enc(
                torch.empty(1, 1, at.config.hidden_size, dtype=torch.bfloat16, device=self.device)
            )
            for layer in at.layers:
                attn = layer.self_attn
                attn.register_buffer("_softplus_scale", F.softplus(attn.per_dim_scale), persistent=False)
                rel_k = (
                    attn.relative_k_proj(self._pos_emb_const)
                    .view(-1, attn.num_heads, attn.head_dim)
                    .float()
                    .permute(1, 2, 0)
                )
                attn.register_buffer("_rel_k_perm", rel_k, persistent=False)

                q_w = attn.q_proj.linear.weight
                k_w = attn.k_proj.linear.weight
                v_w = attn.v_proj.linear.weight
                attn.register_buffer("_qkv_w", torch.cat([q_w, k_w, v_w], dim=0).contiguous(), persistent=False)
                attn._qkv_dim = q_w.shape[0]
                attn._qkv_imin = float(attn.q_proj.input_min.item())
                attn._qkv_imax = float(attn.q_proj.input_max.item())
                attn._q_omin = float(attn.q_proj.output_min.item())
                attn._q_omax = float(attn.q_proj.output_max.item())
                attn._k_omin = float(attn.k_proj.output_min.item())
                attn._k_omax = float(attn.k_proj.output_max.item())
                attn._v_omin = float(attn.v_proj.output_min.item())
                attn._v_omax = float(attn.v_proj.output_max.item())
                attn._invalid_val = float(attn.config.attention_invalid_logits_value)
                del attn.q_proj.linear.weight, attn.k_proj.linear.weight, attn.v_proj.linear.weight

        def _fast_audio_attn_forward(attn_self, hidden_states: torch.Tensor, position_embeddings: torch.Tensor, attention_mask: Optional[torch.Tensor] = None):
            batch_size, seq_length, _ = hidden_states.shape
            hidden_shape = (batch_size, seq_length, attn_self.num_heads, attn_self.head_dim)
            D = attn_self._qkv_dim

            x_c = torch.clamp(hidden_states, attn_self._qkv_imin, attn_self._qkv_imax)
            qkv = F.linear(x_c, attn_self._qkv_w)
            query_states = torch.clamp(qkv[..., :D], attn_self._q_omin, attn_self._q_omax).float().view(hidden_shape)
            key_states = torch.clamp(qkv[..., D:2*D], attn_self._k_omin, attn_self._k_omax).float().view(hidden_shape)
            value_states = torch.clamp(qkv[..., 2*D:], attn_self._v_omin, attn_self._v_omax).float().view(hidden_shape)

            query_states = query_states * attn_self.q_scale * attn_self._softplus_scale
            key_states = key_states * attn_self.k_scale

            query_states = attn_self._convert_to_block(query_states)
            key_states = attn_self._extract_block_context(key_states)
            value_states = attn_self._extract_block_context(value_states)
            num_blocks = query_states.shape[1]

            queries = query_states.permute(0, 3, 1, 2, 4)
            matrix_ac = queries @ key_states.permute(0, 3, 1, 4, 2)

            queries_flat = queries.reshape(batch_size, attn_self.num_heads, -1, attn_self.head_dim)
            matrix_bd = (queries_flat @ attn_self._rel_k_perm).reshape(
                batch_size, attn_self.num_heads, num_blocks, attn_self.chunk_size, -1
            )
            matrix_bd = attn_self._rel_shift(matrix_bd)

            attn_weights = (matrix_ac + matrix_bd) / attn_self.softcap
            attn_weights = torch.tanh(attn_weights) * attn_self.softcap

            if attention_mask is not None:
                attn_weights = attn_weights.masked_fill(attention_mask, attn_self._invalid_val)

            attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(value_states.dtype)
            attn_output = attn_weights @ value_states.permute(0, 3, 1, 2, 4)
            attn_output = attn_output.permute(0, 2, 3, 1, 4).reshape(batch_size, num_blocks * attn_self.chunk_size, -1)
            attn_output = attn_output[:, :seq_length].contiguous()
            return attn_self.post(attn_output.to(hidden_states.dtype)), None

        Gemma4AudioAttention.forward = _fast_audio_attn_forward

        # 4. Eliminate redundant 1e10 clamp kernels in Audio FFN, LightConv1d, and AudioLayer
        def _fast_ffn_forward(ffn_self, hidden_states: torch.Tensor) -> torch.Tensor:
            residual = hidden_states
            hidden_states = ffn_self.pre_layer_norm(hidden_states)
            hidden_states = ffn_self.ffw_layer_2(ffn_self.act_fn(ffn_self.ffw_layer_1(hidden_states)))
            hidden_states = ffn_self.post_layer_norm(hidden_states) * ffn_self.post_layer_scale
            return hidden_states + residual

        def _fast_lconv_forward(conv_self, hidden_states: torch.Tensor) -> torch.Tensor:
            residual = hidden_states
            hidden_states = F.glu(conv_self.linear_start(conv_self.pre_layer_norm(hidden_states)), dim=-1)
            hidden_states = conv_self.depthwise_conv1d(hidden_states.transpose(1, 2)).transpose(1, 2)
            hidden_states = conv_self.linear_end(conv_self.act_fn(conv_self.conv_norm(hidden_states)))
            return hidden_states + residual

        def _fast_layer_forward(layer_self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor], position_embeddings: torch.Tensor, **kwargs) -> torch.Tensor:
            hidden_states = layer_self.feed_forward1(hidden_states)
            residual = hidden_states
            h_attn, _ = layer_self.self_attn(
                hidden_states=layer_self.norm_pre_attn(hidden_states),
                position_embeddings=position_embeddings,
                attention_mask=attention_mask,
            )
            hidden_states = layer_self.norm_post_attn(h_attn) + residual
            hidden_states = layer_self.lconv1d(hidden_states)
            hidden_states = layer_self.feed_forward2(hidden_states)
            return layer_self.norm_out(hidden_states)

        Gemma4AudioFeedForward.forward = _fast_ffn_forward
        Gemma4AudioLightConv1d.forward = _fast_lconv_forward
        Gemma4AudioLayer.forward = _fast_layer_forward

        # 5. Fuse Q/K/V and Gate/Up projections across all 24 Language Model layers + direct SDPA
        with torch.no_grad():
            for layer in lm.layers:
                attn = layer.self_attn
                q_w = attn.q_proj.weight
                k_w = attn.k_proj.weight
                v_w = attn.v_proj.weight
                attn.register_buffer("_qkv_w", torch.cat([q_w, k_w, v_w], dim=0).contiguous(), persistent=False)
                attn._q_dim = q_w.shape[0]
                attn._kv_dim = k_w.shape[0]
                attn._o_w = attn.o_proj.weight
                del attn.q_proj.weight, attn.k_proj.weight, attn.v_proj.weight

                mlp = layer.mlp
                g_w = mlp.gate_proj.weight
                u_w = mlp.up_proj.weight
                mlp.register_buffer("_gu_w", torch.cat([g_w, u_w], dim=0).contiguous(), persistent=False)
                mlp._inter_dim = mlp.intermediate_size
                mlp._down_w = mlp.down_proj.weight
                del mlp.gate_proj.weight, mlp.up_proj.weight

                ple_b = layer.ple_block
                ple_b._gate_w = ple_b.per_layer_input_gate.weight
                ple_b._proj_w = ple_b.per_layer_projection.weight
                layer._scalar_val = float(layer.layer_scalar.item())

        def _fused_lm_mlp_forward(mlp_self, x: torch.Tensor) -> torch.Tensor:
            gu = F.linear(x, mlp_self._gu_w)
            I = mlp_self._inter_dim
            return F.linear(mlp_self.act_fn(gu[..., :I]) * gu[..., I:], mlp_self._down_w)

        def _fused_lm_attn_forward(attn_self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, **kwargs):
            input_shape = hidden_states.shape[:-1]
            hidden_shape = (*input_shape, -1, attn_self.head_dim)
            Dq = attn_self._q_dim
            Dkv = attn_self._kv_dim

            qkv = F.linear(hidden_states, attn_self._qkv_w)
            query_states = attn_self.q_norm(qkv[..., :Dq].view(hidden_shape).transpose(1, 2))
            key_states = attn_self.k_norm(qkv[..., Dq:Dq+Dkv].view(hidden_shape).transpose(1, 2))
            value_states = attn_self.v_norm(qkv[..., Dq+Dkv:].view(hidden_shape).transpose(1, 2))

            cos, sin = position_embeddings
            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

            attn_output = F.scaled_dot_product_attention(
                query_states,
                key_states,
                value_states,
                attn_mask=attention_mask,
                dropout_p=0.0,
                scale=attn_self.scaling,
                is_causal=False,
                enable_gqa=attn_self.num_key_value_groups > 1,
            )
            attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous()
            return F.linear(attn_output, attn_self._o_w), None

        def _fused_ple_block_forward(ple_self, hidden_states: torch.Tensor, per_layer_input: torch.Tensor) -> torch.Tensor:
            h = ple_self.act_fn(F.linear(hidden_states, ple_self._gate_w)) * per_layer_input
            h = ple_self.post_per_layer_input_norm(F.linear(h, ple_self._proj_w))
            return hidden_states + h

        def _fused_lm_layer_forward(layer_self, hidden_states: torch.Tensor, per_layer_input: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, **kwargs) -> torch.Tensor:
            h_attn, _ = layer_self.self_attn(
                hidden_states=layer_self.input_layernorm(hidden_states),
                attention_mask=attention_mask,
                position_embeddings=position_embeddings,
            )
            hidden_states = hidden_states + layer_self.post_attention_layernorm(h_attn)
            h_mlp = layer_self.mlp(layer_self.pre_feedforward_layernorm(hidden_states))
            hidden_states = hidden_states + layer_self.post_feedforward_layernorm(h_mlp)
            hidden_states = layer_self.ple_block(hidden_states, per_layer_input)
            if layer_self._scalar_val != 1.0:
                hidden_states = hidden_states * layer_self.layer_scalar
            return hidden_states

        EmbeddingGemma2MLP.forward = _fused_lm_mlp_forward
        EmbeddingGemma2Attention.forward = _fused_lm_attn_forward
        EmbeddingGemma2TextPLEBlock.forward = _fused_ple_block_forward
        EmbeddingGemma2EncoderLayer.forward = _fused_lm_layer_forward

        # 6. Precompute prompt token embeddings & MRL 256d projection matrix slice
        self.audio_token_id = self.model.config.audio_token_id
        self.pad_id = getattr(self.model.config.text_config, "pad_token_id", 0) or 0

        with torch.no_grad():
            sample_prompt = self.processor(
                text="<|audio|>",
                audio=np.zeros(1600, dtype=np.float32),
                sampling_rate=self.target_sr,
                return_tensors="pt"
            )["input_ids"][0]
            prefix_ids = sample_prompt[:2].unsqueeze(0).to(self.device)
            suffix_ids = sample_prompt[-2:].unsqueeze(0).to(self.device)
            pad_ids = torch.tensor([[self.pad_id]], dtype=torch.long, device=self.device)

            embed_fn = self.model.get_input_embeddings()
            self._prefix_emb = embed_fn(prefix_ids).detach()  # (1, 2, 512)
            self._suffix_emb = embed_fn(suffix_ids).detach()  # (1, 2, 512)
            self._pad_emb = embed_fn(pad_ids).detach()        # (1, 1, 512)
            self._proj_w_256 = lm.embedding_projection.weight[:256, :].detach().contiguous()

        # Pre-cache STFT / Mel parameters for zero-copy CPU Mel extraction
        fe = self.feature_extractor
        self._mel_pad_left = fe.frame_length // 2
        self._mel_frame_length = fe.frame_length
        self._mel_hop_length = fe.hop_length
        self._mel_fft_length = fe.fft_length
        self._mel_window = fe.window
        self._mel_filters = fe.mel_filters
        self._mel_floor = fe.mel_floor

        # Caches for attention masks, base embeddings, and RoPE (cos, sin)
        self._audio_mask_cache: Dict[Tuple[int, int], Optional[torch.Tensor]] = {}
        self._batched_win_cache: Dict[int, Tuple[torch.Tensor, torch.Tensor]] = {}
        self._lm_entry_cache: Dict[Tuple[int, int], Tuple[Dict[str, Optional[torch.Tensor]], torch.Tensor]] = {}
        self._rope_cache: Dict[int, Dict[str, Tuple[torch.Tensor, torch.Tensor]]] = {}
        gc.collect()

    def resample_if_needed(self, wav: torch.Tensor, orig_sr: int) -> torch.Tensor:
        if orig_sr != self.target_sr:
            if not hasattr(self, "_resamplers"):
                self._resamplers: Dict[int, T.Resample] = {}
            resampler = self._resamplers.get(orig_sr)
            if resampler is None:
                resampler = T.Resample(orig_sr, self.target_sr)
                self._resamplers[orig_sr] = resampler
            wav = resampler(wav)
        return wav

    def preprocess_audio_cpu(self, wav: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, int, int]:
        """
        CPU-only preprocessing stage (safe to run in background producer thread):
        Computes Mel-spectrogram features and masks matching Gemma4AudioFeatureExtractor bit-for-bit,
        bypassing redundant string construction, tokenizer calls, and padded STFT frames,
        and pads to 64-token buckets.
        Returns (f_in_bf16, fm_in_bool, n_audio, n_audio_pad).
        """
        if isinstance(wav, torch.Tensor):
            wav_np = wav.detach().cpu().numpy().reshape(-1)
        else:
            wav_np = np.asarray(wav, dtype=np.float32).reshape(-1)

        if wav_np.dtype != np.float32:
            wav_np = wav_np.astype(np.float32)

        raw_len = min(wav_np.shape[0], 480_000)
        hop_length = self._mel_hop_length
        pad_left = self._mel_pad_left
        frame_length = self._mel_frame_length

        valid_frames = max(0, (raw_len - 160 + hop_length - 1) // hop_length)
        n_audio = (valid_frames + 3) // 4
        n_audio_pad = ((n_audio + self.audio_bucket - 1) // self.audio_bucket) * self.audio_bucket
        mel_target = n_audio_pad * 4
        copy_len = min(valid_frames, mel_target)

        f_in = torch.zeros((1, mel_target, 128), dtype=torch.bfloat16)
        fm_in = torch.zeros((1, mel_target), dtype=torch.bool)

        if copy_len > 0:
            needed_len = pad_left + (copy_len - 1) * hop_length + frame_length
            buf = np.zeros(needed_len, dtype=np.float32)
            buf[pad_left : pad_left + raw_len] = wav_np[:raw_len]
            frames = np.lib.stride_tricks.as_strided(
                buf,
                shape=(copy_len, frame_length),
                strides=(buf.strides[0] * hop_length, buf.strides[0]),
            )
            stft = np.fft.rfft(frames * self._mel_window, n=self._mel_fft_length, axis=-1)
            mel_spec = np.matmul(np.abs(stft), self._mel_filters)
            log_mel = np.log(mel_spec + self._mel_floor).astype(np.float32)
            f_in[0, :copy_len] = torch.from_numpy(log_mel).to(torch.bfloat16)
            fm_in[0, :copy_len] = True

        return f_in, fm_in, n_audio, n_audio_pad

    def _get_batched_audio_mask(self, n_t: torch.Tensor, n_audio_pad: int) -> torch.Tensor:
        """
        Builds the 5D boolean blocked local-attention mask (B, 1, num_blocks, chunk_size, context_size)
        directly on device for a same-bucket batch.
        """
        at = self.model.audio_tower
        entry = self._batched_win_cache.get(n_audio_pad)
        if entry is None:
            C = at.config.attention_chunk_size
            L = at.config.attention_context_left - 1
            R = at.config.attention_context_right
            CTX = C + L + R
            NB = (n_audio_pad + C - 1) // C
            b = torch.arange(NB, device=self.device).view(NB, 1, 1)
            c = torch.arange(C, device=self.device).view(1, C, 1)
            j = torch.arange(CTX, device=self.device).view(1, 1, CTX)
            q = b * C + c
            k = b * C + j - L
            dist = q - k
            win = ((dist >= 0) & (dist < L)) | ((dist < 0) & (-dist < R))
            base_valid = win & (q < n_audio_pad) & (k >= 0) & (k < n_audio_pad)
            entry = (base_valid.unsqueeze(0), k.unsqueeze(0))
            self._batched_win_cache[n_audio_pad] = entry
        base_valid, k_idx = entry
        valid = base_valid & (k_idx < n_t.view(-1, 1, 1, 1))
        return (~valid).unsqueeze(1)

    def extract_from_preprocessed(
        self,
        f_in: torch.Tensor,
        fm_in: torch.Tensor,
        n_audio: int,
        n_audio_pad: int
    ) -> torch.Tensor:
        """
        Executes MPS GPU forward pass on pre-padded Mel features using cached masks, RoPE,
        and prompt embeddings, returning the bit-exact (n_audio, 256) bfloat16 tensor on CPU.
        """
        f_mps = f_in.to(self.device, dtype=torch.bfloat16)
        fm_mps = fm_in.to(self.device)
        tok_target = n_audio_pad + 4
        am_key = (n_audio, n_audio_pad)

        at = self.model.audio_tower
        lm = self.model.language_model

        with torch.no_grad():
            # 1. Audio Tower Forward
            hidden_states, output_mask = at.subsample_conv_projection(f_mps, fm_mps)
            if am_key in self._audio_mask_cache:
                audio_attn_mask = self._audio_mask_cache[am_key]
            else:
                m4d = self._g4_create_mask(
                    config=at.config,
                    inputs_embeds=hidden_states,
                    attention_mask=output_mask,
                    and_mask_function=self._sliding_window_mask_fn(
                        (at.config.attention_context_left - 1, at.config.attention_context_right)
                    ),
                )
                audio_attn_mask = (at._convert_4d_mask_to_blocked_5d(m4d) != 0) if m4d is not None else None
                self._audio_mask_cache[am_key] = audio_attn_mask

            for encoder_layer in at.layers:
                hidden_states = encoder_layer(
                    hidden_states,
                    attention_mask=audio_attn_mask,
                    position_embeddings=self._pos_emb_const,
                )
            audio_proj = self.model.embed_audio(inputs_embeds=at.output_proj(hidden_states))

            # 2. Language Model Forward with cached masks, base embeddings, and RoPE
            lm_entry = self._lm_entry_cache.get(am_key)
            if lm_entry is None:
                am_mps = torch.zeros((1, tok_target), dtype=torch.long, device=self.device)
                am_mps[:, : n_audio + 4] = 1
                dummy_emb = torch.empty((1, tok_target, lm.config.hidden_size), dtype=torch.bfloat16, device=self.device)
                mk = {"config": lm.config, "inputs_embeds": dummy_emb, "attention_mask": am_mps}
                am_map = {
                    "full_attention": self._eg2_create_mask(**mk),
                    "sliding_attention": self._eg2_create_sliding_mask(**mk),
                }
                base_embeds = self._pad_emb.expand(1, tok_target, -1).clone()
                base_embeds[:, :2, :] = self._prefix_emb
                base_embeds[:, 2 + n_audio : 4 + n_audio, :] = self._suffix_emb
                lm_entry = (am_map, base_embeds)
                self._lm_entry_cache[am_key] = lm_entry

            am_map, base_embeds = lm_entry
            inputs_embeds = base_embeds.clone()
            inputs_embeds[:, 2 : 2 + n_audio, :] = audio_proj[:, :n_audio, :]

            pos_emb = self._rope_cache.get(tok_target)
            if pos_emb is None:
                pos_ids = torch.arange(tok_target, device=self.device).unsqueeze(0)
                pos_emb = {lt: lm.rotary_emb(inputs_embeds, pos_ids, lt) for lt in lm.unique_layer_types}
                self._rope_cache[tok_target] = pos_emb

            per_layer_inputs = lm.ple(inputs_embeds)
            h = inputs_embeds
            for i, encoder_layer in enumerate(lm.layers):
                lt = lm.config.layer_types[i]
                h = encoder_layer(
                    h,
                    per_layer_inputs[:, :, i, :],
                    attention_mask=am_map[lt],
                    position_embeddings=pos_emb[lt],
                )

            h_normed = lm.norm(h)
            mrl_256 = F.linear(h_normed[:, 2 : 2 + n_audio, :], self._proj_w_256).cpu()[0]

        norm_256 = F.normalize(mrl_256.float(), p=2, dim=-1).to(torch.bfloat16)
        return norm_256

    def extract_batch_same_bucket(
        self,
        batch_items: List[Tuple[torch.Tensor, torch.Tensor, int, int]]
    ) -> List[torch.Tensor]:
        """
        Executes a batched MPS GPU forward pass for multiple preprocessed audio samples that share
        the exact same `n_audio_pad` bucket length (zero cross-bucket padding distortion).
        Falls back to the bit-exact single-sample path when len(batch_items) == 1.
        """
        B = len(batch_items)
        if B == 1:
            return [self.extract_from_preprocessed(*batch_items[0])]

        n_audio_pad = batch_items[0][3]
        n_list = [it[2] for it in batch_items]
        tok_target = n_audio_pad + 4

        f_mps = torch.cat([it[0] for it in batch_items], dim=0).to(self.device, dtype=torch.bfloat16)
        fm_mps = torch.cat([it[1] for it in batch_items], dim=0).to(self.device)
        n_t = torch.tensor(n_list, device=self.device)

        at = self.model.audio_tower
        lm = self.model.language_model

        with torch.no_grad():
            hidden_states, _ = at.subsample_conv_projection(f_mps, fm_mps)
            audio_attn_mask = self._get_batched_audio_mask(n_t, n_audio_pad)

            for encoder_layer in at.layers:
                hidden_states = encoder_layer(
                    hidden_states,
                    attention_mask=audio_attn_mask,
                    position_embeddings=self._pos_emb_const,
                )
            audio_proj = self.model.embed_audio(inputs_embeds=at.output_proj(hidden_states))

            t = torch.arange(tok_target, device=self.device).view(1, tok_target, 1)
            n3 = n_t.view(B, 1, 1)
            inputs_embeds = self._pad_emb.expand(B, tok_target, -1)
            shifted = F.pad(audio_proj, (0, 0, 2, 2))
            inputs_embeds = torch.where((t >= 2) & (t < 2 + n3), shifted, inputs_embeds)
            inputs_embeds = torch.where(t == 0, self._prefix_emb[:, 0:1], inputs_embeds)
            inputs_embeds = torch.where(t == 1, self._prefix_emb[:, 1:2], inputs_embeds)
            inputs_embeds = torch.where(t == 2 + n3, self._suffix_emb[:, 0:1], inputs_embeds)
            inputs_embeds = torch.where(t == 3 + n3, self._suffix_emb[:, 1:2], inputs_embeds)

            am_mps = (torch.arange(tok_target, device=self.device).unsqueeze(0) < (n_t + 4).unsqueeze(1)).long()
            mk = {"config": lm.config, "inputs_embeds": inputs_embeds, "attention_mask": am_mps}
            am_map = {
                "full_attention": self._eg2_create_mask(**mk),
                "sliding_attention": self._eg2_create_sliding_mask(**mk),
            }

            pos_emb = self._rope_cache.get(tok_target)
            if pos_emb is None:
                pos_ids = torch.arange(tok_target, device=self.device).unsqueeze(0)
                pos_emb = {lt: lm.rotary_emb(inputs_embeds[0:1], pos_ids, lt) for lt in lm.unique_layer_types}
                self._rope_cache[tok_target] = pos_emb

            per_layer_inputs = lm.ple(inputs_embeds)
            h = inputs_embeds
            for i, encoder_layer in enumerate(lm.layers):
                lt = lm.config.layer_types[i]
                h = encoder_layer(
                    h,
                    per_layer_inputs[:, :, i, :],
                    attention_mask=am_map[lt],
                    position_embeddings=pos_emb[lt],
                )

            mrl_256 = F.linear(lm.norm(h)[:, 2 : 2 + n_audio_pad, :], self._proj_w_256).cpu()

        return [
            F.normalize(mrl_256[i, :n].float(), p=2, dim=-1).to(torch.bfloat16)
            for i, n in enumerate(n_list)
        ]

    def extract_features(self, wav: torch.Tensor) -> torch.Tensor:
        """
        Extracts genuine (Time, 256) bfloat16 feature tensor from raw waveform in the shared multimodal space.
        Maintains full backward compatibility with standalone callers.
        """
        f_in, fm_in, n_audio, n_audio_pad = self.preprocess_audio_cpu(wav)
        return self.extract_from_preprocessed(f_in, fm_in, n_audio, n_audio_pad)


def format_time_hms(seconds: Optional[float]) -> str:
    """초 단위를 'X시간 Y분 Z초', 'Y분 Z초', 'Z초' 형태로 사람이 읽기 쉽게 변환합니다."""
    if seconds is None or seconds < 0 or not math.isfinite(seconds):
        return "계산 중..."
    s = int(round(seconds))
    h = s // 3600
    m = (s % 3600) // 60
    sec = s % 60
    if h > 0:
        return f"{h}시간 {m:02d}분 {sec:02d}초"
    elif m > 0:
        return f"{m}분 {sec:02d}초"
    else:
        return f"{sec}초"


# ==============================================================================
# 5. Main Execution Entry
# ==============================================================================

def main():
    parser = argparse.ArgumentParser(description="Step 2: Pre-extract and Cache Audio Features from Real 300M Audio Encoder")
    parser.add_argument("--config", type=str, default="configs/config.yaml", help="Path to config file")
    parser.add_argument("--mode", type=str, choices=["streaming", "zip", "aihub_relay"], default="streaming",
                        help="Data source mode: 'streaming' (Zeroth-Korean), 'zip' (AI Hub single zip), or 'aihub_relay' (100GB batch relay)")
    parser.add_argument("--source_dir", type=str, default=None, help="Path to source directory or split zip volumes (e.g. TS.z01)")
    parser.add_argument("--zip_path", type=str, default=None, help="Path to AI Hub ZIP archive (required if mode=zip)")
    parser.add_argument("--audio_dir", type=str, default=None, help="Path to directory containing batch .wav audio files")
    parser.add_argument("--aihub_labels_zip", type=str, default=None, help="Path to AI Hub labeling TL.zip archive")
    parser.add_argument("--aihub_labels_dir", type=str, default=None, help="Path to AI Hub labeling directory containing json files")
    parser.add_argument("--delete_source_after_extract", action="store_true", help="Delete source .wav file after feature extraction to free disk space")
    parser.add_argument("--clean_cache", action="store_true", help="Clean existing cached features in data/cache_features before starting")
    parser.add_argument("--max_samples", type=int, default=None, help="Number of audio samples to extract and cache (default: all available in archive)")
    parser.add_argument("--split", type=str, default="train", help="Dataset split for streaming mode")
    parser.add_argument("--no_augment", action="store_true", help="Disable lecture room acoustic augmentation")
    parser.add_argument("--batch_size", type=int, default=8, help="Max same-bucket GPU batch size (default: 8; set 1 for strict single-sample mode)")
    args = parser.parse_args()

    config_path = Path(args.config)
    if not config_path.is_absolute():
        config_path = Path(__file__).resolve().parent.parent / config_path

    with open(config_path, "r", encoding="utf-8") as f:
        cfg = yaml.safe_load(f)

    base_dir = Path(__file__).resolve().parent.parent
    cache_dir = base_dir / cfg["data"].get("cache_dir", "data/cache_features")

    # Clean previous cached features only if explicitly requested
    if args.clean_cache and cache_dir.exists():
        print(f"[Step 2] Cleaning previous cached features in {cache_dir}...")
        for old_f in cache_dir.glob("*.pt"):
            try:
                old_f.unlink()
            except Exception:
                pass
    cache_dir.mkdir(parents=True, exist_ok=True)

    # Build O(1) in-memory set of existing .pt filenames in a single scandir pass (no extra stat() calls)
    with os.scandir(cache_dir) as it:
        existing_pt_files: Set[str] = {entry.name for entry in it if entry.name.endswith(".pt")}
    start_index = len(existing_pt_files)
    if start_index > 0:
        print(f"[Step 2] Incremental Relay Mode: {start_index:,} existing features indexed in RAM (O(1) resume lookup).")

    # Initialize Augmenter (Reverb disabled; ambient noise only)
    aug_cfg = cfg["audio"].get("augmentation", {})
    do_augment = aug_cfg.get("enabled", True) and not args.no_augment
    augmenter = LectureRoomAugmenter(
        sample_rate=16000,
        rir_prob=0.0,
        noise_prob=aug_cfg.get("noise_prob", 0.4) if do_augment else 0.0,
        snr_db_low=aug_cfg.get("noise_snr_db_low", 15.0),
        snr_db_high=aug_cfg.get("noise_snr_db_high", 30.0)
    )

    # Initialize Genuine Feature Extractor (300M Audio Encoder + 271M LM)
    device_name = cfg["system"].get("device", "mps")
    device = torch.device(device_name if (device_name == "mps" and torch.backends.mps.is_available()) else "cpu")
    extractor = GenuineAudioFeatureExtractor(device=device, model_name=cfg["model"].get("name", "google/embeddinggemma-2"))

    # Determine Data Stream
    if args.mode == "aihub_relay":
        audio_src = args.source_dir or args.audio_dir or args.zip_path
        if not audio_src:
            default_audio_src = Path("/Users/yeongchan/Downloads/137.한국어 대학 강의 데이터/01-1.정식개방데이터/Training/01.원천데이터")
            if default_audio_src.exists():
                audio_src = str(default_audio_src)
            else:
                raise ValueError("--source_dir or --audio_dir required when --mode aihub_relay is specified.")
        labels_src = args.aihub_labels_zip or args.aihub_labels_dir
        if not labels_src:
            default_labels_zip = Path("/Users/yeongchan/Downloads/137.한국어 대학 강의 데이터/01-1.정식개방데이터/Training/02.라벨링데이터/TL.zip")
            if default_labels_zip.exists():
                labels_src = str(default_labels_zip)
            else:
                raise ValueError("AI Hub labels source (TL.zip or folder) must be provided via --aihub_labels_zip or --aihub_labels_dir.")
        relay_reader = AIHubLectureRelayReader(
            source_dir=audio_src,
            labels_source=labels_src,
            cache_dir=str(cache_dir),
            max_samples=args.max_samples,
            existing_files=existing_pt_files
        )
        raw_stream = relay_reader.stream_samples()
    elif args.mode == "zip":
        if not args.zip_path:
            raise ValueError("--zip_path is required when --mode zip is specified.")
        zip_reader = AIHubZipReader(args.zip_path)
        raw_stream = ((w, s, t, None, None) for w, s, t in zip_reader.stream_samples(max_samples=args.max_samples))
    else:
        dataset_name = cfg["data"].get("hf_dataset_name", "Bingsu/zeroth-korean")
        raw_stream = ((w, s, t, None, None) for w, s, t in stream_zeroth_korean(dataset_name=dataset_name, split=args.split, max_samples=args.max_samples))

    # Stage 1 Generator: Perform all CPU-bound work (resampling, noise augmentation,
    # Mel-spectrogram extraction, bucket padding, and early CTC length validation) in the background thread!
    def _cpu_preprocess_generator():
        seq_counter = start_index
        for wav, sr, text, short_key, src_path in raw_stream:
            clean_text = normalize_text(text)
            if not clean_text:
                continue

            num_syllables = len(clean_text.replace(" ", ""))
            if num_syllables == 0:
                continue

            if short_key:
                target_fname = f"feat_aihub_{short_key}.pt"
            else:
                while f"feat_{seq_counter:06d}.pt" in existing_pt_files:
                    seq_counter += 1
                target_fname = f"feat_{seq_counter:06d}.pt"

            if target_fname in existing_pt_files:
                continue

            # 1. Resample to 16,000Hz mono
            wav = extractor.resample_if_needed(wav, orig_sr=sr)
            if wav.dim() > 1 and wav.size(0) > 1:
                wav = wav.mean(dim=0, keepdim=True)

            # 2. Apply Lecture Room Augmentation (Ambient Noise only)
            if do_augment:
                wav = augmenter(wav)

            # 3. Extract Mel-spectrogram & bucket padding on CPU
            f_in, fm_in, n_audio, n_audio_pad = extractor.preprocess_audio_cpu(wav)

            # 4. Early CTC validity check BEFORE running expensive GPU inference (T >= L)
            if n_audio < num_syllables:
                continue

            cache_file = cache_dir / target_fname
            if not short_key:
                seq_counter += 1

            yield f_in, fm_in, n_audio, n_audio_pad, num_syllables, clean_text, cache_file, target_fname, src_path

    # Wrap CPU preprocessing generator in asynchronous prefetcher (Stage 1) and disk saver (Stage 3)
    preprocessed_stream = PrefetchedStream(_cpu_preprocess_generator(), max_queue_size=32)
    disk_saver = AsyncDiskSaver(max_queue_size=64, delete_source_after_extract=args.delete_source_after_extract)

    saved_count = start_index
    new_saved = 0
    max_bs = max(1, int(args.batch_size))
    token_budget = 1536  # Memory-safe token budget for 16GB Unified Memory
    bucket_queues: Dict[int, List[Tuple[Any, ...]]] = defaultdict(list)
    buffered_count = 0
    max_buffered = max_bs * 6

    # Determine total expected extraction count for progress bar & remaining time calculation
    target_new: Optional[int] = None
    if args.max_samples and args.max_samples > 0:
        target_new = args.max_samples
    elif args.mode == "aihub_relay" and "relay_reader" in locals() and hasattr(relay_reader, "labels_map") and relay_reader.labels_map:
        target_new = max(0, len(relay_reader.labels_map) - start_index)
    elif args.mode == "zip" and "zip_reader" in locals() and hasattr(zip_reader, "valid_keys") and zip_reader.valid_keys:
        target_new = max(0, len(zip_reader.valid_keys) - start_index)

    pbar = tqdm(
        total=target_new,
        desc="Extracting Genuine Features",
        dynamic_ncols=True
    )

    start_time = time.time()

    def _flush_bucket(n_audio_pad_key: int) -> bool:
        nonlocal saved_count, new_saved, buffered_count
        items = bucket_queues.get(n_audio_pad_key)
        if not items:
            return False

        bs_cap = max(1, min(max_bs, token_budget // max(n_audio_pad_key, 64)))
        batch = items[:bs_cap]
        del items[:bs_cap]
        buffered_count -= len(batch)

        # Stage 2: Dedicated MPS GPU Forward Pass (Same-Bucket Batched or Bit-Exact Single)
        pre_inputs = [(it[0], it[1], it[2], it[3]) for it in batch]
        feats_list = extractor.extract_batch_same_bucket(pre_inputs)

        for it, features_256 in zip(batch, feats_list):
            _, _, _, _, num_syllables, clean_text, cache_file, target_fname, src_path = it
            if features_256.size(0) < num_syllables:
                continue

            # Stage 3: Enqueue non-blocking atomic disk write
            existing_pt_files.add(target_fname)
            disk_saver.save(
                {
                    "feature": features_256,
                    "raw_text": clean_text
                },
                cache_file,
                src_path
            )

            saved_count += 1
            new_saved += 1
            pbar.update(1)

            if new_saved % 250 == 0:
                gc.collect()
                if torch.backends.mps.is_available():
                    torch.mps.empty_cache()

            if new_saved % 500 == 0:
                import shutil
                free_gb = shutil.disk_usage(cache_dir).free / (1024 ** 3)
                if free_gb < 5.0:
                    print(f"\n⚠️ [Step 2 Disk Safety Alert] Remaining disk space is {free_gb:.2f} GB (< 5GB safety threshold). Gracefully concluding extraction.")
                    return True

            if args.max_samples and args.max_samples > 0 and new_saved >= args.max_samples:
                return True

        return False

    try:
        stop_early = False
        for item in preprocessed_stream:
            n_audio_pad = item[3]
            if max_bs <= 1:
                bucket_queues[n_audio_pad].append(item)
                buffered_count += 1
                if _flush_bucket(n_audio_pad):
                    stop_early = True
                    break
            else:
                bucket_queues[n_audio_pad].append(item)
                buffered_count += 1
                bs_cap = max(1, min(max_bs, token_budget // max(n_audio_pad, 64)))
                if len(bucket_queues[n_audio_pad]) >= bs_cap:
                    if _flush_bucket(n_audio_pad):
                        stop_early = True
                        break
                elif buffered_count >= max_buffered:
                    # Flush the fullest bucket to bound latency and RAM
                    oldest_bucket = max(bucket_queues.keys(), key=lambda k: len(bucket_queues[k]))
                    if _flush_bucket(oldest_bucket):
                        stop_early = True
                        break

        if not stop_early:
            for k in sorted(list(bucket_queues.keys())):
                while bucket_queues[k]:
                    if _flush_bucket(k):
                        break
    finally:
        preprocessed_stream.close()
        disk_saver.close()
        pbar.close()

    total_elapsed = time.time() - start_time
    avg_speed = new_saved / max(total_elapsed, 1e-5)
    print(f"\n[Step 2 Completed Successfully!]")
    print(f"  Newly extracted features: {new_saved:,} (Total in cache: {saved_count:,})")
    print(f"  Total time elapsed: {format_time_hms(total_elapsed)} (Average speed: {avg_speed:.2f} samples/s)")
    print(f"  Target dimension: 256d (bfloat16), Sample rate: 16kHz")


if __name__ == "__main__":
    main()