//! Speech recognition on top of Canary (ONNX).
//!
//! The fact that shapes this crate: **decoding a window costs more than
//! recording it**. On this machine a 6 s window takes about 800 ms to decode
//! and only advances 400 ms of audio. Any design that processes every window
//! in order keeps falling behind the speaker without limit.
//!
//! The way out is dropping work: the decoding thread drains its whole queue,
//! keeps only the most recent window and discards the rest. Latency stays
//! bounded by one decode instead of growing unchecked, and what is lost are
//! partial transcriptions that were going to be overwritten anyway.

use std::time::{Duration, Instant};

use crossbeam_channel::{Receiver, Sender};

use asist_core::config::AsrConfig;
use asist_core::error::{Error, Result};
use asist_core::event::TurnId;

pub use canary_rs::{Canary, CanarySession, ExecutionConfig, ExecutionProvider, StreamConfig};

/// Work arriving at the decoder.
#[derive(Debug)]
pub enum AsrJob {
    /// New audio for the sliding window.
    Window {
        turn: TurnId,
        samples: Vec<f32>,
        at: Instant,
    },
    /// Closed utterance, to be transcribed in full.
    Utterance {
        turn: TurnId,
        samples: Vec<f32>,
        at: Instant,
    },
    /// The utterance was empty: resets the window state.
    Reset { turn: TurnId },
}

/// What the decoder returns.
#[derive(Debug, Clone)]
pub enum AsrResult {
    Partial {
        turn: TurnId,
        committed: String,
        volatile: String,
        /// Windows discarded as stale before this one.
        dropped: usize,
        decode: Duration,
    },
    Final {
        turn: TurnId,
        text: String,
        audio_secs: f32,
        decode: Duration,
        /// Instant the user stopped talking. It is the origin perceived latency
        /// is measured from, and it cannot be taken here: by the time the
        /// transcription is ready almost a second has passed.
        spoken_at: Instant,
    },
    Empty {
        turn: TurnId,
    },
    Error {
        turn: TurnId,
        message: String,
    },
}

/// Engine loaded and ready to decode.
pub struct Recognizer {
    model: Canary,
    config: AsrConfig,
}

impl Recognizer {
    /// Loads the model from `config.model_dir`.
    pub fn load(config: &AsrConfig) -> Result<Self> {
        if !config.model_dir.is_dir() {
            return Err(Error::Asr(format!(
                "the model directory does not exist: {}. Run scripts/bootstrap.sh",
                config.model_dir.display()
            )));
        }
        let started = Instant::now();
        let model = Canary::from_pretrained(
            config.model_dir.to_string_lossy().as_ref(),
            Some(execution_config(config)),
        )
        .map_err(|e| Error::Asr(format!("could not load Canary: {e}")))?;

        tracing::info!(
            target: "asr",
            dir = %config.model_dir.display(),
            provider = %config.execution_provider,
            ms = started.elapsed().as_millis(),
            "modelo cargado"
        );
        Ok(Self {
            model,
            config: config.clone(),
        })
    }

    /// Standalone session to transcribe in one go (tests and checks).
    pub fn transcribe(&self, samples: &[f32], sample_rate: u32) -> Result<String> {
        let mut session = self.model.session();
        session
            .transcribe_samples(
                samples,
                sample_rate as usize,
                1,
                &self.config.source_lang,
                &self.config.target_lang,
            )
            .map(|r| normalize(&r.text))
            .map_err(|e| Error::Asr(e.to_string()))
    }

