//! Unit tests for the weight-free code paths: pure math, formatting, and the
//! audio I/O WAV round-trip. These run on every `cargo test` — no HF auth or
//! model downloads required.

use approx::assert_relative_eq;
use rtx_csm::{
    audio_io::{TARGET_SAMPLE_RATE, load_mono_24k, write_wav_24k_mono},
    config::ModelConfig,
    util::{all_zero_codebooks, rms},
};

#[test]
fn rms_zero_on_identical_signals() {
    let a = vec![0.1, -0.2, 0.3, -0.4, 0.5];
    assert_relative_eq!(rms(&a, &a), 0.0, epsilon = 1e-6);
}

#[test]
fn rms_matches_hand_calc() {
    let a = vec![0.0, 1.0, 2.0, 3.0];
    let b = vec![0.0, 0.0, 0.0, 0.0];
    // sum_sq = 0 + 1 + 4 + 9 = 14, n=4, rms = sqrt(14/4) ≈ 1.8708286
    assert_relative_eq!(rms(&a, &b), (14.0f32 / 4.0).sqrt(), epsilon = 1e-5);
}

#[test]
fn all_zero_codebooks_true_only_on_all_zero() {
    use candle_core::{Device, Tensor};
    let dev = Device::Cpu;
    let zeros = Tensor::zeros((1, 32), candle_core::DType::I64, &dev).unwrap();
    assert!(all_zero_codebooks(&zeros).unwrap());

    // One non-zero → false.
    let mut v = vec![0i64; 32];
    v[7] = 42;
    let nonzero = Tensor::from_vec(v, (1, 32), &dev).unwrap();
    assert!(!all_zero_codebooks(&nonzero).unwrap());
}

#[test]
fn speaker_formatting_brackets_id() {
    use rtx_csm::tokenizer::CsmTokenizer;
    assert_eq!(CsmTokenizer::format_segment(0, "hello"), "[0]hello");
    assert_eq!(CsmTokenizer::format_segment(1, "world"), "[1]world");
    assert_eq!(CsmTokenizer::format_segment(0, ""), "[0]");
}

#[test]
fn config_csm_1b_defaults() {
    let cfg = ModelConfig::csm_1b();
    assert_eq!(cfg.text_vocab_size, 128_256);
    assert_eq!(cfg.audio_vocab_size, 2051);
    assert_eq!(cfg.audio_num_codebooks, 32);
    assert_eq!(cfg.sample_rate, 24_000);
    assert_eq!(cfg.frame_rate_hz, 12.5);
    assert_relative_eq!(cfg.frame_duration_ms(), 80.0, epsilon = 1e-6);
}

#[test]
fn wav_write_read_roundtrip() {
    let tmp = tempfile::NamedTempFile::with_suffix(".wav").unwrap();
    // Synth a 0.1-second sine wave at 440 Hz so the test file is ~4.8 KB.
    let n = (TARGET_SAMPLE_RATE as f32 * 0.1) as usize;
    let input: Vec<f32> = (0..n)
        .map(|i| {
            0.5 * (2.0 * std::f32::consts::PI * 440.0 * i as f32 / TARGET_SAMPLE_RATE as f32).sin()
        })
        .collect();

    write_wav_24k_mono(tmp.path(), &input).unwrap();
    let read_back = load_mono_24k(tmp.path()).unwrap();

    // hound writes 16-bit PCM so we lose some precision; expect same length and
    // small per-sample error.
    assert_eq!(read_back.len(), input.len());
    let err = rms(&input, &read_back);
    assert!(err < 1e-3, "wav round-trip rms {err} too high");
}

#[test]
fn empty_prompt_has_shape_one_zero_cb_plus_one() {
    use candle_core::{DType, Device};
    // Directly validate that a zero-length audio tensor has the right shape —
    // catches regressions in the empty-audio branch of prompt.rs.
    let dev = Device::Cpu;
    let cb = 32;
    let t = candle_core::Tensor::zeros((1, 0, cb + 1), DType::U32, &dev).unwrap();
    assert_eq!(t.dims(), &[1, 0, 33]);
}