//! NVIDIA Parakeet TDT 0.6B v3 STT engine via ONNX Runtime (pure Rust).
//!
//! Implements the full NeMo TDT speech-to-text pipeline in Rust, mirroring the
//! reference implementation in `istupakov/onnx-asr`:
//!
//! ```text
//! 16 kHz mono waveform
//!   └─ nemo128.onnx         → mel features    [1, 128, T]
//!        └─ encoder-model   → encoder outputs [1, 1024, T']
//!             └─ decoder_joint (greedy TDT) → token IDs
//!                  └─ vocab.txt              → text
//! ```
//!
//! The TDT decoder is transitive: at each encoder frame it predicts a token and
//! a "duration" (how many frames to advance). A special blank token (id 8192,
//! `<blk>`) means "no token emitted for this frame".
//!
//! A worker task consumes microphone PCM frames, runs RMS-based Voice Activity
//! Detection (VAD), and triggers inference once an end-of-utterance (EOU)
//! silence window is detected.

use crate::config::{EOU_SILENCE_MS, STT_SAMPLE_RATE, VAD_RMS_THRESHOLD};
use anyhow::{Context, Result};
use ndarray::{Array1, ArrayD, IxDyn};
use ort::environment::Environment;
use ort::session::builder::GraphOptimizationLevel;
use ort::session::{Session, SessionInputValue};
use ort::value::{Value, ValueType};
use parking_lot::Mutex;
use std::path::Path;
use std::sync::Arc;
use tokio::sync::mpsc;
use tracing::{debug, error, info};

// --- Model files (relative to the model directory) -------------------------

const CONFIG_FILE: &str = "config.json";
const VOCAB_FILE: &str = "vocab.txt";
const PREPROCESSOR_FILE: &str = "nemo128.onnx";
/// Quantized encoder (preferred when present: ~2-3x faster on CPU, no external
/// data file). Drop this file to fall back to the fp32 encoder.
const ENCODER_INT8_FILE: &str = "encoder-model.int8.onnx";
const ENCODER_FILE: &str = "encoder-model.onnx";
const DECODER_FILE: &str = "decoder_joint-model.onnx";

// --- ONNX tensor names (fixed by the exported graphs) ----------------------

const PP_IN_WAVEFORMS: &str = "waveforms";
const PP_IN_LENS: &str = "waveforms_lens";
const PP_OUT_FEATURES: &str = "features";
const PP_OUT_LENS: &str = "features_lens";

const ENC_IN_SIGNAL: &str = "audio_signal";
const ENC_IN_LENGTH: &str = "length";
const ENC_OUT_OUTPUTS: &str = "outputs";
const ENC_OUT_LENGTHS: &str = "encoded_lengths";

const DEC_IN_ENCODER: &str = "encoder_outputs";
const DEC_IN_TARGETS: &str = "targets";
const DEC_IN_TARGET_LENGTH: &str = "target_length";
const DEC_IN_STATE1: &str = "input_states_1";
const DEC_IN_STATE2: &str = "input_states_2";
const DEC_OUT_OUTPUTS: &str = "outputs";
const DEC_OUT_STATE1: &str = "output_states_1";
const DEC_OUT_STATE2: &str = "output_states_2";

/// A single decoded turn from the STT engine.
#[derive(Debug, Clone)]
pub struct Transcript {
    /// Token IDs emitted by the TDT decoder.
    pub token_ids: Vec<i64>,
    /// Decoded, human-readable transcript.
    pub text: String,
    /// Approximate duration of the utterance in milliseconds.
    pub duration_ms: u64,
}

/// The three ONNX sessions making up the Parakeet TDT pipeline.
struct Sessions {
    /// Keeps the ONNX Runtime environment alive for the lifetime of the sessions.
    _environment: Environment,
    preprocessor: Session,
    encoder: Session,
    decoder_joint: Session,
}