    /// Decoder loop. Runs on its own thread until `jobs` is closed.
    pub fn run(self, jobs: Receiver<AsrJob>, out: Sender<AsrResult>) {
        let mut stream = match self.model.stream(
            self.config.source_lang.clone(),
            self.config.target_lang.clone(),
            stream_config(&self.config),
        ) {
            Ok(stream) => stream,
            Err(err) => {
                let _ = out.send(AsrResult::Error {
                    turn: TurnId::default(),
                    message: format!("could not open the stream: {err}"),
                });
                return;
            }
        };
        let mut session = self.model.session();
        let mut committed = String::new();
        let step_samples = (self.config.step * 16_000.0).max(1.0) as usize;

        while let Some(batch) = drain(&jobs) {
            // All the audio left behind an utterance close belongs to a turn that
            // is about to be transcribed in full: spending a window on it would
            // be work doomed to be overwritten.
            let mut pending: Vec<f32> = Vec::new();
            let mut pending_turn = TurnId::default();
            let mut newest = Instant::now();
            let mut dropped = 0usize;

            for job in batch {
                match job {
                    AsrJob::Window { turn, samples, at } => {
                        pending_turn = turn;
                        pending.extend_from_slice(&samples);
                        newest = at;
                    }
                    AsrJob::Reset { turn } => {
                        dropped += pending.len() / step_samples;
                        pending.clear();
                        stream.reset();
                        committed.clear();
                        if out.send(AsrResult::Empty { turn }).is_err() {
                            return;
                        }
                    }
                    AsrJob::Utterance { turn, samples, at } => {
                        dropped += pending.len() / step_samples;
                        pending.clear();
                        stream.reset();
                        committed.clear();
                        let result = decode_final(&mut session, &self.config, turn, &samples, at);
                        if out.send(result).is_err() {
                            return;
                        }
                    }
                }
            }

            if pending.is_empty() || !self.config.partials {
                continue;
            }
            // Only the newest window survives; the older ones no longer describe
            // what is being said now.
            dropped += (pending.len() / step_samples).saturating_sub(1);

            let started = Instant::now();
            let chunks = match stream.push_samples(&pending, 16_000, 1) {
                Ok(chunks) => chunks,
                Err(err) => {
                    let _ = out.send(AsrResult::Error {
                        turn: pending_turn,
                        message: err.to_string(),
                    });
                    continue;
                }
            };
            let Some(chunk) = chunks.last() else { continue };
            if is_degenerate(&chunk.result.text) {
                continue;
            }

            append_delta(&mut committed, chunk.delta_text.trim());
            let volatile = volatile_tail(&committed, chunk.result.text.trim());
            let decode = started.elapsed();
            tracing::debug!(
                target: "asr",
                turn = pending_turn.0,
                ms = decode.as_millis(),
                delay_ms = newest.elapsed().as_millis(),
                dropped = dropped,
                "ventana"
            );
            if out
                .send(AsrResult::Partial {
                    turn: pending_turn,
                    committed: committed.clone(),
                    volatile,
                    dropped,
                    decode,
                })
                .is_err()
            {
                return;
            }
        }
    }
}

fn decode_final(
    session: &mut CanarySession,
    config: &AsrConfig,
    turn: TurnId,
    samples: &[f32],
    at: Instant,
) -> AsrResult {
    let started = Instant::now();
    let audio_secs = samples.len() as f32 / 16_000.0;
    match session.transcribe_samples(samples, 16_000, 1, &config.source_lang, &config.target_lang) {
        Ok(result) => {
            let text = normalize(&result.text);
            let decode = started.elapsed();
            tracing::debug!(
                target: "asr",
                turn = turn.0,
                ms = decode.as_millis(),
                audio_s = audio_secs,
                rtf = decode.as_secs_f32() / audio_secs.max(0.001),
                delay_ms = at.elapsed().as_millis(),
                "final transcription"
            );
            if text.is_empty() {
                AsrResult::Empty { turn }
            } else {
                AsrResult::Final {
                    turn,
                    text,
                    audio_secs,
                    decode,
                    spoken_at: at,
                }
            }
        }
        Err(err) => AsrResult::Error {
            turn,
            message: err.to_string(),
        },
    }
}

fn execution_config(config: &AsrConfig) -> ExecutionConfig {
    let provider = match config.execution_provider.trim().to_lowercase().as_str() {
        "cuda" => ExecutionProvider::Cuda,
        "tensorrt" => ExecutionProvider::TensorRT,
        "rocm" => ExecutionProvider::ROCm,
        "openvino" => ExecutionProvider::OpenVINO,
        "webgpu" => ExecutionProvider::WebGPU,
        "coreml" => ExecutionProvider::CoreML,
        "directml" => ExecutionProvider::DirectML,
        _ => ExecutionProvider::Cpu,
    };
    ExecutionConfig::new()
        .with_execution_provider(provider)
        .with_threads(config.inter_threads, config.intra_threads)
}

