turbospark-runtime 0.3.0

//! Synthetic whisper runner tests: the determinism contract, special-token
//! suppression, stop conditions, and multi-window segmentation, all on
//! deterministic pseudo-random weights (no checkpoint on disk).

#![cfg(target_os = "macos")]

use compute::whisper::WhisperSelfKv;
use model_io::whisper_config::{WhisperConfig, WhisperSpecialTokens};

use crate::whisper::decode::{argmax_suppressed, build_prompt};
use crate::whisper::{WhisperRunner, WINDOW_STRIDE_SAMPLES};
use audio::whisper::{WHISPER_SAMPLE_RATE, WHISPER_WINDOW_SAMPLES};

/// A tiny valid whisper config (the kernels only need consistent shapes).
fn tiny_config() -> WhisperConfig {
    WhisperConfig {
        model_type: "whisper".to_string(),
        d_model: 16,
        num_hidden_layers: Some(2),
        encoder_layers: None,
        decoder_layers: 2,
        num_attention_heads: 4,
        decoder_attention_heads: 4,
        encoder_ffn_dim: 32,
        decoder_ffn_dim: 32,
        vocab_size: 51864,
        n_mels: 80,
        max_source_positions: 1500,
        max_target_positions: 448,
        activation_function: "gelu".to_string(),
        layer_norm_eps: 1e-5,
        scale_embedding: false,
        decoder_start_token_id: Some(50257),
        eos_token_id: Some(50256),
        bos_token_id: Some(50257),
        pad_token_id: Some(50256),
        forced_decoder_ids: None,
        quantization: None,
    }
}

/// The English-only special-token table, as whisper-tiny.en's tokenizer
/// resolves (no language/task specials beyond transcribe).
fn en_tokens() -> WhisperSpecialTokens {
    WhisperSpecialTokens {
        eot: 50256,
        sot: 50257,
        first_language: 50258,
        last_language: 50356,
        translate: None,
        transcribe: 50358,
        startoflm: None,
        no_speech: None,
        no_timestamps: 50362,
        first_timestamp: 50363,
        multilingual: false,
    }
}