/// Parakeet TDT ONNX pipeline. Cloning is cheap (Arc + Mutex).
#[derive(Clone)]
pub struct ParakeetStt {
    sessions: Arc<Mutex<Sessions>>,
    /// Token id -> piece, with the SentencePiece word marker (`▁`) replaced by a
    /// regular space.
    vocab: Arc<Vec<String>>,
    /// Index of the blank token (`<blk>`).
    blank_idx: i64,
    /// Number of token logits (excludes the trailing duration logits).
    vocab_size: usize,
    /// Maximum number of tokens emitted per encoder frame before forcing an advance.
    max_tokens_per_step: usize,
    /// Encoder hidden size (1024 for Parakeet TDT 0.6B).
    d_model: usize,
    /// Number of decoder prediction-network state layers.
    state_layers: usize,
    /// Hidden size of each decoder state layer.
    state_hidden: usize,
}

impl ParakeetStt {
    /// Load the Parakeet TDT ONNX pipeline from `model_dir`, which must contain
    /// `nemo128.onnx`, an `encoder-model*.onnx`, `decoder_joint-model.onnx`,
    /// `vocab.txt` and `config.json`.
    pub fn load(model_dir: &Path) -> Result<Self> {
        info!(dir = %model_dir.display(), "loading Parakeet TDT ONNX models");

        // `config.json` carries `max_tokens_per_step` (default 10).
        let config_json: serde_json::Value =
            serde_json::from_slice(&std::fs::read(model_dir.join(CONFIG_FILE)).with_context(
                || format!("failed to read {}", model_dir.join(CONFIG_FILE).display()),
            )?)
            .context("invalid Parakeet config.json")?;
        let max_tokens_per_step = config_json
            .get("max_tokens_per_step")
            .and_then(|v| v.as_u64())
            .unwrap_or(10) as usize;

        // Vocabulary: id -> piece, with `▁` (U+2581) mapped to a real space.
        let vocab = Arc::new(load_vocab(&model_dir.join(VOCAB_FILE))?);
        let vocab_size = vocab.len();
        let blank_idx = vocab
            .iter()
            .position(|p| p == "<blk>")
            .map(|i| i as i64)
            .context("vocab.txt does not contain the <blk> token")?;

        // Prefer the int8 encoder when available; fall back to fp32.
        let int8_path = model_dir.join(ENCODER_INT8_FILE);
        let encoder_path = if int8_path.exists() {
            int8_path
        } else {
            model_dir.join(ENCODER_FILE)
        };

        let environment = crate::onnx::environment()?;

        let preprocessor = build_session(&environment, &model_dir.join(PREPROCESSOR_FILE))?;
        let encoder = build_session(&environment, &encoder_path)?;
        let decoder_joint = build_session(&environment, &model_dir.join(DECODER_FILE))?;

        // Introspect the graphs for the dimensions we need at runtime.
        let d_model = output_dim(&encoder, ENC_OUT_OUTPUTS, 1)?;
        let decoder_out_dim = output_dim(&decoder_joint, DEC_OUT_OUTPUTS, 3)?;
        let num_durations = decoder_out_dim
            .checked_sub(vocab_size)
            .context("decoder output is smaller than the vocabulary")?;
        let (state_layers, state_hidden) = input_state_dims(&decoder_joint, DEC_IN_STATE1)?;

        info!(
            encoder = %encoder_path.display(),
            vocab_size,
            blank_idx,
            num_durations,
            d_model,
            state_layers,
            state_hidden,
            "Parakeet TDT pipeline loaded"
        );

        Ok(Self {
            sessions: Arc::new(Mutex::new(Sessions {
                _environment: environment,
                preprocessor,
                encoder,
                decoder_joint,
            })),
            vocab,
            blank_idx,
            vocab_size,
            max_tokens_per_step,
            d_model,
            state_layers,
            state_hidden,
        })
    }

