//! Phase 8.8 — end-to-end Moonshine transcription. Loads real audio,
//! runs encoder + greedy decode + tokenizer, prints the transcript.
//!
//! Usage:
//! ```bash
//! cargo run -p rtx-csm --release --features metal --example moonshine_transcribe -- \
//!     --in /tmp/asr_test.flac
//! ```

use anyhow::{Context, Result};
use candle_core::{Device, Tensor};
use clap::Parser;
use hf_hub::api::sync::Api;
use rtx_csm::{audio_io, moonshine};
use std::path::PathBuf;
use std::time::Instant;

#[derive(Debug, Parser)]
struct Cli {
    #[arg(long = "in", default_value = "/tmp/asr_test.flac")]
    input: PathBuf,
    /// Maximum decoded tokens (Moonshine config caps at 194).
    #[arg(long, default_value_t = 100)]
    max_tokens: usize,
}

fn main() -> Result<()> {
    let cli = Cli::parse();
    let device = if candle_core::utils::metal_is_available() {
        Device::new_metal(0)?
    } else {
        Device::Cpu
    };
    eprintln!("device: {device:?}");

    // Download both safetensors and tokenizer.json.
    let api = Api::new()?;
    let repo = api.model("UsefulSensors/moonshine-tiny".to_string());
    let weights = repo.get("model.safetensors")?;
    let tok_path = repo.get("tokenizer.json")?;
    eprintln!("weights:   {}", weights.display());
    eprintln!("tokenizer: {}", tok_path.display());

    // Audio: load + resample to 16 kHz.
    let pcm = audio_io::load_mono_at_rate(&cli.input, 16_000).context("load audio")?;
    let audio_secs = pcm.len() as f32 / 16_000.0;
    eprintln!(
        "audio:     {} samples ({audio_secs:.2}s @ 16 kHz)",
        pcm.len()
    );

    let cfg = moonshine::MoonshineConfig::tiny();
    let load_t = Instant::now();
    let (encoder, decoder) = moonshine::load_full(&weights, &device, &cfg)?;
    eprintln!("load:      {:.2}s", load_t.elapsed().as_secs_f32());
    let tok =
        moonshine::load_tokenizer(&tok_path).map_err(|e| anyhow::anyhow!("tokenizer: {e}"))?;
    eprintln!("tokenizer loaded (vocab {})", tok.get_vocab_size(true));

    // Encode audio.
    let pcm_t = Tensor::from_vec(pcm.clone(), (1, 1, pcm.len()), &device)?;
    let enc_t = Instant::now();
    let enc = encoder.forward(&pcm_t)?;
    eprintln!(
        "encode:    {:.0} ms (output {:?})",
        enc_t.elapsed().as_millis(),
        enc.shape()
    );

    // Greedy decode (cached path — O(T) per token instead of O(T²)).
    let dec_t = Instant::now();
    let token_ids = decoder.generate_cached(&enc, &cfg, cli.max_tokens)?;
    let dec_ms = dec_t.elapsed().as_millis();
    eprintln!(
        "decode:    {dec_ms} ms ({} tokens, {:.1} ms/token, KV-cached)",
        token_ids.len(),
        dec_ms as f64 / token_ids.len().max(1) as f64
    );

    // Detokenize.
    let text = tok
        .decode(&token_ids, true)
        .map_err(|e| anyhow::anyhow!("detok: {e}"))?;
    println!();
    println!("=== transcript ===");
    println!("{text}");
    println!();
    println!("token ids: {token_ids:?}");

    // Realtime factor.
    let rtf = (enc_t.elapsed().as_secs_f64() + dec_ms as f64 / 1000.0) / audio_secs as f64;
    println!();
    println!(
        "realtime factor: {rtf:.3}x  ({:.1}s of compute / {:.1}s of audio)",
        enc_t.elapsed().as_secs_f64() + dec_ms as f64 / 1000.0,
        audio_secs
    );

    Ok(())
}