#[path = "shared/utils.rs"]
mod utils;

use canary_rs::{Canary, StreamConfig};
use std::time::Instant;

/// Benchmarks the pieces that drive live-streaming latency:
///
/// * encoder cost as a function of window length,
/// * decoder cost per generated token,
/// * end-to-end cost of one streaming window vs. one full-utterance decode.
///
/// Usage: `bench_live <wav> [lang]`, honouring the same `CANARY_*` environment
/// variables as the other examples.
fn main() -> Result<(), Box<dyn std::error::Error>> {
    let wav = match std::env::args().nth(1) {
        Some(path) => path,
        None => return Err("usage: bench_live <wav> [lang]".into()),
    };
    let lang = std::env::args().nth(2).unwrap_or_else(|| "en".to_string());

    let config = utils::execution_config_from_env();
    let model_dir =
        std::env::var("CANARY_MODEL_DIR").unwrap_or_else(|_| "canary-180m-flash".to_string());

    let t0 = Instant::now();
    let model = Canary::from_pretrained(&model_dir, Some(config))?;
    println!("model load: {:.2?}", t0.elapsed());

    let mut reader = hound::WavReader::open(&wav)?;
    let spec = reader.spec();
    let samples: Vec<f32> = reader
        .samples::<i16>()
        .map(|s| s.map(|v| v as f32 / 32768.0))
        .collect::<std::result::Result<_, _>>()?;
    let sr = spec.sample_rate as usize;
    let duration = samples.len() as f32 / sr as f32;
    println!("audio: {:.2}s @ {} Hz, {} ch", duration, sr, spec.channels);

    let mut session = model.session();
    let warm = &samples[..(sr).min(samples.len())];
    let _ = session.transcribe_samples(warm, sr, 1, &lang, &lang)?;

    // A live window always ends mid-utterance, so mirror that here: take the
    // *last* `window` seconds and append the same tail silence the stream uses.
    let tail = StreamConfig::default().tail_silence_duration;
    println!("\nrolling windows (tail silence {:.2}s):", tail);
    for window_s in [2.0f32, 4.0, 6.0, 8.0] {
        if window_s > duration {
            continue;
        }
        let start = samples.len() - (sr as f32 * window_s) as usize;
        let mut win = samples[start..].to_vec();
        win.resize(win.len() + (sr as f32 * tail) as usize, 0.0);

        let mut best = f64::MAX;
        let mut tokens = 0;
        for _ in 0..3 {
            let t = Instant::now();
            let r = session.transcribe_samples(&win, sr, 1, &lang, &lang)?;
            best = best.min(t.elapsed().as_secs_f64());
            tokens = r.tokens.len();
        }
        // No tokens means the model saw nothing worth transcribing, and the time is encoder-only.
        let per_token = match tokens {
            0 => "encoder only".to_string(),
            n => format!("{:.1} ms/token", best * 1000.0 / n as f64),
        };
        println!(
            "  {:>4.1}s window -> {:>6.0} ms  ({:>3} tokens, {})",
            window_s,
            best * 1000.0,
            tokens,
            per_token
        );
    }

    let t = Instant::now();
    let r = session.transcribe_samples(&samples, sr, 1, &lang, &lang)?;
    let full = t.elapsed().as_secs_f64();
    println!(
        "\nfull utterance ({:.1}s) -> {:.0} ms ({} tokens, {:.1}x realtime)",
        duration,
        full * 1000.0,
        r.tokens.len(),
        duration as f64 / full
    );
    println!("text: {}", r.text);

    Ok(())
}