    /// Run the full pipeline on a complete utterance (16 kHz mono float PCM
    /// normalized to [-1, 1]). Returns the emitted token IDs and decoded text.
    pub fn infer(&self, audio: &[f32]) -> Result<(Vec<i64>, String)> {
        let n = audio.len();
        if n == 0 {
            return Ok((Vec::new(), String::new()));
        }

        let mut guard = self.sessions.lock();
        let sessions = &mut *guard;

        // --- 1. Mel-feature preprocessor (nemo128.onnx) ---------------------
        let (feature_shape, feature_data, feature_len) = {
            let waveforms =
                Value::from_array(ArrayD::from_shape_vec(IxDyn(&[1, n]), audio.to_vec())?)
                    .context("failed to build waveform tensor")?;
            let waveforms_lens = Value::from_array(Array1::from_vec(vec![n as i64]))
                .context("failed to build waveform length tensor")?;

            let outputs = sessions
                .preprocessor
                .run(vec![
                    (
                        PP_IN_WAVEFORMS.to_string(),
                        SessionInputValue::from(waveforms),
                    ),
                    (
                        PP_IN_LENS.to_string(),
                        SessionInputValue::from(waveforms_lens),
                    ),
                ])
                .context("nemo128 preprocessor inference failed")?;

            let (shape, data) = outputs
                .get(PP_OUT_FEATURES)
                .context("preprocessor produced no 'features' output")?
                .try_extract_tensor::<f32>()
                .context("failed to extract preprocessor features")?;
            let (_, lens) = outputs
                .get(PP_OUT_LENS)
                .context("preprocessor produced no 'features_lens' output")?
                .try_extract_tensor::<i64>()
                .context("failed to extract preprocessor feature lengths")?;

            let shape: Vec<usize> = shape.iter().map(|&d| d as usize).collect();
            (shape, data.to_vec(), lens[0])
        };

        // --- 2. Encoder -----------------------------------------------------
        // Encoder output layout is [1, d_model, T'] (already transposed by the
        // exported graph). We re-pack it to [T', d_model] so each frame is
        // contiguous for the decoder loop.
        let (frames, encoded_len) = {
            let signal = Value::from_array(
                ArrayD::from_shape_vec(IxDyn(&feature_shape), feature_data)
                    .context("failed to reshape mel features")?,
            )
            .context("failed to build encoder input tensor")?;
            let length = Value::from_array(Array1::from_vec(vec![feature_len]))
                .context("failed to build encoder length tensor")?;

            let outputs = sessions
                .encoder
                .run(vec![
                    (ENC_IN_SIGNAL.to_string(), SessionInputValue::from(signal)),
                    (ENC_IN_LENGTH.to_string(), SessionInputValue::from(length)),
                ])
                .context("encoder inference failed")?;

            let (shape, data) = outputs
                .get(ENC_OUT_OUTPUTS)
                .context("encoder produced no 'outputs'")?
                .try_extract_tensor::<f32>()
                .context("failed to extract encoder outputs")?;
            let (_, lens) = outputs
                .get(ENC_OUT_LENGTHS)
                .context("encoder produced no 'encoded_lengths'")?
                .try_extract_tensor::<i64>()
                .context("failed to extract encoded lengths")?;

            let t = shape[2] as usize;
            let mut frames = vec![0f32; t * self.d_model];
            for tt in 0..t {
                for d in 0..self.d_model {
                    frames[tt * self.d_model + d] = data[d * t + tt];
                }
            }
            // Clamp: low-precision encoders may report more frames than present.
            let encoded_len = (lens[0] as usize).min(t);
            (frames, encoded_len)
        };

        // --- 3. Greedy TDT decoding ----------------------------------------
        let mut tokens: Vec<i64> = Vec::new();
        let mut state1 = vec![0f32; self.state_layers * self.state_hidden];
        let mut state2 = state1.clone();
        let mut prev_token = self.blank_idx;
        let mut emitted = 0usize;
        let mut t = 0usize;

        while t < encoded_len {
            let frame = &frames[t * self.d_model..(t + 1) * self.d_model];

            let enc_frame = Value::from_array(ArrayD::from_shape_vec(
                IxDyn(&[1, self.d_model, 1]),
                frame.to_vec(),
            )?)
            .context("failed to build decoder encoder-frame tensor")?;
            let targets = Value::from_array(ArrayD::from_shape_vec(
                IxDyn(&[1, 1]),
                vec![prev_token as i32],
            )?)
            .context("failed to build decoder targets tensor")?;
            let target_length = Value::from_array(Array1::from_vec(vec![1i32]))
                .context("failed to build decoder target-length tensor")?;
            let s1 = Value::from_array(ArrayD::from_shape_vec(
                IxDyn(&[self.state_layers, 1, self.state_hidden]),
                state1.clone(),
            )?)
            .context("failed to build decoder state1 tensor")?;
            let s2 = Value::from_array(ArrayD::from_shape_vec(
                IxDyn(&[self.state_layers, 1, self.state_hidden]),
                state2.clone(),
            )?)
            .context("failed to build decoder state2 tensor")?;

            let (token, step, new_state1, new_state2) = {
                let outputs = sessions
                    .decoder_joint
                    .run(vec![
                        (
                            DEC_IN_ENCODER.to_string(),
                            SessionInputValue::from(enc_frame),
                        ),
                        (DEC_IN_TARGETS.to_string(), SessionInputValue::from(targets)),
                        (
                            DEC_IN_TARGET_LENGTH.to_string(),
                            SessionInputValue::from(target_length),
                        ),
                        (DEC_IN_STATE1.to_string(), SessionInputValue::from(s1)),
                        (DEC_IN_STATE2.to_string(), SessionInputValue::from(s2)),
                    ])
                    .context("decoder_joint inference failed")?;

                let (_, logits) = outputs
                    .get(DEC_OUT_OUTPUTS)
                    .context("decoder produced no 'outputs'")?
                    .try_extract_tensor::<f32>()
                    .context("failed to extract decoder logits")?;
                if logits.len() < self.vocab_size {
                    anyhow::bail!("decoder logits shorter than vocabulary");
                }
                // Token logits come first, then the duration logits (TDT).
                let token = argmax(&logits[..self.vocab_size]) as i64;
                let step = argmax(&logits[self.vocab_size..]);

                let (_, next_s1) = outputs
                    .get(DEC_OUT_STATE1)
                    .context("decoder produced no 'output_states_1'")?
                    .try_extract_tensor::<f32>()
                    .context("failed to extract decoder state1")?;
                let (_, next_s2) = outputs
                    .get(DEC_OUT_STATE2)
                    .context("decoder produced no 'output_states_2'")?
                    .try_extract_tensor::<f32>()
                    .context("failed to extract decoder state2")?;

                (token, step, next_s1.to_vec(), next_s2.to_vec())
            };

            // Emit the token unless it is the blank symbol.
            if token != self.blank_idx {
                state1 = new_state1;
                state2 = new_state2;
                tokens.push(token);
                prev_token = token;
                emitted += 1;
            }

            // Advance: a positive duration jumps several frames; otherwise we
            // only move on when a blank was emitted or we hit the per-step cap.
            if step > 0 {
                t += step;
                emitted = 0;
            } else if token == self.blank_idx || emitted == self.max_tokens_per_step {
                t += 1;
                emitted = 0;
            }
        }

        let text = decode_text(&self.vocab, &tokens);
        Ok((tokens, text))
    }
}