fn stream_config(config: &AsrConfig) -> StreamConfig {
    StreamConfig::new()
        .with_window_duration(config.window)
        .with_step_duration(config.step)
        .with_emit_partial(true)
        .with_pad_partial(false)
        .with_stability_window(config.stability)
        // The reason for the whole design: never drag a queue of old windows.
        .with_max_windows_per_push(1)
}

/// Grabs everything queued at once, blocking until the first item arrives.
fn drain(jobs: &Receiver<AsrJob>) -> Option<Vec<AsrJob>> {
    let mut batch = vec![jobs.recv().ok()?];
    while let Ok(job) = jobs.try_recv() {
        batch.push(job);
    }
    Some(batch)
}

/// Joins the stable text with the new fragment, keeping the spacing right.
pub fn append_delta(committed: &mut String, delta: &str) {
    if delta.is_empty() {
        return;
    }
    let needs_space = !committed.is_empty()
        && !committed.ends_with(' ')
        && !delta.starts_with(|c: char| c.is_ascii_punctuation());
    if needs_space {
        committed.push(' ');
    }
    committed.push_str(delta);
}

/// Part of the current window that is not stable yet.
pub fn volatile_tail(committed: &str, window: &str) -> String {
    // The window repeats the end of what is already committed; the
    // interesting part is what follows it.
    match window.rfind(last_words(committed, 3).as_str()) {
        Some(idx) if !committed.is_empty() => window[idx + last_words(committed, 3).len()..]
            .trim_start()
            .to_string(),
        _ => window.to_string(),
    }
}

fn last_words(text: &str, n: usize) -> String {
    let words: Vec<&str> = text.split_whitespace().collect();
    words[words.len().saturating_sub(n)..].join(" ")
}

/// Canary sometimes gets stuck repeating a token when the window is
/// almost silence. Showing that would confuse more than help.
pub fn is_degenerate(text: &str) -> bool {
    let words: Vec<&str> = text.split_whitespace().collect();
    if words.len() < 6 {
        return false;
    }
    let distinct: std::collections::HashSet<&&str> = words.iter().collect();
    distinct.len() * 4 <= words.len()
}

/// Fixes the punctuation spacing the decoder sometimes leaves.
pub fn normalize(text: &str) -> String {
    let mut out = String::with_capacity(text.len());
    for c in text.trim().chars() {
        if matches!(c, ',' | '.' | '!' | '?' | ';' | ':') {
            while out.ends_with(' ') {
                out.pop();
            }
        }
        out.push(c);
    }
    out
}

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

    #[test]
    fn delta_is_joined_with_a_space_except_before_punctuation() {
        let mut s = String::from("hola");
        append_delta(&mut s, "mundo");
        assert_eq!(s, "hola mundo");
        append_delta(&mut s, ",");
        assert_eq!(s, "hola mundo,");
        append_delta(&mut s, "");
        assert_eq!(s, "hola mundo,");
    }

    #[test]
    fn volatile_tail_is_what_follows_the_committed_text() {
        assert_eq!(volatile_tail("hola qué tal", "hola qué tal estás"), "estás");
    }

    #[test]
    fn without_committed_text_the_whole_window_is_volatile() {
        assert_eq!(volatile_tail("", "hola qué tal"), "hola qué tal");
    }

    #[test]
    fn degenerate_repetition_is_detected() {
        assert!(is_degenerate("sí sí sí sí sí sí sí sí"));
        assert!(!is_degenerate("hola qué tal estás hoy amigo"));
        assert!(
            !is_degenerate("sí sí"),
            "a short sentence is not degenerate"
        );
    }

    #[test]
    fn punctuation_loses_the_space_before_it() {
        assert_eq!(normalize("hola , qué tal ?"), "hola, qué tal?");
        assert_eq!(normalize("  ya está  "), "ya está");
    }
}