fn synthetic_runner() -> WhisperRunner {
    let config = tiny_config();
    let d = config.d_model;
    let mut rng = 20_261_003u64;
    fn lcg(rng: &mut u64) -> f32 {
        *rng = rng
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        ((*rng >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
    }
    fn vec(rng: &mut u64, n: usize) -> Vec<f32> {
        (0..n).map(|_| lcg(rng)).collect()
    }
    let enc_layers = (0..config.encoder_layers())
        .map(|_| crate::whisper::weights::WhisperEncoderLayerOwned {
            q: vec(&mut rng, d * d),
            q_bias: vec(&mut rng, d),
            k: vec(&mut rng, d * d),
            v: vec(&mut rng, d * d),
            v_bias: vec(&mut rng, d),
            out: vec(&mut rng, d * d),
            out_bias: vec(&mut rng, d),
            ln1_weight: vec(&mut rng, d),
            ln1_bias: vec(&mut rng, d),
            fc1: vec(&mut rng, 32 * d),
            fc1_bias: vec(&mut rng, 32),
            fc2: vec(&mut rng, d * 32),
            fc2_bias: vec(&mut rng, d),
            ln2_weight: vec(&mut rng, d),
            ln2_bias: vec(&mut rng, d),
        })
        .collect();
    let dec_layers = (0..config.decoder_layers())
        .map(|_| crate::whisper::weights::WhisperDecoderLayerOwned {
            self_q: vec(&mut rng, d * d),
            self_q_bias: vec(&mut rng, d),
            self_k: vec(&mut rng, d * d),
            self_v: vec(&mut rng, d * d),
            self_v_bias: vec(&mut rng, d),
            self_out: vec(&mut rng, d * d),
            self_out_bias: vec(&mut rng, d),
            ln_self_weight: vec(&mut rng, d),
            ln_self_bias: vec(&mut rng, d),
            cross_q: vec(&mut rng, d * d),
            cross_q_bias: vec(&mut rng, d),
            cross_k: vec(&mut rng, d * d),
            cross_v: vec(&mut rng, d * d),
            cross_v_bias: vec(&mut rng, d),
            cross_out: vec(&mut rng, d * d),
            cross_out_bias: vec(&mut rng, d),
            ln_cross_weight: vec(&mut rng, d),
            ln_cross_bias: vec(&mut rng, d),
            fc1: vec(&mut rng, 32 * d),
            fc1_bias: vec(&mut rng, 32),
            fc2: vec(&mut rng, d * 32),
            fc2_bias: vec(&mut rng, d),
            ln_fc_weight: vec(&mut rng, d),
            ln_fc_bias: vec(&mut rng, d),
        })
        .collect();
    let weights = crate::whisper::weights::WhisperWeights {
        conv1: vec(&mut rng, d * 80 * 3),
        conv1_bias: vec(&mut rng, d),
        conv2: vec(&mut rng, d * d * 3),
        conv2_bias: vec(&mut rng, d),
        enc_positions: compute::whisper::sinusoids(1500, d),
        enc_layers,
        enc_ln_weight: vec(&mut rng, d),
        enc_ln_bias: vec(&mut rng, d),
        embed_tokens: vec(&mut rng, 51864 * d),
        dec_positions: compute::whisper::sinusoids(448, d),
        dec_layers,
        dec_ln_weight: vec(&mut rng, d),
        dec_ln_bias: vec(&mut rng, d),
    };
    WhisperRunner::from_parts(config, en_tokens(), weights)
}

#[test]
fn english_only_prompt_drops_the_language_block() {
    let tokens = en_tokens();
    assert_eq!(build_prompt(&tokens, 50259), vec![50257, 50362]);
    let mut multilingual = tokens;
    multilingual.multilingual = true;
    assert_eq!(
        build_prompt(&multilingual, 50259),
        vec![50257, 50259, 50358, 50362]
    );
}

#[test]
fn suppression_masks_specials_and_keeps_eot() {
    let tokens = en_tokens();
    let mut logits = vec![-1.0f32; 51864];
    logits[100] = 7.0; // best text token
    logits[50300] = 99.0; // a language token: masked
    logits[51000] = 98.0; // a timestamp: masked
    logits[50256] = 6.0; // eot: allowed, but loses to 100
    assert_eq!(argmax_suppressed(&logits, &tokens, 51864), 100);
    logits[100] = 1.0;
    assert_eq!(argmax_suppressed(&logits, &tokens, 51864), 50256);
}

#[test]
fn transcribe_is_deterministic_across_runs() {
    let runner = synthetic_runner();
    // Two seconds of a fixed synthetic wave: the tiny random model will
    // emit garbage tokens, but the SAME tokens every time.
    let pcm: Vec<f32> = (0..2 * WHISPER_SAMPLE_RATE as usize)
        .map(|i| (i as f32 * 0.01).sin())
        .collect();
    let a = runner.transcribe(&pcm, None).unwrap();
    let b = runner.transcribe(&pcm, None).unwrap();
    assert_eq!(a, b);
}

#[test]
fn cooperative_controls_report_windows_and_cancel_in_the_token_loop() {
    use super::WhisperProgress;
    let runner = synthetic_runner();
    let pcm = vec![0.25; WHISPER_SAMPLE_RATE as usize];
    let expected = runner.transcribe(&pcm, None).unwrap();
    let mut reached_token = false;
    let mut windows = vec![];
    let cancelled = runner.transcribe_with_progress(&pcm, None, &|| false, |event| {
        match event {
            WhisperProgress::Window { index, total } => windows.push((index, total)),
            WhisperProgress::Token {
                window,
                emitted,
                budget,
            } => {
                assert_eq!((window, emitted), (0, 0));
                assert!(budget > 0);
                reached_token = true;
                return false;
            }
        }
        true
    });
    assert!(reached_token, "cancellation must reach the greedy decoder");
    assert_eq!(windows, vec![(0, 1)]);
    assert!(cancelled.unwrap_err().contains("cancelled"));
    let replay = runner
        .transcribe_with_progress(&pcm, None, &|| false, |_| true)
        .unwrap();
    assert_eq!(
        replay, expected,
        "cancelled KV must not leak into a later request"
    );
}

#[test]
fn long_audio_produces_strided_windows_with_absolute_offsets() {
    let runner = synthetic_runner();
    // 62 seconds: window 1 covers [0, 30), window 2 [29, 59), window 3
    // [58, 62] with zero padding; the exclusive segment ends follow the
    // 29-second stride.
    let len = 62 * WHISPER_SAMPLE_RATE as usize;
    let pcm = vec![0.25f32; len];
    let out = runner.transcribe(&pcm, None).unwrap();
    // The synthetic model emits deterministic garbage and may emit an
    // empty window; what the contract fixes is that every segment sits on
    // the 29-second stride grid, in order, covering the audio.
    let starts: Vec<f64> = out.segments.iter().map(|s| s.start_seconds).collect();
    assert!(
        starts.iter().all(|s| [0.0, 29.0, 58.0].contains(s)),
        "starts {starts:?}"
    );
    let ends: Vec<f64> = out.segments.iter().map(|s| s.end_seconds).collect();
    assert!(
        ends.iter().all(|e| [29.0, 58.0, 62.0].contains(e)),
        "ends {ends:?}"
    );
    assert!(out
        .segments
        .windows(2)
        .all(|w| w[0].start_seconds < w[1].start_seconds));
    // Segment indices are dense and ordered.
    for (i, seg) in out.segments.iter().enumerate() {
        assert_eq!(seg.index, i);
    }
}

#[test]
fn repetition_limit_stops_the_decode() {
    // A model whose logits always prefer the same text token: the decode
    // must stop at the repetition limit instead of looping forever.
    let runner = synthetic_runner();
    let pcm: Vec<f32> = (0..WHISPER_SAMPLE_RATE as usize)
        .map(|i| (i as f32 * 0.05).sin())
        .collect();
    let out = runner.transcribe(&pcm, None).unwrap();
    // One window, one segment (or none if every token is eot); whatever
    // the text, the run must terminate promptly and deterministically.
    assert!(out.segments.len() <= 1);
    let again = runner.transcribe(&pcm, None).unwrap();
    assert_eq!(out, again);
}

#[test]
fn self_kv_slots_track_positions() {
    // The KV cache must grow exactly one slot per decoded position; the
    // incremental-vs-full oracle in compute covers the values, this pins
    // the slot accounting through the runner path.
    let mut state = WhisperSelfKv::default();
    let d = 16;
    state.k.resize(3 * d, 0.0);
    state.v.resize(3 * d, 0.0);
    state.len += 3;
    assert_eq!(state.len, 3);
    assert_eq!(state.k.len(), 3 * d);
}

#[test]
fn window_stride_is_one_second_of_overlap() {
    assert_eq!(
        WINDOW_STRIDE_SAMPLES,
        WHISPER_WINDOW_SAMPLES - WHISPER_SAMPLE_RATE as usize
    );
    assert_eq!(WINDOW_STRIDE_SAMPLES, 464_000);
}

// ---- The on-disk open path: config, tokenizer, and safetensors by name. ----
// `synthetic_runner` above builds weights in memory; `WhisperRunner::open`
// is the path a real install takes, and nothing else exercised it.

pub(super) const TINY_CONFIG_JSON: &str = r#"{"model_type":"whisper","d_model":16,"encoder_layers":2,"decoder_layers":2,
    "encoder_attention_heads":4,"decoder_attention_heads":4,
    "encoder_ffn_dim":32,"decoder_ffn_dim":32,"vocab_size":51864,
    "num_mel_bins":80,"max_source_positions":1500,"max_target_positions":448,
    "activation_function":"gelu","scale_embedding":false}"#;

/// A tokenizer.json the HF `Tokenizer::from_file` accepts, carrying exactly
/// the specials `WhisperSpecialTokens::from_tokenizer_json` requires of an
/// English-only distribution.
fn tiny_tokenizer_json() -> String {
    let special = |id: u32, content: &str| {
        serde_json::json!({
            "id": id, "content": content, "single_word": false,
            "lstrip": false, "rstrip": false, "normalized": false, "special": true
        })
    };
    let specials = [
        (50256, "<|endoftext|>"),
        (50257, "<|startoftranscript|>"),
        (50258, "<|en|>"),
        (50358, "<|transcribe|>"),
        (50362, "<|notimestamps|>"),
        (50363, "<|0.00|>"),
    ];
    // Contiguous model IDs keep tokenizers from reallocating added IDs.
    let mut vocab = serde_json::Map::new();
    for id in 0..51864u32 {
        let name = specials
            .iter()
            .find(|(special_id, _)| *special_id == id)
            .map(|(_, content)| (*content).to_string())
            .unwrap_or_else(|| format!("token{id}"));
        vocab.insert(name, serde_json::json!(id));
    }
    serde_json::json!({
        "version": "1.0",
        "truncation": null,
        "padding": null,
        "added_tokens": [
            special(50256, "<|endoftext|>"),
            special(50257, "<|startoftranscript|>"),
            special(50258, "<|en|>"),
            special(50358, "<|transcribe|>"),
            special(50362, "<|notimestamps|>"),
            special(50363, "<|0.00|>"),
        ],
        "normalizer": null,
        "pre_tokenizer": null,
        "post_processor": null,
        "decoder": null,
        "model": {
            "type": "WordLevel",
            "vocab": vocab,
            "unk_token": "<|endoftext|>"
        }
    })
    .to_string()
}

/// Every tensor `load_from_safetensors` reads, with its expected length.
fn tiny_tensor_shapes() -> Vec<(String, usize)> {
    let d = 16usize;
    let ffn = 32usize;
    let mut tensors: Vec<(String, usize)> = vec![
        ("encoder.conv1.weight".into(), d * 80 * 3),
        ("encoder.conv1.bias".into(), d),
        ("encoder.conv2.weight".into(), d * d * 3),
        ("encoder.conv2.bias".into(), d),
        ("encoder.embed_positions.weight".into(), 1500 * d),
        ("encoder.layer_norm.weight".into(), d),
        ("encoder.layer_norm.bias".into(), d),
        ("decoder.embed_tokens.weight".into(), 51864 * d),
        ("decoder.embed_positions.weight".into(), 448 * d),
        ("decoder.layer_norm.weight".into(), d),
        ("decoder.layer_norm.bias".into(), d),
    ];
    let encoder: &[(&str, usize)] = &[
        ("self_attn.q_proj.weight", d * d),
        ("self_attn.q_proj.bias", d),
        ("self_attn.k_proj.weight", d * d),
        ("self_attn.v_proj.weight", d * d),
        ("self_attn.v_proj.bias", d),
        ("self_attn.out_proj.weight", d * d),
        ("self_attn.out_proj.bias", d),
        ("self_attn_layer_norm.weight", d),
        ("self_attn_layer_norm.bias", d),
        ("fc1.weight", ffn * d),
        ("fc1.bias", ffn),
        ("fc2.weight", d * ffn),
        ("fc2.bias", d),
        ("final_layer_norm.weight", d),
        ("final_layer_norm.bias", d),
    ];
    for layer in 0..2 {
        for (suffix, len) in encoder {
            tensors.push((format!("encoder.layers.{layer}.{suffix}"), *len));
        }
    }
    let decoder: &[(&str, usize)] = &[
        ("self_attn.q_proj.weight", d * d),
        ("self_attn.q_proj.bias", d),
        ("self_attn.k_proj.weight", d * d),
        ("self_attn.v_proj.weight", d * d),
        ("self_attn.v_proj.bias", d),
        ("self_attn.out_proj.weight", d * d),
        ("self_attn.out_proj.bias", d),
        ("self_attn_layer_norm.weight", d),
        ("self_attn_layer_norm.bias", d),
        ("encoder_attn.q_proj.weight", d * d),
        ("encoder_attn.q_proj.bias", d),
        ("encoder_attn.k_proj.weight", d * d),
        ("encoder_attn.v_proj.weight", d * d),
        ("encoder_attn.v_proj.bias", d),
        ("encoder_attn.out_proj.weight", d * d),
        ("encoder_attn.out_proj.bias", d),
        ("encoder_attn_layer_norm.weight", d),
        ("encoder_attn_layer_norm.bias", d),
        ("fc1.weight", ffn * d),
        ("fc1.bias", ffn),
        ("fc2.weight", d * ffn),
        ("fc2.bias", d),
        ("final_layer_norm.weight", d),
        ("final_layer_norm.bias", d),
    ];
    for layer in 0..2 {
        for (suffix, len) in decoder {
            tensors.push((format!("decoder.layers.{layer}.{suffix}"), *len));
        }
    }
    tensors
}

/// Values seeded from the tensor's own name, so a tensor loaded into the
/// wrong field, truncated, or duplicated cannot match the expectation.
fn values_for(name: &str, len: usize) -> Vec<f32> {
    let mut seed = 0xcbf2_9ce4_8422_2325u64;
    for byte in name.bytes() {
        seed ^= byte as u64;
        seed = seed.wrapping_mul(0x100_0000_01b3);
    }
    (0..len)
        .map(|_| {
            seed = seed
                .wrapping_mul(6364136223846793005)
                .wrapping_add(1442695040888963407);
            ((seed >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
        })
        .collect()
}

/// The safetensors framing `SafetensorsFile::open` reads: u64 LE header
/// length, JSON header, then the tensor blobs back to back.
fn fixture_shape(name: &str, len: usize) -> Vec<usize> {
    let name = name.strip_prefix("model.").unwrap_or(name);
    match name {
        "encoder.conv1.weight" => vec![16, 80, 3],
        "encoder.conv2.weight" => vec![16, 16, 3],
        "encoder.embed_positions.weight" => vec![1500, 16],
        "decoder.embed_positions.weight" => vec![448, 16],
        "decoder.embed_tokens.weight" => vec![51864, 16],
        _ if name.ends_with("fc1.weight") => vec![32, 16],
        _ if name.ends_with("fc2.weight") => vec![16, 32],
        _ if name.ends_with("_proj.weight") => vec![16, 16],
        _ => vec![len],
    }
}

fn write_safetensors(path: &std::path::Path, tensors: &[(String, Vec<f32>)]) {
    let mut header = serde_json::Map::new();
    let mut offset = 0usize;
    let mut blobs = Vec::new();
    for (name, data) in tensors {
        let bytes: Vec<u8> = data.iter().flat_map(|f| f.to_le_bytes()).collect();
        header.insert(
            name.clone(),
            serde_json::json!({
                "dtype": "F32",
                "shape": fixture_shape(name, data.len()),
                "data_offsets": [offset, offset + bytes.len()],
            }),
        );
        offset += bytes.len();
        blobs.push(bytes);
    }
    let header_json = serde_json::to_string(&header).unwrap();
    let mut file = Vec::new();
    file.extend_from_slice(&(header_json.len() as u64).to_le_bytes());
    file.extend_from_slice(header_json.as_bytes());
    for blob in &blobs {
        file.extend_from_slice(blob);
    }
    std::fs::write(path, &file).unwrap();
}

fn speech_install_dir(tag: &str) -> std::path::PathBuf {
    static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
    std::env::temp_dir().join(format!(
        "turbospark-whisper-open-{tag}-{}-{}",
        std::process::id(),
        NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
    ))
}

/// Writes a complete tiny install; `mutate_tensors` may rename, drop, or
/// resize entries before the file is framed.
fn write_speech_install(
    dir: &std::path::Path,
    mutate_tensors: impl FnOnce(Vec<(String, Vec<f32>)>) -> Vec<(String, Vec<f32>)>,
) {
    std::fs::create_dir_all(dir).unwrap();
    std::fs::write(dir.join("config.json"), TINY_CONFIG_JSON).unwrap();
    std::fs::write(dir.join("tokenizer.json"), tiny_tokenizer_json()).unwrap();
    let tensors: Vec<(String, Vec<f32>)> = tiny_tensor_shapes()
        .into_iter()
        .map(|(name, len)| {
            let values = values_for(&name, len);
            (name, values)
        })
        .collect();
    write_safetensors(&dir.join("model.safetensors"), &mutate_tensors(tensors));
    write_test_receipt(dir);
}

fn write_test_receipt(dir: &std::path::Path) {
    use model_io::speech_receipt::{FileEntry, SpeechInstallReceipt};
    let config =
        WhisperConfig::from_json_str(&std::fs::read_to_string(dir.join("config.json")).unwrap())
            .unwrap();
    let files = ["config.json", "tokenizer.json", "model.safetensors"]
        .into_iter()
        .map(|name| {
            let path = dir.join(name);
            (
                name.to_string(),
                FileEntry {
                    size: std::fs::metadata(&path).unwrap().len(),
                    sha256: model_io::hash_file(&path, 1 << 20).unwrap(),
                },
            )
        })
        .collect();
    SpeechInstallReceipt::new(
        "whisper",
        config.n_mels,
        config.memory_ceiling_bytes(),
        None,
        None,
        files,
    )
    .write_to_dir(dir)
    .unwrap();
}

#[test]
fn open_admission_precedes_tokenizer_and_weight_loading() {
    let dir = speech_install_dir("admission");
    std::fs::create_dir_all(&dir).unwrap();
    std::fs::write(dir.join("config.json"), TINY_CONFIG_JSON).unwrap();
    // No receipt, tokenizer, or tensor file exists. Memory refusal must win.
    let error = WhisperRunner::open_with_budget(&dir, || Ok(1))
        .err()
        .unwrap();
    assert!(error.contains("memory admission refused"), "{error}");
    let oversized = TINY_CONFIG_JSON.replace("\"encoder_layers\":2", "\"encoder_layers\":512");
    assert_ne!(oversized, TINY_CONFIG_JSON);
    std::fs::write(dir.join("config.json"), oversized).unwrap();
    let error =
        WhisperRunner::open_with_budget(&dir, || panic!("architecture gate must run first"))
            .err()
            .unwrap();
    assert!(
        error.contains("encoder_layers") && error.contains("maximum"),
        "{error}"
    );
    std::fs::remove_dir_all(dir).unwrap();
}

#[test]
fn open_budget_counts_mapped_source_and_metal_before_materialization() {
    let dir = speech_install_dir("budget-boundary");
    write_speech_install(&dir, |tensors| tensors);
    let config = WhisperConfig::from_json_str(TINY_CONFIG_JSON).unwrap();
    let mapped = std::fs::metadata(dir.join("model.safetensors"))
        .unwrap()
        .len();
    let estimate = super::admission::estimated_open_bytes(&config, mapped);
    assert!(estimate > config.memory_ceiling_bytes() * 2 + mapped);
    let error = WhisperRunner::open_with_budget(&dir, || Ok(estimate - 1))
        .err()
        .unwrap();
    assert!(error.contains("memory admission refused"), "{error}");
    let runner = WhisperRunner::open_with_budget(&dir, || Ok(estimate)).unwrap();
    assert_eq!(runner.weights.enc_layers.len(), 2);
    drop(runner);
    std::fs::remove_dir_all(dir).unwrap();
}

#[test]
fn open_requires_matching_receipt_metadata_and_content() {
    use model_io::speech_receipt::{SpeechInstallReceipt, SPEECH_RECEIPT_FILENAME};
    let dir = speech_install_dir("receipt");
    write_speech_install(&dir, |tensors| tensors);
    let original = SpeechInstallReceipt::read_from_dir(&dir).unwrap();
    std::fs::remove_file(dir.join(SPEECH_RECEIPT_FILENAME)).unwrap();
    let error = WhisperRunner::open_with_budget(&dir, || Ok(u64::MAX))
        .err()
        .unwrap();
    assert!(error.contains(SPEECH_RECEIPT_FILENAME), "{error}");
    for change in ["estimate", "mel", "family", "missing", "size", "hash"] {
        let mut receipt = original.clone();
        match change {
            "estimate" => receipt.memory_ceiling_bytes = 1,
            "mel" => receipt.n_mels = 128,
            "family" => receipt.speech_family = "qwen3_asr".into(),
            "missing" => {
                receipt.files.remove("tokenizer.json");
            }
            "size" => receipt.files.get_mut("config.json").unwrap().size += 1,
            "hash" => receipt.files.get_mut("model.safetensors").unwrap().sha256 = "0".repeat(64),
            _ => unreachable!(),
        }
        receipt.write_to_dir(&dir).unwrap();
        let error = WhisperRunner::open_with_budget(&dir, || Ok(u64::MAX))
            .err()
            .unwrap();
        assert!(error.contains("receipt"), "{change}: {error}");
    }
    original.write_to_dir(&dir).unwrap();
    // Preserve lengths, so size-only verification cannot catch the edits.
    for name in ["config.json", "tokenizer.json", "model.safetensors"] {
        let path = dir.join(name);
        let original_bytes = std::fs::read(&path).unwrap();
        let mut changed = original_bytes.clone();
        if name == "config.json" {
            let json =
                TINY_CONFIG_JSON.replace("\"scale_embedding\":false", "\"scale_embedding\":true ");
            assert_ne!(json.as_bytes(), original_bytes);
            changed = json.into_bytes();
        } else if name == "tokenizer.json" {
            let at = changed.iter().position(|b| *b == b'k').unwrap();
            changed[at] = b'K';
        } else {
            *changed.last_mut().unwrap() ^= 1;
        }
        assert_eq!(changed.len(), original_bytes.len());
        std::fs::write(&path, changed).unwrap();
        let error = WhisperRunner::open_with_budget(&dir, || Ok(u64::MAX))
            .err()
            .unwrap();
        assert!(
            (error.contains("checksum") || error.contains("SHA-256")) && error.contains(name),
            "{name}: {error}"
        );
        std::fs::write(path, original_bytes).unwrap();
    }
    std::fs::remove_dir_all(dir).unwrap();
}

#[test]
fn open_loads_named_tensors_specials_and_tokenizer_from_disk() {
    let dir = speech_install_dir("full");
    write_speech_install(&dir, |tensors| tensors);
    let runner = WhisperRunner::open(&dir).unwrap_or_else(|error| panic!("{error}"));

    assert_eq!(runner.config.d_model, 16);
    assert_eq!(runner.config.encoder_layers(), 2);
    assert_eq!(runner.config.decoder_layers(), 2);
    assert!(runner.tokenizer.is_some(), "open must load tokenizer.json");

    // The specials come from tokenizer.json, not the multilingual defaults.
    let tokens = &runner.tokens;
    assert_eq!(tokens.eot, 50256);
    assert_eq!(tokens.sot, 50257);
    assert_eq!(tokens.transcribe, 50358);
    assert_eq!(tokens.no_timestamps, 50362);
    assert_eq!(tokens.first_timestamp, 50363);
    assert_eq!(
        (tokens.first_language, tokens.last_language),
        (50258, 50258)
    );
    assert!(
        !tokens.multilingual,
        "English text vocabulary selects English prompt"
    );

    let weights = &runner.weights;
    let expect = |name: &str, got: &[f32]| {
        assert_eq!(got, &values_for(name, got.len())[..], "{name} mismatched");
    };
    expect("encoder.conv1.weight", &weights.conv1);
    expect("encoder.conv1.bias", &weights.conv1_bias);
    expect("encoder.conv2.weight", &weights.conv2);
    expect("encoder.conv2.bias", &weights.conv2_bias);
    expect("encoder.embed_positions.weight", &weights.enc_positions);
    expect("encoder.layer_norm.weight", &weights.enc_ln_weight);
    expect("encoder.layer_norm.bias", &weights.enc_ln_bias);
    expect("decoder.embed_tokens.weight", &weights.embed_tokens);
    expect("decoder.embed_positions.weight", &weights.dec_positions);
    expect("decoder.layer_norm.weight", &weights.dec_ln_weight);
    expect("decoder.layer_norm.bias", &weights.dec_ln_bias);
    assert_eq!(weights.enc_layers.len(), 2);
    assert_eq!(weights.dec_layers.len(), 2);

    let layer = &weights.enc_layers[1];
    let p = "encoder.layers.1.";
    expect(&format!("{p}self_attn.q_proj.weight"), &layer.q);
    expect(&format!("{p}self_attn.q_proj.bias"), &layer.q_bias);
    expect(&format!("{p}self_attn.k_proj.weight"), &layer.k);
    expect(&format!("{p}self_attn.v_proj.weight"), &layer.v);
    expect(&format!("{p}self_attn.v_proj.bias"), &layer.v_bias);
    expect(&format!("{p}self_attn.out_proj.weight"), &layer.out);
    expect(&format!("{p}self_attn.out_proj.bias"), &layer.out_bias);
    expect(
        &format!("{p}self_attn_layer_norm.weight"),
        &layer.ln1_weight,
    );
    expect(&format!("{p}self_attn_layer_norm.bias"), &layer.ln1_bias);
    expect(&format!("{p}fc1.weight"), &layer.fc1);
    expect(&format!("{p}fc1.bias"), &layer.fc1_bias);
    expect(&format!("{p}fc2.weight"), &layer.fc2);
    expect(&format!("{p}fc2.bias"), &layer.fc2_bias);
    expect(&format!("{p}final_layer_norm.weight"), &layer.ln2_weight);
    expect(&format!("{p}final_layer_norm.bias"), &layer.ln2_bias);

    let layer = &weights.dec_layers[1];
    let p = "decoder.layers.1.";
    expect(&format!("{p}self_attn.q_proj.weight"), &layer.self_q);
    expect(&format!("{p}self_attn.k_proj.weight"), &layer.self_k);
    expect(&format!("{p}self_attn.v_proj.weight"), &layer.self_v);
    expect(&format!("{p}self_attn.out_proj.weight"), &layer.self_out);
    expect(&format!("{p}encoder_attn.q_proj.weight"), &layer.cross_q);
    expect(&format!("{p}encoder_attn.k_proj.weight"), &layer.cross_k);
    expect(&format!("{p}encoder_attn.v_proj.weight"), &layer.cross_v);
    expect(
        &format!("{p}encoder_attn.out_proj.weight"),
        &layer.cross_out,
    );
    expect(
        &format!("{p}encoder_attn_layer_norm.weight"),
        &layer.ln_cross_weight,
    );
    expect(&format!("{p}fc1.weight"), &layer.fc1);
    expect(&format!("{p}fc2.weight"), &layer.fc2);
    expect(&format!("{p}final_layer_norm.weight"), &layer.ln_fc_weight);
    expect(&format!("{p}final_layer_norm.bias"), &layer.ln_fc_bias);

    std::fs::remove_dir_all(&dir).ok();
}

#[test]
fn open_falls_back_to_the_model_prefixed_tensor_name() {
    let dir = speech_install_dir("prefixed");
    write_speech_install(&dir, |mut tensors| {
        let index = tensors
            .iter()
            .position(|(name, _)| name == "encoder.conv1.weight")
            .expect("conv1 in the tensor set");
        tensors[index].0 = "model.encoder.conv1.weight".to_string();
        tensors
    });
    let runner = WhisperRunner::open(&dir).unwrap_or_else(|error| panic!("{error}"));
    assert_eq!(
        runner.weights.conv1,
        values_for("encoder.conv1.weight", 16 * 80 * 3)
    );
    std::fs::remove_dir_all(&dir).ok();
}

#[test]
fn open_refuses_missing_and_wrong_sized_tensors_by_name() {
    let dir = speech_install_dir("missing");
    write_speech_install(&dir, |tensors| {
        tensors
            .into_iter()
            .filter(|(name, _)| name != "decoder.layers.1.fc2.bias")
            .collect()
    });
    let error = match WhisperRunner::open(&dir) {
        Ok(_) => panic!("a missing tensor must refuse the open"),
        Err(error) => error,
    };
    assert!(
        error.contains("decoder.layers.1.fc2.bias"),
        "the refusal must name the missing tensor: {error}"
    );
    std::fs::remove_dir_all(&dir).ok();

    let dir = speech_install_dir("sized");
    write_speech_install(&dir, |mut tensors| {
        let index = tensors
            .iter()
            .position(|(name, _)| name == "encoder.conv1.weight")
            .expect("conv1 in the tensor set");
        tensors[index].1.truncate(8);
        tensors
    });
    let error = match WhisperRunner::open(&dir) {
        Ok(_) => panic!("a wrong-sized tensor must refuse the open"),
        Err(error) => error,
    };
    assert!(
        error.contains("encoder.conv1.weight"),
        "the refusal must name the wrong-sized tensor: {error}"
    );
    std::fs::remove_dir_all(&dir).ok();
}

#[test]
fn open_refuses_malformed_norm_and_projection_shapes() {
    for name in [
        "decoder.layers.1.encoder_attn.k_proj.weight",
        "encoder.layers.0.final_layer_norm.bias",
    ] {
        let dir = speech_install_dir("bad-shape");
        write_speech_install(&dir, |mut tensors| {
            tensors
                .iter_mut()
                .find(|(key, _)| key == name)
                .unwrap()
                .1
                .pop();
            tensors
        });
        let result = WhisperRunner::open(&dir);
        assert!(result.is_err(), "malformed {name} must refuse at open");
        let error = result.err().unwrap();
        assert!(error.contains(name), "{error}");
        std::fs::remove_dir_all(dir).unwrap();
    }
}

#[test]
fn malformed_mel_rows_and_nonfinite_pcm_refuse_before_compute() {
    let runner = synthetic_runner();
    assert!(runner
        .encode_window(&[vec![0.0; 80], vec![0.0; 81]])
        .is_err());
    assert!(runner
        .encode_window(&[vec![0.0; 80], vec![0.0; 79]])
        .is_err());
    assert!(runner.transcribe(&[f32::NAN], None).is_err());
    assert!(runner.transcribe(&[0.0], Some("de")).is_err());
}

#[test]
fn open_uses_loaded_special_ids_and_refuses_config_mismatch() {
    let dir = speech_install_dir("resolved-tokens");
    write_speech_install(&dir, |tensors| tensors);
    let mut json: serde_json::Value = serde_json::from_str(&tiny_tokenizer_json()).unwrap();
    // This raw ID conflicts with the model vocabulary. The loader resolves
    // the existing token's actual ID, which the runner must use.
    json["added_tokens"][1]["id"] = serde_json::json!(42);
    std::fs::write(dir.join("tokenizer.json"), json.to_string()).unwrap();
    write_test_receipt(&dir);
    let runner = WhisperRunner::open(&dir).unwrap();
    assert_eq!(
        runner.tokens.sot,
        runner
            .tokenizer
            .as_ref()
            .unwrap()
            .token_to_id("<|startoftranscript|>")
            .unwrap()
    );
    assert_eq!(runner.tokens.sot, 50257);
    let mut config: serde_json::Value = serde_json::from_str(TINY_CONFIG_JSON).unwrap();
    config["eos_token_id"] = serde_json::json!(42);
    std::fs::write(dir.join("config.json"), config.to_string()).unwrap();
    write_test_receipt(&dir);
    let error = WhisperRunner::open(&dir).err().unwrap();
    assert!(error.contains("eos_token_id"), "{error}");
    std::fs::remove_dir_all(dir).ok();
}

#[test]
fn decoder_rejects_invalid_prompt_and_encoded_shape() {
    let runner = synthetic_runner();
    for (encoded, seq, prompt) in [
        (vec![0.0; 16], 1, vec![]),
        (vec![0.0; 15], 1, vec![50257]),
        (vec![], 0, vec![50257]),
        (vec![0.0; 16], 1, vec![50257; 449]),
    ] {
        assert!(crate::whisper::WindowDecoder::new(&runner, &encoded, seq, &prompt).is_err());
    }
    let mut runner = runner;
    runner.config.max_target_positions = 1;
    let decoder = crate::whisper::WindowDecoder::new(&runner, &[0.0; 16], 1, &[50257]).unwrap();
    let result = decoder.decode().unwrap();
    assert_eq!(result.stop, crate::whisper::StopReason::TokenBudget);
    assert!(result.text_tokens.is_empty());
}
/// Real-model, cross-device comparison, gated on `TS_STT_TEST_MODEL`
/// naming a speech install (the same env the FFI surface test uses). The
/// Metal engine's encoder hidden states are bounded against the CPU
/// reference's, and the full transcripts must carry the same segment
/// texts (greedy argmax can legitimately flip on f32 noise only when the
/// logits tie, which this test has never seen on the witness models --
/// if it starts failing on text, run the encoder bound first).
#[test]
#[ignore]
fn real_model_metal_matches_cpu() {
    let Some(model_dir) = std::env::var_os("TS_STT_TEST_MODEL") else {
        eprintln!("TS_STT_TEST_MODEL not set; skipping real-model metal comparison");
        return;
    };
    let model_dir = std::path::PathBuf::from(model_dir);
    assert!(
        model_dir.exists(),
        "TS_STT_TEST_MODEL must name an existing install"
    );

    // CPU-forced runner: the field is pub(crate), tests flip it directly.
    let mut cpu_runner = WhisperRunner::open(&model_dir).expect("runner opens");
    cpu_runner.metal = None;

    // The metal runner and a matching engine.
    let runner = WhisperRunner::open(&model_dir).expect("runner opens");
    let mut engine =
        crate::whisper::metal::WhisperMetalEngine::new(&runner.weights, &runner.config)
            .expect("engine builds");

    // 30 seconds of deterministic pseudo-random PCM: not speech, but the
    // numerics comparison does not care.
    let pcm: Vec<f32> = (0..audio::whisper::WHISPER_WINDOW_SAMPLES)
        .map(|i| ((i % 1600) as f32 / 1600.0 - 0.5) * 0.2)
        .collect();
    let n_mels = runner.config.n_mels;
    let mel = audio::whisper::whisper_log_mel_window(&pcm, n_mels).expect("mel");
    let frames = mel.len();
    let mut mel_band = vec![0.0f32; n_mels * frames];
    for (t, frame) in mel.iter().enumerate() {
        for (b, &v) in frame.iter().enumerate() {
            mel_band[b * frames + t] = v;
        }
    }

    // Encoder hidden states, CPU versus Metal.
    let cpu_enc = cpu_runner.encode_window(&mel).expect("cpu encode");
    engine.encode(&mel_band, frames).expect("metal encode");
    engine.build_cross_caches().expect("cross caches");
    let seq = runner.config.max_source_positions;
    let d = runner.config.d_model;
    let gpu_enc = engine.read_enc_out();
    assert_eq!(cpu_enc.len(), seq * d);
    assert_eq!(gpu_enc.len(), cpu_enc.len());
    let mut max_err = 0.0f32;
    let mut sum_err = 0.0f64;
    for (a, b) in cpu_enc.iter().zip(&gpu_enc) {
        max_err = max_err.max((a - b).abs());
        sum_err += f64::from((a - b).abs());
    }
    println!(
        "encoder hidden states: max abs {max_err:.3e}, mean abs {:.3e}",
        sum_err / cpu_enc.len() as f64
    );

    // Warmed-prompt logits on both paths. The Metal engine's final warm
    // step already produced its logits (the step pass ends in the final
    // norm and the vocab projection); the CPU side re-walks the reference
    // stack over the prompt and projects through the runner.
    let prompt = crate::whisper::decode::build_prompt(&runner.tokens, 0);
    engine.begin_decode();
    let mut gpu_logits = Vec::new();
    for (pos, &token) in prompt.iter().enumerate() {
        gpu_logits = engine
            .step_with_logits(
                &runner.weights.embed_tokens,
                &runner.weights.dec_positions,
                token,
                pos,
            )
            .expect("metal warm step");
    }

    let vocab = runner.config.vocab_size;

    // CPU warmed logits: walk the reference decoder layer stack over the
    // prompt the same way WindowDecoder::new does.
    let mut cross = Vec::new();
    for layer in &cpu_runner.weights.dec_layers {
        let w = borrow_decoder_for_test(layer);
        cross.push(cross_kv(&cpu_enc, seq, &w, d));
    }
    let mut state: Vec<WhisperSelfKv> = (0..cpu_runner.weights.dec_layers.len())
        .map(|_| WhisperSelfKv::default())
        .collect();
    let scale_embed = runner.config.embed_scale();
    let mut input = Vec::new();
    for (pos, &token) in prompt.iter().enumerate() {
        let row = token as usize * d;
        input = cpu_runner.weights.embed_tokens[row..row + d]
            .iter()
            .zip(&cpu_runner.weights.dec_positions[pos * d..pos * d + d])
            .map(|(&e, &p)| e * scale_embed + p)
            .collect();
        for (i, layer) in cpu_runner.weights.dec_layers.iter().enumerate() {
            let w = borrow_decoder_for_test(layer);
            input = whisper_decoder_layer_step(
                &input,
                pos,
                &w,
                &mut state[i],
                &cross[i],
                d,
                runner.config.decoder_attention_heads(),
                runner.config.layer_norm_eps,
            );
        }
    }
    let mut cpu_logits = cpu_runner.logits(&input).expect("cpu logits");

    let logit_err = max_abs(
        &cpu_logits[..4096.min(vocab)],
        &gpu_logits[..4096.min(vocab)],
    );
    println!("first 4096 logits max abs: {logit_err:.3e}");
    let mut top = Vec::new();
    for _i in 0..8 {
        let ci = cpu_logits
            .iter()
            .take(vocab)
            .enumerate()
            .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
            .map(|(i, v)| (i, *v))
            .unwrap();
        let gi = gpu_logits
            .iter()
            .take(vocab)
            .enumerate()
            .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
            .map(|(i, v)| (i, *v))
            .unwrap();
        top.push((ci, gi));
        // zero out best cpu entry for the next pass
        cpu_logits[ci.0] = f32::NEG_INFINITY;
        gpu_logits[gi.0] = f32::NEG_INFINITY;
    }
    println!("cpu top8: {top:?}");
}

/// Per-stage encoder bisection: conv+pos, layer 0, layer 1.
#[test]
#[ignore]
fn real_model_metal_encoder_bisect() {
    let Some(model_dir) = std::env::var_os("TS_STT_TEST_MODEL") else {
        eprintln!("TS_STT_TEST_MODEL not set; skipping");
        return;
    };
    let model_dir = std::path::PathBuf::from(model_dir);
    let cpu_runner = WhisperRunner::open(&model_dir).expect("runner opens");
    let mut engine =
        crate::whisper::metal::WhisperMetalEngine::new(&cpu_runner.weights, &cpu_runner.config)
            .expect("engine builds");

    let pcm: Vec<f32> = (0..audio::whisper::WHISPER_WINDOW_SAMPLES)
        .map(|i| ((i % 1600) as f32 / 1600.0 - 0.5) * 0.2)
        .collect();
    let n_mels = cpu_runner.config.n_mels;
    let mel = audio::whisper::whisper_log_mel_window(&pcm, n_mels).expect("mel");
    let frames = mel.len();
    let mut mel_band = vec![0.0f32; n_mels * frames];
    for (t, frame) in mel.iter().enumerate() {
        for (b, &v) in frame.iter().enumerate() {
            mel_band[b * frames + t] = v;
        }
    }
    let d = cpu_runner.config.d_model;
    let seq = cpu_runner.config.max_source_positions;

    // CPU: conv front end + positional add, per compute::whisper.
    let conv_w = compute::whisper::WhisperConvWeights {
        conv1_weight: &cpu_runner.weights.conv1,
        conv1_bias: &cpu_runner.weights.conv1_bias,
        conv2_weight: &cpu_runner.weights.conv2,
        conv2_bias: &cpu_runner.weights.conv2_bias,
    };
    let conv_out = compute::whisper::whisper_conv_frontend(&mel_band, n_mels, frames, &conv_w, d);
    let mut cpu_hidden = vec![0.0f32; seq * d];
    for t in 0..seq {
        for i in 0..d {
            cpu_hidden[t * d + i] =
                conv_out[i * seq + t] + cpu_runner.weights.enc_positions[t * d + i];
        }
    }

    // GPU: conv + pos only (stop_after 0).
    engine
        .encode_partial(&mel_band, frames, Some(0))
        .expect("partial 0");
    let gpu_hidden = engine.read_hidden();
    let mut max_err = max_abs(&cpu_hidden, &gpu_hidden);
    println!("after conv+pos: max abs {max_err:.3e}");

    // Where does it diverge? First coordinates above threshold.
    for (idx, (c, g)) in cpu_hidden.iter().zip(&gpu_hidden).enumerate() {
        if (c - g).abs() > 0.05 {
            let t = idx / d;
            let i = idx % d;
            println!("first hidden divergence at t {t} dim {i}: cpu {c:.6} gpu {g:.6}");
            break;
        }
    }
    // Direct conv1/conv2 comparisons: CPU conv_out ([d, seq] band-major) vs GPU.
    let (conv1_gpu, conv2_gpu) = engine.read_conv_out();
    let cpu_conv1 = compute::whisper::conv1d3_gelu(
        &mel_band,
        n_mels,
        frames,
        &cpu_runner.weights.conv1,
        &cpu_runner.weights.conv1_bias,
        d,
        1,
        1,
    );
    println!("conv1: max abs {}", max_abs(&cpu_conv1, &conv1_gpu));
    // Whole-zero t blocks and where the max error sits.
    let zero_blocks: Vec<usize> = (0..frames)
        .filter(|&t| conv1_gpu[t * d..(t + 1) * d].iter().all(|v| *v == 0.0))
        .collect();
    println!(
        "conv1 all-zero t blocks: {} (first 20: {:?})",
        zero_blocks.len(),
        &zero_blocks[..zero_blocks.len().min(20)]
    );
    let (mut max_i, mut max_v) = (0usize, 0.0f32);
    for (i, (c, g)) in cpu_conv1.iter().zip(&conv1_gpu).enumerate() {
        let e = (c - g).abs();
        if e > max_v {
            max_v = e;
            max_i = i;
        }
    }
    println!(
        "conv1 max err at channel {} t {}: cpu {:.6} gpu {:.6}",
        max_i / frames,
        max_i % frames,
        cpu_conv1[max_i],
        conv1_gpu[max_i]
    );
    println!("conv2: max abs {}", max_abs(&conv_out, &conv2_gpu));
    for (idx, (c, g)) in conv_out.iter().zip(&conv2_gpu).enumerate() {
        if (c - g).abs() > 0.05 {
            let ch = idx / seq;
            let t = idx % seq;
            println!("first conv2 divergence at channel {ch} t {t}: cpu {c:.6} gpu {g:.6}");
            break;
        }
    }
    // mel upload sanity: first 8 values read back.
    let mel_read = engine.read_mel_in();
    println!(
        "mel readback first 8: {:?} vs cpu {:?} (len {} vs {})",
        &mel_read[..8],
        &mel_band[..8],
        mel_read.len(),
        mel_band.len()
    );

    // Layer 0.
    let l0 = &cpu_runner.weights.enc_layers[0];
    let w = compute::whisper::WhisperEncoderLayerWeights {
        q_weight: &l0.q,
        q_bias: &l0.q_bias,
        k_weight: &l0.k,
        v_weight: &l0.v,
        v_bias: &l0.v_bias,
        out_weight: &l0.out,
        out_bias: &l0.out_bias,
        ln1_weight: &l0.ln1_weight,
        ln1_bias: &l0.ln1_bias,
        fc1_weight: &l0.fc1,
        fc1_bias: &l0.fc1_bias,
        fc2_weight: &l0.fc2,
        fc2_bias: &l0.fc2_bias,
        ln2_weight: &l0.ln2_weight,
        ln2_bias: &l0.ln2_bias,
    };
    let cpu_l0 = compute::whisper::whisper_encoder_layer(
        &cpu_hidden,
        seq,
        &w,
        d,
        cpu_runner.config.num_attention_heads,
        cpu_runner.config.encoder_ffn_dim,
        cpu_runner.config.layer_norm_eps,
    );
    engine
        .encode_partial(&mel_band, frames, Some(1))
        .expect("partial 1");
    let gpu_l0 = engine.read_hidden();
    max_err = max_abs(&cpu_l0, &gpu_l0);
    println!("after layer 0: max abs {max_err:.3e}");

    // Layer 1.
    let l1 = &cpu_runner.weights.enc_layers[1];
    let w1 = compute::whisper::WhisperEncoderLayerWeights {
        q_weight: &l1.q,
        q_bias: &l1.q_bias,
        k_weight: &l1.k,
        v_weight: &l1.v,
        v_bias: &l1.v_bias,
        out_weight: &l1.out,
        out_bias: &l1.out_bias,
        ln1_weight: &l1.ln1_weight,
        ln1_bias: &l1.ln1_bias,
        fc1_weight: &l1.fc1,
        fc1_bias: &l1.fc1_bias,
        fc2_weight: &l1.fc2,
        fc2_bias: &l1.fc2_bias,
        ln2_weight: &l1.ln2_weight,
        ln2_bias: &l1.ln2_bias,
    };
    let cpu_l1 = compute::whisper::whisper_encoder_layer(
        &cpu_l0,
        seq,
        &w1,
        d,
        cpu_runner.config.num_attention_heads,
        cpu_runner.config.encoder_ffn_dim,
        cpu_runner.config.layer_norm_eps,
    );
    engine
        .encode_partial(&mel_band, frames, Some(2))
        .expect("partial 2");
    let gpu_l1 = engine.read_hidden();
    max_err = max_abs(&cpu_l1, &gpu_l1);
    println!("after layer 1: max abs {max_err:.3e}");
}

/// The full-transcript equality gate on real audio-like input. Slow: the
/// CPU reference path takes minutes per window. Split from the fast
/// encoder/logits comparison so iteration on numerics does not pay it.
#[test]
#[ignore]
fn real_model_metal_matches_cpu_transcript() {
    let Some(model_dir) = std::env::var_os("TS_STT_TEST_MODEL") else {
        eprintln!("TS_STT_TEST_MODEL not set; skipping");
        return;
    };
    let model_dir = std::path::PathBuf::from(model_dir);
    assert!(model_dir.exists(), "TS_STT_TEST_MODEL must exist");

    let mut cpu_runner = WhisperRunner::open(&model_dir).expect("runner opens");
    cpu_runner.metal = None;
    let runner = WhisperRunner::open(&model_dir).expect("runner opens");
    let mut engine =
        crate::whisper::metal::WhisperMetalEngine::new(&runner.weights, &runner.config)
            .expect("engine builds");

    // 30 seconds of deterministic pseudo-random PCM.
    let pcm: Vec<f32> = (0..audio::whisper::WHISPER_WINDOW_SAMPLES)
        .map(|i| ((i % 1600) as f32 / 1600.0 - 0.5) * 0.2)
        .collect();
    let cpu_trans = cpu_runner
        .transcribe(&pcm, Some("en"))
        .expect("cpu transcribe");
    let metal_trans = runner
        .transcribe_metal(&mut engine, &pcm, Some("en"), &|| false, &mut |_| true)
        .expect("metal transcribe");
    let cpu_text: Vec<&str> = cpu_trans.segments.iter().map(|s| s.text.as_str()).collect();
    let metal_text: Vec<&str> = metal_trans
        .segments
        .iter()
        .map(|s| s.text.as_str())
        .collect();
    println!("cpu text:   {cpu_text:?}");
    println!("metal text: {metal_text:?}");
    assert_eq!(cpu_text, metal_text, "same segment texts on both devices");
}

use compute::whisper::{cross_kv, whisper_decoder_layer_step, WhisperDecoderLayerWeights};

fn max_abs(a: &[f32], b: &[f32]) -> f32 {
    a.iter()
        .zip(b)
        .map(|(x, y)| (x - y).abs())
        .fold(0.0f32, f32::max)
}

fn borrow_decoder_for_test(
    layer: &crate::whisper::weights::WhisperDecoderLayerOwned,
) -> WhisperDecoderLayerWeights<'_> {
    WhisperDecoderLayerWeights {
        self_q_weight: &layer.self_q,
        self_q_bias: &layer.self_q_bias,
        self_k_weight: &layer.self_k,
        self_v_weight: &layer.self_v,
        self_v_bias: &layer.self_v_bias,
        self_out_weight: &layer.self_out,
        self_out_bias: &layer.self_out_bias,
        ln_self_weight: &layer.ln_self_weight,
        ln_self_bias: &layer.ln_self_bias,
        cross_q_weight: &layer.cross_q,
        cross_q_bias: &layer.cross_q_bias,
        cross_k_weight: &layer.cross_k,
        cross_v_weight: &layer.cross_v,
        cross_v_bias: &layer.cross_v_bias,
        cross_out_weight: &layer.cross_out,
        cross_out_bias: &layer.cross_out_bias,
        ln_cross_weight: &layer.ln_cross_weight,
        ln_cross_bias: &layer.ln_cross_bias,
        fc1_weight: &layer.fc1,
        fc1_bias: &layer.fc1_bias,
        fc2_weight: &layer.fc2,
        fc2_bias: &layer.fc2_bias,
        ln_fc_weight: &layer.ln_fc_weight,
        ln_fc_bias: &layer.ln_fc_bias,
    }
}

/// Synthetic cross-device decoder-step gate: the Metal engine and the CPU
/// reference stack must agree on the decoder stream after every warm step
/// on deterministic pseudo-random weights. Fast (d_model 16); this is the
/// regression that catches composition bugs the per-kernel parity tests
/// cannot see. Requires Metal; skipped on failure to build the engine.
#[test]
fn synthetic_metal_decoder_step_matches_cpu() {
    use compute::whisper::{cross_kv, whisper_decoder_layer_step, WhisperSelfKv};

    let mut runner = synthetic_runner();
    runner.metal = None;
    let weights = &runner.weights;
    let config = runner.config.clone();
    let d = config.d_model;
    let seq = config.max_source_positions;

    let mut engine = match crate::whisper::metal::WhisperMetalEngine::new(weights, &config) {
        Ok(engine) => engine,
        Err(e) => {
            eprintln!("no Metal engine ({e}); skipping");
            return;
        }
    };

    // A zero mel window is a valid encoder input.
    let frames = 2 * seq;
    let mel = vec![0.0f32; config.n_mels * frames];
    let mel_band = mel; // already band-major
    let cpu_enc = runner
        .encode_window(&audio_mel_from_band(&mel_band, config.n_mels))
        .expect("cpu encode");
    engine.encode(&mel_band, frames).expect("metal encode");
    engine.build_cross_caches().expect("cross caches");
    let _ = cpu_enc;

    // Cross-cache composition check: the engine's transposed value cache
    // must equal the transpose of the natural-layout cross V. This is the
    // check that catches a missing transpose dispatch: the attention step
    // would read uninitialized memory and diverge only in the decode
    // stream, where the cause is far from the symptom.
    {
        let engine_cross_v = engine.read_dec_scratch("cross_v");
        let engine_cross_v_t = engine.read_dec_scratch("cross_v_t");
        let mut want_t = vec![0.0f32; seq * d];
        for t in 0..seq {
            for i in 0..d {
                want_t[i * seq + t] = engine_cross_v[t * d + i];
            }
        }
        let err = max_abs(&want_t, &engine_cross_v_t);
        assert_eq!(err, 0.0, "cross_v_t must be the exact transpose of cross_v");
    }

    // Warm the English-only prompt through both stacks, comparing the
    // stream after every step.
    let prompt = crate::whisper::decode::build_prompt(&runner.tokens, 0);
    let mut state: Vec<WhisperSelfKv> = (0..weights.dec_layers.len())
        .map(|_| WhisperSelfKv::default())
        .collect();
    let mut cross = Vec::new();
    for layer in &weights.dec_layers {
        let w = borrow_decoder_for_test(layer);
        cross.push(cross_kv(&cpu_enc, seq, &w, d));
    }
    let scale_embed = config.embed_scale();
    engine.begin_decode();
    for (pos, &token) in prompt.iter().enumerate() {
        let row = token as usize * d;
        let input: Vec<f32> = weights.embed_tokens[row..row + d]
            .iter()
            .zip(&weights.dec_positions[pos * d..pos * d + d])
            .map(|(&e, &p)| e * scale_embed + p)
            .collect();
        let mut h = input.clone();
        for (i, layer) in weights.dec_layers.iter().enumerate() {
            let w = borrow_decoder_for_test(layer);
            h = whisper_decoder_layer_step(
                &h,
                pos,
                &w,
                &mut state[i],
                &cross[i],
                d,
                config.decoder_attention_heads(),
                config.layer_norm_eps,
            );
        }
        engine
            .step_with_logits(&weights.embed_tokens, &weights.dec_positions, token, pos)
            .expect("metal step");
        let gpu_hidden = engine.read_dec_hidden();
        let err = max_abs(&h, &gpu_hidden);
        println!("synthetic step {pos}: dec_hidden max abs {err:.3e}");
        // The GEMV path reduces each dot product as a 32-lane simd tree,
        // a different f32 summation order from the reference's sequential
        // loop. Measured: 1.1e-3 on this d=16 fixture, and on the real
        // tiny.en (d=384, 4 decoder layers) the same noise reads 1.3e-2
        // on logits near 10 with the CPU's top-8 tokens reproduced
        // exactly (real_model_metal_matches_cpu). The bound pins that
        // measured noise; a structural bug reads orders of magnitude
        // higher (the binding-index bug read 7.5 on step 0).
        assert!(
            err < 4e-3,
            "synthetic step {pos}: stream diverged, max abs {err:.3e}"
        );
    }
}

/// Rebuilds the frame-major mel the CPU encoder expects from a
/// band-major window.
fn audio_mel_from_band(band: &[f32], n_mels: usize) -> Vec<Vec<f32>> {
    let frames = band.len() / n_mels;
    (0..frames)
        .map(|t| (0..n_mels).map(|b| band[b * frames + t]).collect())
        .collect()
}