/// Build a session with full graph optimization.
fn build_session(environment: &Environment, path: &Path) -> Result<Session> {
    let builder = Session::builder(environment)
        .map_err(|e| anyhow::anyhow!("failed to create ONNX session builder: {e}"))?;
    builder
        .with_optimization_level(GraphOptimizationLevel::Level3)
        .map_err(|e| anyhow::anyhow!("failed to set ONNX optimization level: {e}"))?
        .commit_from_file(path)
        .with_context(|| format!("failed to load ONNX model {}", path.display()))
}

/// Read the size of `axis` from the named output of `session`.
fn output_dim(session: &Session, name: &str, axis: usize) -> Result<usize> {
    let outlet = session
        .outputs()
        .iter()
        .find(|o| o.name() == name)
        .with_context(|| format!("model has no output named '{name}'"))?;
    match outlet.dtype() {
        ValueType::Tensor { shape, .. } => {
            let dim = *shape
                .get(axis)
                .with_context(|| format!("output '{name}' has no axis {axis}"))?;
            usize::try_from(dim).with_context(|| format!("output '{name}' axis {axis} is dynamic"))
        }
        _ => anyhow::bail!("output '{name}' is not a tensor"),
    }
}

/// Read `(layers, hidden)` from a decoder state input shaped `[layers, batch, hidden]`.
fn input_state_dims(session: &Session, name: &str) -> Result<(usize, usize)> {
    let outlet = session
        .inputs()
        .iter()
        .find(|i| i.name() == name)
        .with_context(|| format!("decoder has no input named '{name}'"))?;
    match outlet.dtype() {
        ValueType::Tensor { shape, .. } => {
            let layers = usize::try_from(*shape.first().context("state input has no axes")?)
                .context("state input layer count is dynamic")?;
            let hidden = usize::try_from(*shape.get(2).context("state input has < 3 axes")?)
                .context("state input hidden size is dynamic")?;
            Ok((layers, hidden))
        }
        _ => anyhow::bail!("input '{name}' is not a tensor"),
    }
}

/// Index of the maximum value (ties resolved to the lowest index).
fn argmax(values: &[f32]) -> usize {
    let mut best = 0usize;
    let mut best_val = f32::NEG_INFINITY;
    for (i, &v) in values.iter().enumerate() {
        if v > best_val {
            best_val = v;
            best = i;
        }
    }
    best
}

/// Load `vocab.txt` (lines of `<piece> <id>`) into an id-indexed table, mapping
/// the SentencePiece word marker `▁` (U+2581) to a real space.
fn load_vocab(path: &Path) -> Result<Vec<String>> {
    let text = std::fs::read_to_string(path)
        .with_context(|| format!("failed to read {}", path.display()))?;

    let mut entries: Vec<(usize, String)> = Vec::new();
    let mut max_id = 0usize;
    for line in text.lines() {
        if line.is_empty() {
            continue;
        }
        let (piece, id) = line
            .rsplit_once(' ')
            .with_context(|| format!("malformed vocab line: {line:?}"))?;
        let id: usize = id
            .trim()
            .parse()
            .with_context(|| format!("invalid token id in vocab line: {line:?}"))?;
        max_id = max_id.max(id);
        entries.push((id, piece.replace('\u{2581}', " ")));
    }

    let mut vocab = vec![String::new(); max_id + 1];
    for (id, piece) in entries {
        vocab[id] = piece;
    }
    Ok(vocab)
}

/// Decode token IDs to text and normalize SentencePiece spacing.
fn decode_text(vocab: &[String], ids: &[i64]) -> String {
    let mut joined = String::new();
    for &id in ids {
        if id >= 0 {
            if let Some(piece) = vocab.get(id as usize) {
                joined.push_str(piece);
            }
        }
    }
    normalize_spaces(&joined)
}

/// Collapse SentencePiece spacing: drop a leading space, keep a space only when
/// it precedes a word character, and drop spacing before punctuation. This
/// mirrors the `\A\s|\s\B|(\s)\b` substitution used by the reference decoder.
fn normalize_spaces(text: &str) -> String {
    let chars: Vec<char> = text.chars().collect();
    let mut out = String::with_capacity(text.len());
    for (i, &c) in chars.iter().enumerate() {
        if !c.is_whitespace() {
            out.push(c);
            continue;
        }
        if i == 0 {
            continue; // leading whitespace
        }
        let next_is_word = chars
            .get(i + 1)
            .map(|&n| n.is_alphanumeric() || n == '_')
            .unwrap_or(false);
        if next_is_word {
            out.push(' ');
        }
    }
    out
}

/// Input of the STT worker channel.
enum SttInput {
    /// A streaming mic frame: the worker runs its own RMS voice detection.
    Frame(Vec<i16>),
    /// A complete utterance already segmented by an external VAD (see
    /// [`crate::vad`]), which bypasses the internal RMS detection.
    Utterance(Vec<i16>),
}

/// Handle to the running STT worker.
#[derive(Clone)]
pub struct SttWorker {
    tx: mpsc::UnboundedSender<SttInput>,
}

impl SttWorker {
    /// Spawn the STT worker task. `stt` is the loaded Parakeet pipeline.
    /// `transcript_tx` receives one [`Transcript`] per detected EOU.
    pub fn spawn(stt: ParakeetStt, transcript_tx: mpsc::UnboundedSender<Transcript>) -> Self {
        let (tx, mut rx) = mpsc::unbounded_channel::<SttInput>();

        tokio::spawn(async move {
            let mut utterance: Vec<f32> = Vec::new();
            let mut in_utterance = false;
            let mut silence_samples: u64 = 0;
            let eou_samples = (STT_SAMPLE_RATE as u64 * EOU_SILENCE_MS) / 1000;

            while let Some(input) = rx.recv().await {
                // Frames arrive already resampled to 16 kHz by the session layer.
                // Externally segmented utterances need no further detection.
                let frame = match input {
                    SttInput::Utterance(frame) => {
                        let audio: Vec<f32> = frame.iter().map(|&s| s as f32 / 32768.0).collect();
                        debug!(samples = audio.len(), "VAD-segmented utterance received");
                        let stt = stt.clone();
                        let tx = transcript_tx.clone();
                        tokio::task::spawn_blocking(move || {
                            run_inference_blocking(&stt, &audio, tx);
                        });
                        continue;
                    }
                    SttInput::Frame(frame) => frame,
                };

                // Convert i16 -> f32 normalized.
                let normalized: Vec<f32> = frame.iter().map(|&s| s as f32 / 32768.0).collect();
                let rms = compute_rms(&normalized);

                if rms > VAD_RMS_THRESHOLD {
                    in_utterance = true;
                    silence_samples = 0;
                    utterance.extend_from_slice(&normalized);
                } else if in_utterance {
                    silence_samples += normalized.len() as u64;
                    utterance.extend_from_slice(&normalized);
                    if silence_samples >= eou_samples {
                        // EOU detected: run inference on a blocking thread.
                        let utterance_clone = std::mem::take(&mut utterance);
                        let stt = stt.clone();
                        let tx = transcript_tx.clone();

                        tokio::task::spawn_blocking(move || {
                            run_inference_blocking(&stt, &utterance_clone, tx);
                        });

                        in_utterance = false;
                        silence_samples = 0;
                    }
                }
            }
            info!("STT worker channel closed, stopping");
        });

        Self { tx }
    }

    /// Push a frame of 16 kHz mono i16 PCM samples into the worker, which will
    /// segment it with the built-in RMS voice detection.
    pub fn push_frame(&self, samples: Vec<i16>) {
        let _ = self.tx.send(SttInput::Frame(samples));
    }

    /// Push a complete 16 kHz mono i16 utterance segmented by the external VAD.
    /// Transcribed as-is, without any further voice-activity detection.
    pub fn push_utterance(&self, samples: Vec<i16>) {
        if samples.is_empty() {
            return;
        }
        let _ = self.tx.send(SttInput::Utterance(samples));
    }
}

/// Blocking inference helper used inside `spawn_blocking`.
fn run_inference_blocking(stt: &ParakeetStt, audio: &[f32], tx: mpsc::UnboundedSender<Transcript>) {
    let (token_ids, text) = match stt.infer(audio) {
        Ok(result) => result,
        Err(e) => {
            error!(error = %e, "Parakeet inference failed");
            return;
        }
    };

    let duration_ms = (audio.len() as u64 * 1000) / STT_SAMPLE_RATE as u64;
    debug!(tokens = token_ids.len(), duration_ms, text = %text, "EOU transcript ready");

    let _ = tx.send(Transcript {
        token_ids,
        text,
        duration_ms,
    });
}

/// Compute the Root-Mean-Square energy of a normalized f32 frame.
fn compute_rms(samples: &[f32]) -> f32 {
    if samples.is_empty() {
        return 0.0;
    }
    let sum_sq: f32 = samples.iter().map(|&s| s * s).sum();
    (sum_sq / samples.len() as f32).sqrt()
}

/// Convert a duration in milliseconds to a sample count at the STT rate.
#[allow(dead_code)]
pub fn ms_to_samples(ms: u64) -> usize {
    ((STT_SAMPLE_RATE as u64 * ms) / 1000) as usize
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_normalize_spaces() {
        assert_eq!(normalize_spaces(" Hello world"), "Hello world");
        assert_eq!(normalize_spaces(" Hello world!"), "Hello world!");
        assert_eq!(normalize_spaces("a b c"), "a b c");
        assert_eq!(normalize_spaces(" trailing "), "trailing");
        assert_eq!(normalize_spaces("a , b"), "a, b");
    }

    #[test]
    fn test_load_real_vocab() {
        let path = Path::new("models/parakeet/vocab.txt");
        if !path.exists() {
            eprintln!("{} not found, skipping", path.display());
            return;
        }
        let vocab = load_vocab(path).expect("vocab load failed");
        assert_eq!(vocab.len(), 8193);
        assert_eq!(vocab[8192], "<blk>");
        assert_eq!(vocab[0], "<unk>");
    }

    /// Read a canonical 16-bit PCM mono WAV file into normalized f32 samples.
    fn read_wav_mono16(path: &Path) -> Vec<f32> {
        let bytes = std::fs::read(path).expect("failed to read wav");
        assert_eq!(&bytes[0..4], b"RIFF", "not a RIFF file");
        assert_eq!(&bytes[8..12], b"WAVE", "not a WAVE file");
        // Walk the chunk list to find the `data` chunk.
        let mut pos = 12usize;
        while pos + 8 <= bytes.len() {
            let id = &bytes[pos..pos + 4];
            let size = u32::from_le_bytes(bytes[pos + 4..pos + 8].try_into().unwrap()) as usize;
            let body = pos + 8;
            if id == b"data" {
                return bytes[body..body + size]
                    .chunks_exact(2)
                    .map(|c| i16::from_le_bytes([c[0], c[1]]) as f32 / 32768.0)
                    .collect();
            }
            pos = body + size + (size & 1); // chunks are word-aligned
        }
        panic!("no data chunk found");
    }

    /// End-to-end decode test. Requires the full model directory. Run with:
    /// `PARAKEET_TEST_WAV=/path/to/16k_mono.wav cargo test -- --ignored`
    #[test]
    #[ignore = "requires the full Parakeet ONNX model directory + a 16 kHz WAV"]
    fn test_infer_real_audio() {
        let model_dir =
            std::env::var("PARAKEET_MODEL_DIR").unwrap_or_else(|_| "models/parakeet".to_string());
        let wav = std::env::var("PARAKEET_TEST_WAV")
            .unwrap_or_else(|_| "/tmp/stt_test/solis16k.wav".to_string());

        if !Path::new(&model_dir).join(PREPROCESSOR_FILE).exists() {
            eprintln!("model dir {model_dir} not found, skipping");
            return;
        }
        let stt = ParakeetStt::load(Path::new(&model_dir)).expect("model load failed");
        let audio = read_wav_mono16(Path::new(&wav));
        eprintln!("running inference on {} samples", audio.len());

        let (tokens, text) = stt.infer(&audio).expect("inference failed");
        eprintln!("tokens: {} | text: {text:?}", tokens.len());
        assert!(!tokens.is_empty(), "no tokens emitted");
        assert!(!text.trim().is_empty(), "empty transcript");
    }
}